144 lines
5.5 KiB
Python
144 lines
5.5 KiB
Python
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from urllib.parse import urlsplit, urlunsplit
|
||
|
||
from dotenv import load_dotenv
|
||
|
||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||
load_dotenv(PROJECT_ROOT / ".env")
|
||
|
||
|
||
def _path(name: str, default: str) -> Path:
|
||
raw = os.getenv(name, default).strip() or default
|
||
path = Path(raw).expanduser()
|
||
return path if path.is_absolute() else PROJECT_ROOT / path
|
||
|
||
|
||
def _int(name: str, default: int) -> int:
|
||
raw = os.getenv(name)
|
||
if raw is None or not raw.strip():
|
||
return default
|
||
try:
|
||
return int(raw)
|
||
except ValueError as exc:
|
||
raise ValueError(f"{name} должен быть целым числом") from exc
|
||
|
||
|
||
def _float(name: str, default: float) -> float:
|
||
raw = os.getenv(name)
|
||
if raw is None or not raw.strip():
|
||
return default
|
||
try:
|
||
return float(raw)
|
||
except ValueError as exc:
|
||
raise ValueError(f"{name} должен быть числом") from exc
|
||
|
||
|
||
def _user_ids(name: str) -> frozenset[int]:
|
||
raw = os.getenv(name, "")
|
||
values: set[int] = set()
|
||
for part in raw.replace(";", ",").replace(" ", ",").split(","):
|
||
part = part.strip()
|
||
if not part:
|
||
continue
|
||
try:
|
||
values.add(int(part))
|
||
except ValueError as exc:
|
||
raise ValueError(f"{name}: некорректный user_id {part!r}") from exc
|
||
return frozenset(values)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Settings:
|
||
bot_token: str
|
||
telegram_proxy: str | None = None
|
||
comfy_base_url: str = "http://192.168.31.113:8188"
|
||
workflow_t2i: Path = PROJECT_ROOT / "workflows" / "txt2img.json"
|
||
workflow_i2i: Path = PROJECT_ROOT / "workflows" / "img2img.json"
|
||
workflow_mapping: Path = PROJECT_ROOT / "workflows" / "mapping.json"
|
||
storage_dir: Path = PROJECT_ROOT / "storage"
|
||
db_path: Path = PROJECT_ROOT / "data" / "bot.db"
|
||
worker_concurrency: int = 1
|
||
max_queue_size: int = 20
|
||
allowed_user_ids: frozenset[int] = field(default_factory=frozenset)
|
||
progress_interval: float = 3.0
|
||
generation_timeout: float = 600.0
|
||
log_level: str = "INFO"
|
||
default_negative: str = "low quality, blurry, bad anatomy, extra digits"
|
||
max_prompt_length: int = 1000
|
||
max_photo_bytes: int = 15 * 1024 * 1024
|
||
web_page_size: int = 50
|
||
web_username: str = "admin"
|
||
web_password: str = ""
|
||
|
||
@property
|
||
def comfy_ws_url(self) -> str:
|
||
parts = urlsplit(self.comfy_base_url.rstrip("/"))
|
||
scheme = "wss" if parts.scheme == "https" else "ws"
|
||
path = parts.path.rstrip("/") + "/ws"
|
||
return urlunsplit((scheme, parts.netloc, path, "", ""))
|
||
|
||
@classmethod
|
||
def from_env(cls, require_token: bool = True) -> Settings:
|
||
token = os.getenv("BOT_TOKEN", "").strip()
|
||
if require_token and not token:
|
||
raise ValueError("BOT_TOKEN не задан. Добавь его в .env")
|
||
|
||
base_url = os.getenv("COMFY_BASE_URL", "http://192.168.31.113:8188").strip()
|
||
if not base_url.startswith(("http://", "https://")):
|
||
raise ValueError("COMFY_BASE_URL должен начинаться с http:// или https://")
|
||
|
||
proxy = os.getenv("TELEGRAM_PROXY", "").strip() or None
|
||
if proxy and not proxy.startswith(
|
||
("http://", "https://", "socks5://", "socks5h://")
|
||
):
|
||
raise ValueError(
|
||
"TELEGRAM_PROXY должен начинаться с http://, https://, "
|
||
"socks5:// или socks5h://"
|
||
)
|
||
|
||
settings = cls(
|
||
bot_token=token,
|
||
telegram_proxy=proxy,
|
||
comfy_base_url=base_url.rstrip("/"),
|
||
workflow_t2i=_path("WORKFLOW_T2I", "workflows/txt2img.json"),
|
||
workflow_i2i=_path("WORKFLOW_I2I", "workflows/img2img.json"),
|
||
workflow_mapping=_path("WORKFLOW_MAPPING", "workflows/mapping.json"),
|
||
storage_dir=_path("STORAGE_DIR", "storage"),
|
||
db_path=_path("DB_PATH", "data/bot.db"),
|
||
worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)),
|
||
max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)),
|
||
allowed_user_ids=_user_ids("ALLOWED_USER_IDS"),
|
||
progress_interval=max(1.0, _float("PROGRESS_INTERVAL", 3.0)),
|
||
generation_timeout=max(30.0, _float("GENERATION_TIMEOUT", 600.0)),
|
||
log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO",
|
||
default_negative=os.getenv("DEFAULT_NEGATIVE", "").strip()
|
||
or cls.default_negative,
|
||
max_prompt_length=max(1, _int("MAX_PROMPT_LENGTH", 1000)),
|
||
max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)),
|
||
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
|
||
web_username=os.getenv("WEB_USERNAME", "admin").strip() or "admin",
|
||
web_password=os.getenv("WEB_PASSWORD", "").strip(),
|
||
)
|
||
return settings
|
||
|
||
|
||
def get_settings(require_token: bool = True) -> Settings:
|
||
return Settings.from_env(require_token=require_token)
|
||
|
||
|
||
def setup_logging(level: str) -> None:
|
||
logging.basicConfig(
|
||
level=getattr(logging, level.upper(), logging.INFO),
|
||
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
||
datefmt="%Y-%m-%d %H:%M:%S",
|
||
force=True,
|
||
)
|
||
# httpx пишет полный URL запроса, включая токен бота.
|
||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|