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()