import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

import spaces  # MUST be first (after os.environ setup)
import random
import time
from pathlib import Path
from typing import Optional
from dataclasses import dataclass

import torch
import torch.nn as nn
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download, snapshot_download

# ─── model constants ─────────────────────────────────────────────
BASE_MODEL_ID = "circlestone-labs/Anima"
DIFFUSION_FILE = "split_files/diffusion_models/anima-base-v1.0.safetensors"
TEXT_ENCODER_FILE = "split_files/text_encoders/qwen_3_06b_base.safetensors"
VAE_FILE = "split_files/vae/qwen_image_vae.safetensors"
QWEN_TOKENIZER_ID = "Qwen/Qwen3-0.6B"
T5_TOKENIZER_ID = "google/t5-v1_1-xxl"
CONTROLNET_MODEL_ID = "TaihoC/Anima-ControlNet-VACE-Depth"
CONTROLNET_FILE = "anima-vace-depth.safetensors"

DEFAULT_POSITIVE_PREFIX = "masterpiece, best quality, score_7, safe, "
DEFAULT_NEGATIVE = "worst quality, low quality, score_1, score_2, score_3, artist name"

# ─── paths ───────────────────────────────────────────────────────
def _select_writable_dir(env_name, candidates):
    existing = os.getenv(env_name)
    if existing:
        candidates = [existing] + [c for c in candidates if c != existing]
    for c in candidates:
        try:
            p = Path(c)
            p.mkdir(parents=True, exist_ok=True)
            test = p / ".write_test"
            test.write_text("ok")
            test.unlink(missing_ok=True)
            return str(p)
        except Exception:
            continue
    fb = Path("/tmp") / env_name.lower()
    fb.mkdir(parents=True, exist_ok=True)
    return str(fb)

HF_HOME = _select_writable_dir("HF_HOME", ["/data/.cache/huggingface", "/tmp/.cache/huggingface"])
LOCAL_MODEL_DIR = _select_writable_dir("ANIMA_LOCAL_MODEL_DIR", ["/data/models/anima-vace", "/tmp/models/anima-vace"])
os.environ["HF_HOME"] = HF_HOME
os.environ.setdefault("HF_HUB_CACHE", str(Path(HF_HOME) / "hub"))
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", str(Path(LOCAL_MODEL_DIR) / "diffsynth"))

_HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")

# ─── DiffSynth imports (after spaces) ────────────────────────────
from diffsynth.pipelines.anima_image import AnimaImagePipeline, ModelConfig
from diffsynth.models.anima_dit import Block as AnimaBlock, PatchEmbed

# ─── VACE ControlNet module ──────────────────────────────────────
# Ported from TaihoC/ComfyUI-Advanced-ControlNet-Anima (feat/anima-vace-controlnet)
# The VACE architecture: spaced-block-duplication + zero-conv, where
# each control block mirrors the Anima DiT Block and its output hint
# is added to the corresponding base DiT block's residual stream.

_HINTS_KEY = "anima_vace_hints"


class VACEControlBlock(nn.Module):
    """One block of the VACE control branch: a copy of the Anima DiT Block
    plus zero-conv projections. Block 0 additionally fuses the control
    latent with the base model's patch-embedded input via before_proj."""

    def __init__(self, block: nn.Module, x_dim: int, block_id: int = 0, use_after_proj: bool = True):
        super().__init__()
        self.block = block
        self.block_id = block_id
        self.x_dim = x_dim
        self.use_after_proj = use_after_proj
        if block_id == 0:
            self.before_proj = nn.Linear(x_dim, x_dim)
        if use_after_proj:
            self.after_proj = nn.Linear(x_dim, x_dim)

    def init_weights(self):
        if hasattr(self.block, "init_weights"):
            self.block.init_weights()
        if self.block_id == 0:
            nn.init.zeros_(self.before_proj.weight)
            nn.init.zeros_(self.before_proj.bias)
        if self.use_after_proj:
            nn.init.zeros_(self.after_proj.weight)
            nn.init.zeros_(self.after_proj.bias)

    def forward(self, c, x_base, **block_kwargs):
        all_c = []
        if self.block_id == 0:
            c = self.before_proj(c) + x_base
        elif self.use_after_proj:
            all_c = list(torch.unbind(c))
            c = all_c.pop(-1)
        c = self.block(c, **block_kwargs)
        if self.use_after_proj:
            c_skip = self.after_proj(c)
            all_c = all_c + [c_skip, c]
            c = torch.stack(all_c)
        else:
            c_skip = c
        return c_skip, c


