run.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / diffusers / transformers
import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from diffusers import CosmosTransformer3DModel

# Vendored from hugging-apps/anima-turbo-4step-demo (originally from
# CalamitousFelicitousness/Anima-sdnext-diffusers).
from pipeline import AnimaTextToImagePipeline
from modeling_llm_adapter import AnimaLLMAdapter  # noqa: F401  (needed for loading)

# ---------------------------------------------------------------------------
# Key mapping: ComfyUI Anima DiT checkpoint -> diffusers CosmosTransformer3DModel
# ---------------------------------------------------------------------------
# ComfyUI block keys:  net.blocks.{N}.{module}
# diffusers block keys: transformer_blocks.{N}.{module}

_BLOCK_MAP = {
    "self_attn.q_proj": "attn1.to_q",
    "self_attn.k_proj": "attn1.to_k",
    "self_attn.v_proj": "attn1.to_v",
    "self_attn.output_proj": "attn1.to_out.0",
    "self_attn.q_norm": "attn1.norm_q",
    "self_attn.k_norm": "attn1.norm_k",
    "cross_attn.q_proj": "attn2.to_q",
    "cross_attn.k_proj": "attn2.to_k",
    "cross_attn.v_proj": "attn2.to_v",
    "cross_attn.output_proj": "attn2.to_out.0",
    "cross_attn.q_norm": "attn2.norm_q",
    "cross_attn.k_norm": "attn2.norm_k",
    "mlp.layer1": "ff.net.0.proj",
    "mlp.layer2": "ff.net.2",
    "adaln_modulation_self_attn.1": "norm1.linear_1",
    "adaln_modulation_self_attn.2": "norm1.linear_2",
    "adaln_modulation_cross_attn.1": "norm2.linear_1",
    "adaln_modulation_cross_attn.2": "norm2.linear_2",
    "adaln_modulation_mlp.1": "norm3.linear_1",
    "adaln_modulation_mlp.2": "norm3.linear_2",
}

# Non-block keys (embedders / final layer). t_embedder is handled by regex
# below since its submodule index varies.
_TOP_MAP = {
    "net.x_embedder.proj.1.weight": "patch_embed.proj.weight",
    "net.t_embedding_norm.weight": "time_embed.norm.weight",
    "net.final_layer.adaln_modulation.1.weight": "norm_out.linear_1.weight",
    "net.final_layer.adaln_modulation.2.weight": "norm_out.linear_2.weight",
    "net.final_layer.linear.weight": "proj_out.weight",
}


def remap_comfyui_anima_checkpoint(state_dict):
    """Convert a ComfyUI-format Anima DiT checkpoint to diffusers keys.

    Skips net.llm_adapter.* (the 2.9B release only trained DiT layers; the
    adapter comes from the base diffusers repo) and net.pos_embedder.* buffers
    (diffusers computes rotary embeddings on the fly).
    """
    import re

    remapped = {}
    skipped = []
    for key, value in state_dict.items():
        if key.startswith("net.blocks."):
            rest = key[len("net.blocks."):]
            block_n, module_key = rest.split(".", 1)
            # Strip trailing .weight / .bias before lookup, re-append after
            if module_key.endswith(".weight"):
                core, suffix = module_key[:-len(".weight")], ".weight"
            elif module_key.endswith(".bias"):
                core, suffix = module_key[:-len(".bias")], ".bias"
            else:
                core, suffix = module_key, ""
            if core in _BLOCK_MAP:
                remapped[f"transformer_blocks.{block_n}.{_BLOCK_MAP[core]}{suffix}"] = value
            else:
                skipped.append(key)
            continue
        if key.startswith(("net.llm_adapter.", "net.pos_embedder.")):
            skipped.append(key)
            continue
        if key in _TOP_MAP:
            remapped[_TOP_MAP[key]] = value
            continue
        # net.t_embedder.{i}.linear_{1,2}.weight -> time_embed.t_embedder.linear_{1,2}.weight
        m = re.match(r"net\.t_embedder\.\d+\.(linear_[12]\.weight)", key)
        if m:
            remapped[f"time_embed.t_embedder.{m.group(1)}"] = value
            continue
        skipped.append(key)
    if skipped:
        print(f"  Skipped {len(skipped)} keys (llm_adapter/pos_embedder/etc.), e.g.: {skipped[:5]}")
    return remapped


