run.py

import os

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

# spaces MUST be imported before torch, diffusers, or transformers
try:
    import spaces
except ImportError:
    class spaces:
        @staticmethod
        def GPU(duration=60):
            def decorator(fn):
                return fn
            return decorator

import torch
import gradio as gr
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from diffusers import CosmosTransformer3DModel

from pipeline import AnimaTextToImagePipeline
from modeling_llm_adapter import AnimaLLMAdapter  # noqa: F401

_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",
}

_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):
    import re
    remapped = {}
    for key, value in state_dict.items():
        if key.startswith("net.blocks."):
            rest = key[len("net.blocks."):]
            block_n, module_key = rest.split(".", 1)
            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
            continue
        if key.startswith(("net.llm_adapter.", "net.pos_embedder.")):
            continue
        if key in _TOP_MAP:
            remapped[_TOP_MAP[key]] = value
            continue
        m = re.match(r"net\.t_embedder\.\d+\.(linear_[12]\.weight)", key)
        if m:
            remapped[f"time_embed.t_embedder.{m.group(1)}"] = value
    return remapped


print("Loading base Anima pipeline...")
pipe = AnimaTextToImagePipeline.from_pretrained(
    "CalamitousFelicitousness/Anima-sdnext-diffusers",
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    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...")
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))
transformer.load_state_dict(state_dict, strict=False)

pipe.transformer = transformer.to(dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32)
if torch.cuda.is_available():
    pipe.to("cuda")
print("Pipeline ready!")


def _save_image(image) -> dict:
    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_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:
    h = int(height) if height else PRESET_HEIGHT
    w = int(width) if width else PRESET_WIDTH
    steps = int(num_inference_steps) if num_inference_steps else PRESET_STEPS
    pixels = max(h, 1) * max(w, 1)
    per_step = 0.5 * (pixels / (1024 * 1024))
    return max(30, min(int(steps * per_step) + 20, 180))


@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,
):
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    height = (int(height) if height else PRESET_HEIGHT) // 16 * 16
    width = (int(width) if width else PRESET_WIDTH) // 16 * 16
    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()
    else:
        seed = int(seed)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    generator = torch.Generator(device).manual_seed(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), 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,
):
    return _generate_image_gpu(prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed)


demo = gr.Workflow(
    graph="workflow.json",
    bind={"generate_image": generate_image},
)

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