class VACEControlNetAnima(nn.Module):
    """Standalone VACE control branch: PatchEmbed for the control latent
    plus a ModuleList of VACEControlBlock. forward() returns per-base-block
    hint tensors keyed by base block index."""

    def __init__(
        self,
        model_channels=2048,
        num_heads=16,
        num_blocks=28,
        crossattn_emb_channels=1024,
        patch_spatial=2,
        patch_temporal=1,
        latent_channels=16,
        vace_block_every_n=7,
        condition_strategy="spaced",
        use_after_proj=True,
        use_adaln_lora=True,
        adaln_lora_dim=256,
        concat_padding_mask=True,
    ):
        super().__init__()
        self.model_channels = model_channels
        self.num_blocks = num_blocks
        self.latent_channels = latent_channels
        self.vace_block_every_n = vace_block_every_n
        self.condition_strategy = condition_strategy
        self.concat_padding_mask = concat_padding_mask

        if condition_strategy == "spaced":
            self.control_layers = list(range(0, num_blocks, vace_block_every_n))
            self.control_layers_mapping = {i: n for n, i in enumerate(self.control_layers)}
        else:
            n_ctrl = max(1, num_blocks // vace_block_every_n)
            self.control_layers = list(range(0, n_ctrl))
            self.control_layers_mapping = {i: i for i in range(n_ctrl)}
        assert 0 in self.control_layers, "first control block must map to base block 0"

        embed_in_channels = latent_channels + (1 if concat_padding_mask else 0)
        self.control_embedder = PatchEmbed(
            spatial_patch_size=patch_spatial,
            temporal_patch_size=patch_temporal,
            in_channels=embed_in_channels,
            out_channels=model_channels,
            operations=torch.nn,
        )

        self.control_blocks = nn.ModuleList([
            VACEControlBlock(
                block=AnimaBlock(
                    x_dim=model_channels,
                    context_dim=crossattn_emb_channels,
                    num_heads=num_heads,
                    mlp_ratio=4.0,
                    use_adaln_lora=use_adaln_lora,
                    adaln_lora_dim=adaln_lora_dim,
                    operations=torch.nn,
                ),
                x_dim=model_channels,
                block_id=i,
                use_after_proj=use_after_proj,
            )
            for i in range(len(self.control_layers))
        ])

    def forward(self, control_latent, x_base_tokens, emb_B_T_D, crossattn_emb,
                rope_emb=None, adaln_lora=None, extra_per_block_pos_emb=None):
        cdtype = control_latent.dtype
        x_base_tokens = x_base_tokens.to(cdtype)
        emb_B_T_D = emb_B_T_D.to(cdtype)
        crossattn_emb = crossattn_emb.to(cdtype)
        if adaln_lora is not None:
            adaln_lora = adaln_lora.to(cdtype)
        if extra_per_block_pos_emb is not None:
            extra_per_block_pos_emb = extra_per_block_pos_emb.to(cdtype)

        if self.concat_padding_mask and control_latent.shape[1] == self.latent_channels:
            zeros_mask = torch.zeros(
                control_latent.shape[0], 1, *control_latent.shape[2:],
                device=control_latent.device, dtype=control_latent.dtype,
            )
            control_latent = torch.cat([control_latent, zeros_mask], dim=1)

        c = self.control_embedder(control_latent)

        rope_for_blocks = rope_emb
        if rope_emb is not None and rope_emb.ndim == 4:
            rope_for_blocks = rope_emb.unsqueeze(1).unsqueeze(0)

        block_kwargs = dict(
            emb_B_T_D=emb_B_T_D,
            crossattn_emb=crossattn_emb,
            rope_emb_L_1_1_D=rope_for_blocks,
            adaln_lora_B_T_3D=adaln_lora,
            extra_per_block_pos_emb=extra_per_block_pos_emb,
        )

        hints = {}
        for i, cb in enumerate(self.control_blocks):
            c_skip, c = cb(c, x_base_tokens, **block_kwargs)
            hints[self.control_layers[i]] = c_skip
        return hints


# ─── asset download ──────────────────────────────────────────────
@dataclass
class AssetPaths:
    diffusion: str
    text_encoder: str
    vae: str
    qwen_tokenizer_dir: str
    t5_tokenizer_dir: str
    controlnet_ckpt: str


_TOKENIZER_ALLOW = [
    "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json",
    "added_tokens.json", "vocab.json", "merges.txt", "config.json",
    "*.model", "*.tiktoken", "*.jinja",
]


def _repo_local_dir(repo_id, suffix=""):
    safe = repo_id.replace("/", "--")
    if suffix:
        safe = f"{safe}--{suffix}"
    p = Path(LOCAL_MODEL_DIR) / safe
    p.mkdir(parents=True, exist_ok=True)
    return p


def _download_file(repo_id, filename, local_dir):
    return hf_hub_download(repo_id=repo_id, filename=filename, local_dir=str(local_dir), token=_HF_TOKEN)


def _has_tokenizer_files(path):
    if not path.is_dir():
        return False
    markers = ("tokenizer.json", "tokenizer.model", "spiece.model", "vocab.json", "merges.txt")
    return any((path / m).exists() for m in markers)


def _download_tokenizer(repo_id, subfolder="", suffix="tokenizer"):
    cleaned = subfolder.strip("/")
    local_dir = _repo_local_dir(repo_id, suffix if not cleaned else f"{suffix}--{cleaned}")
    if cleaned:
        allow = [f"{cleaned}/{p}" for p in _TOKENIZER_ALLOW]
    else:
        allow = _TOKENIZER_ALLOW
    snapshot_dir = Path(snapshot_download(
        repo_id=repo_id, allow_patterns=allow, local_dir=str(local_dir), token=_HF_TOKEN
    ))
    candidates = []
    if cleaned:
        candidates.append(snapshot_dir / cleaned)
    candidates.append(snapshot_dir)
    for c in candidates:
        if _has_tokenizer_files(c):
            return str(c)
    raise RuntimeError(f"No usable tokenizer files found in {repo_id}")


def _download_assets():
    print(f"[startup] Downloading model assets to {LOCAL_MODEL_DIR}", flush=True)
    start = time.time()
    anima_dir = _repo_local_dir(BASE_MODEL_ID)
    diffusion = _download_file(BASE_MODEL_ID, DIFFUSION_FILE, anima_dir)
    text_encoder = _download_file(BASE_MODEL_ID, TEXT_ENCODER_FILE, anima_dir)
    vae = _download_file(BASE_MODEL_ID, VAE_FILE, anima_dir)
    qwen_dir = _download_tokenizer(QWEN_TOKENIZER_ID, suffix="qwen-tokenizer")
    t5_dir = _download_tokenizer(T5_TOKENIZER_ID, suffix="t5-tokenizer")
    controlnet_ckpt = _download_file(CONTROLNET_MODEL_ID, CONTROLNET_FILE, _repo_local_dir(CONTROLNET_MODEL_ID))
    elapsed = time.time() - start
    print(f"[startup] Downloads ready in {elapsed:.1f}s", flush=True)
    return AssetPaths(
        diffusion=diffusion, text_encoder=text_encoder, vae=vae,
        qwen_tokenizer_dir=qwen_dir, t5_tokenizer_dir=t5_dir,
        controlnet_ckpt=controlnet_ckpt,
    )


# ─── pipeline state ──────────────────────────────────────────────
_PIPE: Optional[AnimaImagePipeline] = None
_VACE: Optional[VACEControlNetAnima] = None
_ASSETS: Optional[AssetPaths] = None
_DEPTH_PROCESSOR = None


def _ensure_assets():
    global _ASSETS
    if _ASSETS is None:
        _ASSETS = _download_assets()
    return _ASSETS


def _load_vace(ckpt_path):
    """Load the VACE ControlNet from a safetensors checkpoint."""
    from safetensors.torch import safe_open, load_file

    meta = {}
    try:
        with safe_open(ckpt_path, framework="pt") as f:
            meta = dict(f.metadata() or {})
    except Exception:
        pass

    model_channels = int(meta.get("vace.model_channels", "2048"))
    num_blocks = int(meta.get("vace.num_blocks", "28"))
    latent_channels = int(meta.get("vace.in_channels", "16"))
    vace_block_every_n = int(meta.get("vace.block_every_n", "7"))
    use_after_proj = meta.get("vace.use_after_proj", "true").lower() == "true"

    sd = load_file(ckpt_path)

    max_ctrl_idx = -1
    for k in sd:
        if k.startswith("control_blocks."):
            max_ctrl_idx = max(max_ctrl_idx, int(k.split(".")[1]))
    if max_ctrl_idx >= 0:
        n_ctrl = max_ctrl_idx + 1
        vace_block_every_n = num_blocks // n_ctrl if n_ctrl > 0 else vace_block_every_n

    model = VACEControlNetAnima(
        model_channels=model_channels,
        num_heads=16,
        num_blocks=num_blocks,
        crossattn_emb_channels=1024,
        patch_spatial=2,
        patch_temporal=1,
        latent_channels=latent_channels,
        vace_block_every_n=vace_block_every_n,
        condition_strategy="spaced",
        use_after_proj=use_after_proj,
        use_adaln_lora=True,
        adaln_lora_dim=256,
        concat_padding_mask=True,
    )
    info = model.load_state_dict(sd, strict=False)
    if info.missing_keys:
        print(f"[vace] missing keys ({len(info.missing_keys)}): {info.missing_keys[:5]}...", flush=True)
    if info.unexpected_keys:
        print(f"[vace] unexpected keys ({len(info.unexpected_keys)}): {info.unexpected_keys[:5]}...", flush=True)
    model = model.to(torch.bfloat16).to("cuda")
    model.eval()
    return model


def _install_vace_hooks(pipe, vace):
    """Register forward hooks on base DiT blocks to inject VACE hints."""
    dit = pipe.dit
    for base_idx in vace.control_layers_mapping:
        block = dit.blocks[base_idx]

        def make_hook(bi):
            def hook(module, args, kwargs, output):
                to = kwargs.get("transformer_options") or {}
                hints = to.get(_HINTS_KEY)
                if not hints:
                    return output
                hint = hints.get(bi)
                if hint is None:
                    return output
                return output + hint.to(dtype=output.dtype)
            return hook

        block.register_forward_hook(make_hook(base_idx), with_kwargs=True)


def _load_pipe():
    global _PIPE, _VACE
    if _PIPE is not None:
        return _PIPE, _VACE
    assets = _ensure_assets()
    torch.set_float32_matmul_precision("high")
    print("[startup] Loading Anima pipeline", flush=True)
    _PIPE = AnimaImagePipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda",
        model_configs=[
            ModelConfig(path=assets.diffusion),
            ModelConfig(path=assets.text_encoder),
            ModelConfig(path=assets.vae),
        ],
        tokenizer_config=ModelConfig(path=assets.qwen_tokenizer_dir),
        tokenizer_t5xxl_config=ModelConfig(path=assets.t5_tokenizer_dir),
    )
    print("[startup] Anima pipeline loaded", flush=True)
    print("[startup] Loading VACE ControlNet", flush=True)
    _VACE = _load_vace(assets.controlnet_ckpt)
    _install_vace_hooks(_PIPE, _VACE)
    print("[startup] VACE ControlNet loaded and hooked", flush=True)
    return _PIPE, _VACE


