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