src/config.py

"""
Runtime configuration for Z-Anime Distill-8-Step (FP8) Gradio + embedded ComfyUI.
Defaults align with SeeSee21/Z-Anime model card and official workflow guidance.
"""
from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Final

# --- Z-Anime Distill-8-Step defaults ---
DEFAULT_WIDTH: Final[int] = 1024
DEFAULT_HEIGHT: Final[int] = 1024
DEFAULT_STEPS: Final[int] = 8
DEFAULT_CFG: Final[float] = 1.0
DEFAULT_SAMPLER: Final[str] = "euler_ancestral"
DEFAULT_SCHEDULER: Final[str] = "beta"
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

# ModelSamplingAuraFlow.shift — matches workflows/Z-Anime-Workflow-v1.json (⚙️ Model Sampling)
DEFAULT_AURA_SHIFT: Final[float] = 3.5

# Model filenames (standard split, FP8)
DISTILL_UNET_NAME: Final[str] = "z-anime-distill-8step-fp8.safetensors"
CLIP_NAME: Final[str] = "qwen_3_4b-fp8.safetensors"
CLIP_TYPE: Final[str] = "qwen_image"
VAE_NAME: Final[str] = "ae.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] = 1.5
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

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",
)

SCHEDULER_CHOICES: Final[tuple[str, ...]] = (
    "normal",
    "karras",
    "exponential",
    "sgm_uniform",
    "simple",
    "ddim_uniform",
    "beta",
    "linear_quadratic",
    "kl_optimal",
)

SAVE_PREFIX: Final[str] = "ZAnimeDistill8_Gradio"

HUB_MODEL_REPO: Final[str] = "SeeSee21/Z-Anime"

# Hub files downloaded at startup into the Comfy model layout.
ARTIFACTS: Final[list[tuple[str, str, str]]] = [
    (
        HUB_MODEL_REPO,
        f"diffusion_models/{DISTILL_UNET_NAME}",
        f"diffusion_models/{DISTILL_UNET_NAME}",
    ),
    (
        HUB_MODEL_REPO,
        f"text_encoder/{CLIP_NAME}",
        f"text_encoders/{CLIP_NAME}",
    ),
    (
        HUB_MODEL_REPO,
        f"vae/{VAE_NAME}",
        f"vae/{VAE_NAME}",
    ),
]

MIN_SIZES: Final[dict[str, int]] = {
    DISTILL_UNET_NAME: 5_000_000_000,
    CLIP_NAME: 3_000_000_000,
    VAE_NAME: 50_000_000,
}

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


def comfy_root() -> str:
    """ComfyUI source tree root (contains nodes.py). Default ./ComfyUI."""
    d = os.environ.get("ZANIME_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/, clip/.
    If ZANIME_MODELS_ROOT is set, it must point to that *models* directory.
    """
    d = os.environ.get("ZANIME_MODELS_ROOT", "").strip()
    if d:
        return os.path.abspath(d)
    return os.path.abspath(os.path.join(comfy_root(), "models"))


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


def disable_aura_sampling() -> bool:
    """If true, skip ModelSamplingAuraFlow (fallback; may diverge from official workflow)."""
    return os.environ.get("ZANIME_DISABLE_AURA_SAMPLING", "").strip() == "1"


def aura_shift() -> float:
    raw = os.environ.get("ZANIME_AURA_SHIFT", "").strip()
    if not raw:
        return float(DEFAULT_AURA_SHIFT)
    try:
        return float(raw)
    except ValueError:
        return float(DEFAULT_AURA_SHIFT)


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)


def clip_mirror_file_path() -> str:
    """Optional mirror under models/clip/ for layouts that expect CLIP-named folders."""
    return os.path.join(model_artifacts_root(), "clip", CLIP_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, ...] = ()