# ─── depth preprocessor ──────────────────────────────────────────
def _load_depth_processor():
    global _DEPTH_PROCESSOR
    if _DEPTH_PROCESSOR is not None:
        return _DEPTH_PROCESSOR
    from transformers import pipeline as hf_pipeline
    # The model card recommends DepthAnything V2; use the Large variant on GPU
    # for crisp depth edges (Small produces soft, low-contrast maps that blur
    # the control signal). Loaded lazily inside the @spaces.GPU context.
    _DEPTH_PROCESSOR = hf_pipeline(
        task="depth-estimation",
        model="depth-anything/Depth-Anything-V2-Large-hf",
        device="cuda",
        torch_dtype=torch.float16,
    )
    print("[startup] Depth processor loaded", flush=True)
    return _DEPTH_PROCESSOR


def _extract_depth(image):
    """Extract a depth map from an input image using Depth Anything V2."""
    processor = _load_depth_processor()
    if image.mode != "RGB":
        image = image.convert("RGB")
    result = processor(image)
    return result["depth"]


# ─── VACE control injection ─────────────────────────────────────
def _run_vace_control(pipe, vace, control_latent, dit, x_noisy, timestep, context, strength):
    """Run the VACE control branch and return per-block hints dict."""
    with torch.no_grad():
        x_5d = x_noisy.unsqueeze(2) if x_noisy.ndim == 4 else x_noisy
        x_B_T_H_W_D, rope_emb, extra_pos_emb = dit.prepare_embedded_sequence(x_5d)

        t = timestep / 1000
        if t.ndim == 1:
            t = t.unsqueeze(1)
        # Match the DiT forward: convert timesteps to x dtype, then compute embedding
        timesteps_emb = dit.t_embedder[0](t.to(x_B_T_H_W_D.dtype))
        # Timesteps may output float32; cast to match TimestepEmbedding weights
        timesteps_emb = timesteps_emb.to(x_B_T_H_W_D.dtype)
        t_emb, adaln_lora = dit.t_embedder[1](timesteps_emb)
        t_emb = dit.t_embedding_norm(t_emb)

        cl = control_latent
        if cl.ndim == 4:
            cl = cl.unsqueeze(2)
        if cl.shape[1] != vace.latent_channels:
            pipe.load_models_to_device(['vae'])
            cl = pipe.vae.encode(cl.movedim(1, -1))
            if cl.ndim == 4:
                cl = cl.unsqueeze(2)

        hints = vace(
            control_latent=cl.to(device=x_B_T_H_W_D.device, dtype=torch.bfloat16),
            x_base_tokens=x_B_T_H_W_D,
            emb_B_T_D=t_emb,
            crossattn_emb=context,
            rope_emb=rope_emb,
            adaln_lora=adaln_lora,
            extra_per_block_pos_emb=extra_pos_emb,
        )
        return {ci: hint * strength for ci, hint in hints.items()}


