221 lines
8.3 KiB
Python
221 lines
8.3 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 _bool(name: str, default: bool) -> bool:
|
||
raw = os.getenv(name)
|
||
if raw is None or not raw.strip():
|
||
return default
|
||
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||
|
||
|
||
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)
|
||
|
||
|
||
def _csv(name: str) -> frozenset[str]:
|
||
raw = os.getenv(name, "")
|
||
values = {
|
||
part.strip().lower()
|
||
for part in raw.replace(";", ",").split(",")
|
||
if part.strip()
|
||
}
|
||
return frozenset(values)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Settings:
|
||
bot_token: str
|
||
telegram_proxy: str | None = None
|
||
telegram_retry_attempts: int = 0
|
||
telegram_retry_max_delay: float = 30.0
|
||
telegram_health_interval: float = 60.0
|
||
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"
|
||
workflow_prompt: Path = PROJECT_ROOT / "workflows" / "prompt_enhance.json"
|
||
prompt_enhance_enabled: bool = False
|
||
prompt_enhance_timeout: float = 120.0
|
||
storage_dir: Path = PROJECT_ROOT / "storage"
|
||
database_url: str = ""
|
||
db_pool_min_size: int = 1
|
||
db_pool_max_size: int = 10
|
||
worker_concurrency: int = 1
|
||
max_queue_size: int = 20
|
||
allowed_user_ids: frozenset[int] = field(default_factory=frozenset)
|
||
admin_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
|
||
moderation_enabled: bool = True
|
||
blocked_public_figures: frozenset[str] = field(default_factory=frozenset)
|
||
payments_enabled: bool = True
|
||
stars_price: int = 5
|
||
premium_stars_price: int = 15
|
||
web_page_size: int = 50
|
||
web_bind_ip: str = "127.0.0.1"
|
||
web_admin_username: str = "admin"
|
||
web_admin_password: str = ""
|
||
krea2_q4_model: str = (
|
||
"Krea2_turbo_uncensored_edit_v1.1-Q4_K_M.gguf"
|
||
)
|
||
krea2_q6_model: str = (
|
||
"Krea2_turbo_uncensored_edit-v1.1-Q6_K.gguf"
|
||
)
|
||
|
||
@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://"
|
||
)
|
||
|
||
database_url = os.getenv("DATABASE_URL", "").strip()
|
||
|
||
settings = cls(
|
||
bot_token=token,
|
||
telegram_proxy=proxy,
|
||
telegram_retry_attempts=max(0, _int("TELEGRAM_RETRY_ATTEMPTS", 0)),
|
||
telegram_retry_max_delay=max(
|
||
1.0, _float("TELEGRAM_RETRY_MAX_DELAY", 30.0)
|
||
),
|
||
telegram_health_interval=max(
|
||
10.0, _float("TELEGRAM_HEALTH_INTERVAL", 60.0)
|
||
),
|
||
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"),
|
||
workflow_prompt=_path(
|
||
"WORKFLOW_PROMPT", "workflows/prompt_enhance.json"
|
||
),
|
||
prompt_enhance_enabled=_bool("PROMPT_ENHANCE_ENABLED", False),
|
||
prompt_enhance_timeout=max(
|
||
10.0, _float("PROMPT_ENHANCE_TIMEOUT", 120.0)
|
||
),
|
||
storage_dir=_path("STORAGE_DIR", "storage"),
|
||
database_url=database_url,
|
||
db_pool_min_size=max(1, _int("DB_POOL_MIN_SIZE", 1)),
|
||
db_pool_max_size=max(1, _int("DB_POOL_MAX_SIZE", 10)),
|
||
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"),
|
||
admin_user_ids=_user_ids("ADMIN_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)),
|
||
moderation_enabled=_bool("MODERATION_ENABLED", True),
|
||
blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"),
|
||
payments_enabled=_bool("PAYMENTS_ENABLED", True),
|
||
stars_price=max(1, _int("STARS_PRICE", 5)),
|
||
premium_stars_price=max(1, _int("PREMIUM_STARS_PRICE", 15)),
|
||
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
|
||
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
|
||
or "127.0.0.1",
|
||
web_admin_username=os.getenv(
|
||
"WEB_ADMIN_USERNAME", "admin"
|
||
).strip()
|
||
or "admin",
|
||
web_admin_password=os.getenv("WEB_ADMIN_PASSWORD", "").strip(),
|
||
krea2_q4_model=os.getenv(
|
||
"KREA2_Q4_MODEL",
|
||
cls.krea2_q4_model,
|
||
).strip()
|
||
or cls.krea2_q4_model,
|
||
krea2_q6_model=os.getenv(
|
||
"KREA2_Q6_MODEL",
|
||
cls.krea2_q6_model,
|
||
).strip()
|
||
or cls.krea2_q6_model,
|
||
)
|
||
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)
|