# ---------------------------------------------------------------------------
# Load base Anima pipeline (diffusers format), then swap in the 2.9B DiT
# ---------------------------------------------------------------------------
print("Loading base Anima pipeline (text encoder, LLM adapter, VAE, scheduler)...")
pipe = AnimaTextToImagePipeline.from_pretrained(
    "CalamitousFelicitousness/Anima-sdnext-diffusers",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

print("Building 40-layer transformer for Anima-2.9B...")
transformer_config = dict(pipe.transformer.config)
transformer_config["num_layers"] = 40
transformer = CosmosTransformer3DModel.from_config(transformer_config)

print("Loading Anima-2.9B checkpoint from Gazingstars123/Anima-2.9B...")
ckpt_path = hf_hub_download(
    "Gazingstars123/Anima-2.9B",
    "Anima-2.9B-preview-v1.safetensors",
)
state_dict = remap_comfyui_anima_checkpoint(load_file(ckpt_path))
missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
# Biases are legitimately absent — the ComfyUI checkpoint is weight-only and
# the Cosmos architecture initializes them to zero.
missing = [k for k in missing if not k.endswith(".bias")]
if missing:
    print(f"  WARNING missing keys ({len(missing)}): {missing[:10]}")
if unexpected:
    print(f"  WARNING unexpected keys ({len(unexpected)}): {unexpected[:10]}")
if missing or unexpected:
    raise RuntimeError("Checkpoint key remap incomplete — see warnings above.")

pipe.transformer = transformer.to(dtype=torch.bfloat16)
pipe.to("cuda")
print("Pipeline loaded!")


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def _save_image(image) -> dict:
    """Serialize a PIL.Image as a JSON pointer the canvas can render."""
    import tempfile

    path = os.path.join(tempfile.gettempdir(), f"anima_{os.urandom(8).hex()}.png")
    image.save(path)
    return {
        "path": path,
        "url": f"/gradio_api/file={path}",
        "orig_name": "anima.png",
        "mime_type": "image/png",
    }


# Preset defaults from the Anima-2.9B model card recommendations:
#   Resolution: 812x1216 class (832x1216 used here — must be divisible by 16)
#   Steps: 28-50, CFG: 3.5-5
PRESET_HEIGHT = 1216
PRESET_WIDTH = 832
PRESET_STEPS = 28
PRESET_CFG = 4.0


def _estimate_duration(prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed) -> int:
    """Rough wall-clock estimate (seconds) for one Anima-2.9B call.

    Scaled by pixel count and steps; clamped so the queue isn't starved.
    See: https://huggingface.co/docs/hub/en/spaces-zerogpu#duration-management
    """
    height = int(height) if height else PRESET_HEIGHT
    width = int(width) if width else PRESET_WIDTH
    steps = int(num_inference_steps) if num_inference_steps else PRESET_STEPS
    pixels = max(height, 1) * max(width, 1)
    # ~0.5s/step at 1024^2 for a 2.9B model, linear-ish in pixels.
    per_step = 0.5 * (pixels / (1024 * 1024))
    seconds = steps * per_step
    return max(30, min(int(seconds) + 20, 180))


def _friendly_gpu_error(err: Exception) -> str:
    """Turn ZeroGPU's terse allocator rejections into a clear message."""
    msg = (str(err) or "").lower()
    capacity_hints = (
        "gpu limit", "reached its gpu limit", "gpu quota", "out of quota",
        "quota", "no gpu", "could not allocate", "gpu is busy", "too many",
        "concurrent",
    )
    if any(h in msg for h in capacity_hints):
        return (
            "⛔ This demo's shared GPU is at capacity right now — it's not a "
            "problem with your prompt or your account. Please wait a minute "
            "and retry; demand clears between bursts."
        )
    if "out of memory" in msg or "oom" in msg or "cuda" in msg:
        return (
            "💥 Image generation ran out of GPU memory. Try a smaller "
            "Height/Width or fewer Inference Steps, then retry."
        )
    return (
        "⚠️ Image generation failed. Please try again in a moment — if it "
        "keeps happening, simplify your prompt or lower the resolution."
    )


@spaces.GPU(duration=_estimate_duration)
def _generate_image_gpu(
    prompt: str,
    height: int = PRESET_HEIGHT,
    width: int = PRESET_WIDTH,
    num_inference_steps: int = PRESET_STEPS,
    guidance_scale: float = PRESET_CFG,
    seed: int = 0,
    randomize_seed: bool = True,
):
    """The GPU-decorated worker. Runs only under a ZeroGPU allocation."""
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    # Fall back to presets for any unset (None) inputs
    height = int(height) if height else PRESET_HEIGHT
    width = int(width) if width else PRESET_WIDTH
    num_inference_steps = int(num_inference_steps) if num_inference_steps else PRESET_STEPS
    guidance_scale = float(guidance_scale) if guidance_scale else PRESET_CFG

    if randomize_seed or seed is None:
        seed = torch.randint(0, 2**31 - 1, (1,)).item()

    generator = torch.Generator("cuda").manual_seed(int(seed))
    image = pipe(
        prompt=prompt,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]

    return _save_image(image), int(seed)


def generate_image(
    prompt: str,
    height: int = PRESET_HEIGHT,
    width: int = PRESET_WIDTH,
    num_inference_steps: int = PRESET_STEPS,
    guidance_scale: float = PRESET_CFG,
    seed: int = 0,
    randomize_seed: bool = True,
):
    """Workflow-facing wrapper around the GPU worker.

    Bound to the canvas as a `fn` operator node. Catches ZeroGPU allocator
    rejections and rewords them into user-facing messages.
    """
    try:
        return _generate_image_gpu(
            prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed
        )
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(_friendly_gpu_error(e)) from e


# The workflow (workflow.json) wires `generate_image` as a `fn` operator:
#   Prompt, Height, Width, Inference Steps, CFG, Seed, Randomize Seed ─▶
#   generate_image (fn operator, kind="fn") ─▶ Output Image, Seed Used
# Reference nodes carry the model card's recommended presets as default values.
demo = gr.Workflow(
    graph="workflow.json",
    bind={"generate_image": generate_image},
)

if __name__ == "__main__":
    demo.launch()