# ─── Gradio inference ───────────────────────────────────────────
@spaces.GPU(duration=180)
def generate(
    input_image,
    prompt: str,
    negative_prompt: str,
    controlnet_strength: float,
    width: int,
    height: int,
    steps: int,
    cfg_scale: float,
    seed: int,
    sigma_shift: float = 5.0,
    progress: gr.Progress = gr.Progress(track_tqdm=False),
):
    """Generate an anime-style image conditioned on the depth map of an input image.

    Args:
        input_image: Image to extract depth from for conditioning.
        prompt: Text prompt describing the desired output.
        negative_prompt: What to avoid in the output.
        controlnet_strength: Strength of depth conditioning (0-2).
        width: Output image width (512-1536, rounded to 16).
        height: Output image height (512-1536, rounded to 16).
        steps: Number of denoising steps.
        cfg_scale: Classifier-free guidance scale.
        seed: Random seed (-1 for random).
        sigma_shift: Flow-matching timestep shift (the ControlNet was trained
            and evaluated at 5.0 per the model card).
    """
    if input_image is None:
        raise gr.Error("Please provide an input image for depth extraction.")
    if not prompt or not prompt.strip():
        prompt = "1girl, solo, long silver hair, blue eyes, blue dress, underwater, floating hair, refraction, portrait"

    prompt = DEFAULT_POSITIVE_PREFIX + prompt.strip()
    negative_prompt = (negative_prompt or DEFAULT_NEGATIVE).strip()
    # Output follows the INPUT image's aspect ratio; the width/height sliders
    # set the pixel budget (target area). This keeps the depth map undistorted.
    import math as _math
    target_area = float(width) * float(height)
    aspect = input_image.width / input_image.height
    width = int(round(_math.sqrt(target_area * aspect) / 16) * 16)
    height = int(round(_math.sqrt(target_area / aspect) / 16) * 16)
    width = max(512, min(1536, width))
    height = max(512, min(1536, height))
    steps = int(max(10, min(60, steps)))
    cfg_scale = float(max(1.0, min(8.0, cfg_scale)))
    sigma_shift = float(max(1.0, min(8.0, sigma_shift)))
    if seed < 0:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)

    pipe, vace = _load_pipe()

    # Extract depth map from input image
    progress(0.1, desc="Extracting depth map")
    depth_map = _extract_depth(input_image)
    depth_map = depth_map.resize((width, height), Image.LANCZOS)
    # Ensure 3-channel RGB for VAE encoding
    if depth_map.mode != "RGB":
        depth_map = depth_map.convert("RGB")

    # VAE-encode the depth map into latent space
    progress(0.2, desc="Encoding depth latent")
    depth_tensor = pipe.preprocess_image(depth_map).to(device=pipe.device, dtype=pipe.torch_dtype)
    with torch.no_grad():
        pipe.load_models_to_device(['vae'])
        control_latent = pipe.vae.encode(depth_tensor.unsqueeze(2), device=pipe.device).squeeze(2)

    # Monkey-patch model_fn to inject VACE control
    original_model_fn = pipe.model_fn
    _vace_state = {"control_latent": control_latent, "strength": controlnet_strength, "vace": vace}

    def patched_model_fn(dit=None, latents=None, timestep=None, prompt_emb=None,
                         t5xxl_ids=None, use_gradient_checkpointing=False,
                         use_gradient_checkpointing_offload=False, **kwargs):
        # The base AnimaDiT runs the raw Qwen hidden states through its
        # LLMAdapter (preprocess_text_embeds, fused with the T5-XXL token
        # embeddings) before the blocks ever see them as crossattn_emb.
        # ComfyUI pre-applies this in model_base.Anima.extra_conds, so the
        # reference VACE branch receives the ADAPTED context — the control
        # blocks were trained against it. Adapt once here and feed the same
        # tensor to both the control branch and the base DiT (t5xxl_ids=None
        # skips re-adaptation inside AnimaDiT.forward).
        context = prompt_emb
        if t5xxl_ids is not None:
            context = dit.preprocess_text_embeds(prompt_emb, t5xxl_ids)

        hints = _run_vace_control(
            pipe, _vace_state["vace"], _vace_state["control_latent"],
            dit, latents, timestep, context, _vace_state["strength"],
        )
        transformer_options = kwargs.get("transformer_options", {})
        transformer_options[_HINTS_KEY] = hints
        kwargs["transformer_options"] = transformer_options

        latents_5d = latents.unsqueeze(2)
        timestep_scaled = timestep / 1000
        model_output = dit(
            x=latents_5d,
            timesteps=timestep_scaled,
            context=context,
            t5xxl_ids=None,
            use_gradient_checkpointing=use_gradient_checkpointing,
            use_gradient_checkpointing_offload=use_gradient_checkpointing_offload,
            **kwargs,
        )
        return model_output.squeeze(2)

    pipe.model_fn = patched_model_fn
    try:
        progress(0.3, desc="Generating")
        with torch.inference_mode():
            image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                cfg_scale=cfg_scale,
                height=height,
                width=width,
                seed=seed,
                num_inference_steps=steps,
                sigma_shift=sigma_shift,
            )
    finally:
        pipe.model_fn = original_model_fn

    info = (
        f"**Seed:** `{seed}`  \n"
        f"**Size:** `{width}×{height}`  \n"
        f"**Steps / CFG / shift:** `{steps}` / `{cfg_scale}` / `{sigma_shift}`  \n"
        f"**ControlNet strength:** `{controlnet_strength}`"
    )
    return image, depth_map, info


