src/config.py

"""
Runtime configuration for Anima RDBT Gradio + embedded ComfyUI backend.
Mirrors defaults from the original RDBT workflow and Comfy model layout.
"""
from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Final

# --- RDBT defaults (match default_workflow_rdbt.json) ---
DEFAULT_WIDTH: Final[int] = 1024
DEFAULT_HEIGHT: Final[int] = 1024
DEFAULT_STEPS: Final[int] = 16
DEFAULT_CFG: Final[float] = 1.0
DEFAULT_SAMPLER: Final[str] = "euler_ancestral"
DEFAULT_SCHEDULER: Final[str] = "simple"
DEFAULT_DENOISE: Final[float] = 1.0
DEFAULT_BATCH_SIZE: Final[int] = 1
DEFAULT_NEGATIVE_PROMPT: Final[str] = ""
DEFAULT_SEED: Final[int] = 42
DEFAULT_RANDOMIZE_SEED: Final[bool] = True

# Model filenames
RDBT_UNET_NAME: Final[str] = "rdbt-anima-p3-v024f-16step-dmd2.safetensors"
CLIP_NAME: Final[str] = "qwen_3_06b_base.safetensors"
CLIP_TYPE: Final[str] = "qwen_image"
VAE_NAME: Final[str] = "qwen_image_vae.safetensors"
UNET_WEIGHT_DTYPE: Final[str] = "default"

# --- UI / validation ranges ---
MIN_WH: Final[int] = 512
MAX_WH: Final[int] = 2048
MIN_STEPS: Final[int] = 1
MAX_STEPS: Final[int] = 50
MIN_CFG: Final[float] = 1.0
MAX_CFG: Final[float] = 2.0
CFG_STEP: Final[float] = 0.1
MIN_DENOISE: Final[float] = 0.0
MAX_DENOISE: Final[float] = 1.0
DENOISE_STEP: Final[float] = 0.01
MIN_BATCH: Final[int] = 1
MAX_BATCH: Final[int] = 4
MIN_SEED: Final[int] = 0
MAX_SEED: Final[int] = 2**32 - 1

# Samplers: broad allowlist; unknown values fall back in validation layer
SAMPLER_CHOICES: Final[tuple[str, ...]] = (
    "euler",
    "euler_ancestral",
    "euler_ancestral_cfg_pp",
    "heun",
    "heunpp2",
    "dpm_2",
    "dpm_2_ancestral",
    "lms",
    "dpm_adaptive",
    "dpm_fast",
    "dpmpp_2m",
    "dpmpp_2m_sde",
    "dpmpp_2m_sde_gpu",
    "dpmpp_3m_sde",
    "dpmpp_3m_sde_gpu",
    "dpmpp_sde",
    "dpmpp_sde_gpu",
    "ddim",
    "uni_pc",
    "uni_pc_bh2",
    "res_multistep",
    "res_multistep_cfg_pp",
    "res_multistep_ancestral",
    "res_multistep_ancestral_cfg_pp",
    "ddpm",
    "lcm",
    "ipndm",
    "ipndm_v",
    "deis",
    "er_sde",
)

# Schedulers: Comfy samplers + schedulers; "simple" matches RDBT card
SCHEDULER_CHOICES: Final[tuple[str, ...]] = (
    "normal",
    "karras",
    "exponential",
    "sgm_uniform",
    "simple",
    "ddim_uniform",
    "beta",
    "linear_quadratic",
    "kl_optimal",
)

# Save prefix (matches old workflow)
SAVE_PREFIX: Final[str] = "AnimaRDBT_Gradio"

# Hub files downloaded at startup into the Comfy model layout.
ARTIFACTS: Final[list[tuple[str, str, str]]] = [
    (
        "circlestone-labs/Anima",
        "split_files/text_encoders/qwen_3_06b_base.safetensors",
        "text_encoders/qwen_3_06b_base.safetensors",
    ),
    (
        "circlestone-labs/Anima",
        "split_files/vae/qwen_image_vae.safetensors",
        "vae/qwen_image_vae.safetensors",
    ),
]

CIVITAI_RDBT: Final[tuple[str, str]] = (
    "https://civitai.com/api/download/models/2847666?type=Model&format=SafeTensor&size=full&fp=fp16",
    "diffusion_models/rdbt-anima-p3-v024f-16step-dmd2.safetensors",
)

MIN_SIZES: Final[dict[str, int]] = {
    "rdbt-anima-p3-v024f-16step-dmd2.safetensors": 3_000_000_000,
    "qwen_3_06b_base.safetensors": 500_000_000,
    "qwen_image_vae.safetensors": 100_000_000,
}

MAX_RETRIES: Final[int] = 8
BACKOFF_CAP_S: Final[int] = 60

def comfy_root() -> str:
    """
    Legacy: root directory; model weights for this Space are expected under <root>/models/.
    If unset, default ./ComfyUI so existing ANIMA_COMFY_ROOT=... Spaces keep the same on-disk layout.
    """
    d = os.environ.get("ANIMA_COMFY_ROOT", "").strip()
    if d:
        return os.path.abspath(d)
    return os.path.abspath(os.path.join(os.getcwd(), "ComfyUI"))


def model_artifacts_root() -> str:
    """
    Directory containing `diffusion_models/`, `text_encoders/`, `vae/` (same layout as before).
    If ANIMA_MODELS_ROOT is set, it must point to that *models* directory.
    If unset, uses <comfy_root>/models (e.g. ./ComfyUI/models).
    """
    d = os.environ.get("ANIMA_MODELS_ROOT", "").strip()
    if d:
        return os.path.abspath(d)
    return os.path.abspath(os.path.join(comfy_root(), "models"))


def skip_civitai() -> bool:
    return os.environ.get("SKIP_CIVITAI", "").strip() == "1"


def skip_startup_bootstrap() -> bool:
    """If true, skip import-time file bootstrap (tests, local dev without downloads)."""
    return os.environ.get("ANIMA_SKIP_STARTUP_BOOTSTRAP", "").strip() == "1"


def text_encoder_file_path() -> str:
    return os.path.join(model_artifacts_root(), "text_encoders", CLIP_NAME)


def vae_file_path() -> str:
    return os.path.join(model_artifacts_root(), "vae", VAE_NAME)


@dataclass(frozen=True)
class GenerationParams:
    prompt: str
    negative_prompt: str
    width: int
    height: int
    steps: int
    cfg: float
    batch_size: int
    sampler_name: str
    scheduler: str
    denoise: float
    seed: int
    warnings: tuple[str, ...] = ()