# ─── startup ─────────────────────────────────────────────────────
# On ZeroGPU, CUDA is only available inside @spaces.GPU decorated functions,
# so we only download assets at startup; the pipeline loads lazily on first call.
try:
    _ASSETS = _download_assets()
    print("[startup] Assets downloaded", flush=True)
except Exception as e:
    print(f"[startup] Download failed: {e}", flush=True)
    _ASSETS = None


# ─── Gradio UI ───────────────────────────────────────────────────
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Anima ControlNet VACE Depth") as demo:
    gr.Markdown(
        "# Anima ControlNet VACE Depth\n"
        "Generate anime-style images conditioned on the depth structure of an input image. "
        "Uses the [Anima](https://huggingface.co/circlestone-labs/Anima) base model with the "
        "[VACE Depth ControlNet](https://huggingface.co/TaihoC/Anima-ControlNet-VACE-Depth) adapter.\n\n"
        "Upload an image, write a prompt, and the model will generate an anime-style image "
        "that follows the depth structure of your input."
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Input image (for depth extraction)", type="pil")
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                value="1girl, solo, long silver hair, blue eyes, blue dress, underwater, floating hair, refraction, portrait, looking at viewer",
            )
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                lines=2,
                value=DEFAULT_NEGATIVE,
            )
            controlnet_strength = gr.Slider(0.0, 2.0, value=1.0, step=0.05,
                                           label="ControlNet strength")
            generate_btn = gr.Button("Generate", variant="primary")

            with gr.Accordion("Advanced settings", open=False):
                gr.Markdown(
                    "*The output matches the input image's aspect ratio; "
                    "width × height set the pixel budget.*"
                )
                with gr.Row():
                    width = gr.Slider(512, 1536, value=1024, step=16, label="Width (budget)")
                    height = gr.Slider(512, 1536, value=1024, step=16, label="Height (budget)")
                with gr.Row():
                    steps = gr.Slider(10, 60, value=50, step=1, label="Steps")
                    cfg_scale = gr.Slider(1.0, 8.0, value=3.5, step=0.1, label="CFG scale")
                with gr.Row():
                    sigma_shift = gr.Slider(1.0, 8.0, value=5.0, step=0.5, label="Flow shift",
                                            info="ControlNet trained at 5.0 (model card)")
                    seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)")

        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated image", type="pil")
            depth_image = gr.Image(label="Depth map (extracted)", type="pil")
            info = gr.Markdown()

    inputs_list = [input_image, prompt, negative_prompt, controlnet_strength,
                   width, height, steps, cfg_scale, seed, sigma_shift]
    generate_btn.click(generate, inputs=inputs_list, outputs=[output_image, depth_image, info],
                       show_progress=True)

    gr.Examples(
        examples=[
            ["examples/depth_landscape.jpg",
             "year 2025, newest, highres, safe, 1girl, standing in a misty forest, sunlight through trees, fantasy landscape, detailed background",
             DEFAULT_NEGATIVE, 1.0, 1024, 1024, 50, 3.5, 42, 5.0],
            ["examples/depth_anime_1.jpg",
             "year 2025, newest, masterpiece, best quality, score_7, safe, 1girl, silver hair, red eyes, gothic dress, standing pose, detailed background",
             DEFAULT_NEGATIVE, 1.0, 1024, 1024, 50, 3.5, 12345, 5.0],
            ["examples/depth_astronaut.jpg",
             "year 2025, newest, highres, safe, 1boy, space suit, helmet, standing on alien planet, stars, nebula, sci-fi landscape, dramatic lighting",
             DEFAULT_NEGATIVE, 0.8, 1024, 1024, 50, 3.5, -1, 5.0],
        ],
        inputs=inputs_list,
        outputs=[output_image, depth_image, info],
        fn=generate,
        cache_examples=False,
        run_on_click=True,
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(
        ssr_mode=False,
        theme=gr.themes.Citrus(),
        css=CSS,
    )