import os import re import time import tempfile from typing import Tuple os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") # spaces MUST be imported before torch, diffusers, or transformers try: import spaces HAS_SPACES = True except ImportError: HAS_SPACES = False 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 (needed for module registry) # --------------------------------------------------------------------------- # Key mapping: ComfyUI Anima DiT checkpoint -> diffusers CosmosTransformer3DModel # --------------------------------------------------------------------------- _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): """Convert a ComfyUI-format Anima DiT checkpoint to diffusers keys.""" 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) 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 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)} non-DiT keys (adapter/pos_embedder)") return remapped # --------------------------------------------------------------------------- # Load base Anima pipeline & swap in 40-layer Anima-2.9B DiT weights # --------------------------------------------------------------------------- print("Loading base Anima pipeline...") device_target = "cuda" if torch.cuda.is_available() else "cpu" dtype_target = torch.bfloat16 if torch.cuda.is_available() else torch.float32 pipe = AnimaTextToImagePipeline.from_pretrained( "CalamitousFelicitousness/Anima-sdnext-diffusers", torch_dtype=dtype_target, trust_remote_code=True, ) print("Building 40-layer Cosmos transformer for Anima-2.9B...") transformer_config = dict(pipe.transformer.config) transformer_config["num_layers"] = 40 transformer = CosmosTransformer3DModel.from_config(transformer_config) print("Downloading and 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) missing = [k for k in missing if not k.endswith(".bias")] if missing: print(f"Warning: missing non-bias keys ({len(missing)}): {missing[:5]}") if unexpected: print(f"Warning: unexpected keys ({len(unexpected)}): {unexpected[:5]}") pipe.transformer = transformer.to(dtype=dtype_target) if torch.cuda.is_available(): pipe.to("cuda") print("Anima-2.9B Pipeline ready!") # --------------------------------------------------------------------------- # ZeroGPU duration estimation & inference helper # --------------------------------------------------------------------------- DEFAULT_HEIGHT = 1216 DEFAULT_WIDTH = 832 DEFAULT_STEPS = 28 DEFAULT_CFG = 4.0 DEFAULT_NEGATIVE = "worst quality, low quality, score_1, score_2, score_3, artist name, blurry, jpeg artifacts, chromatic aberration" def _estimate_duration(prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed, *args, **kwargs) -> int: """Estimate GPU reservation duration based on pixel count and steps.""" h = int(height) if height else DEFAULT_HEIGHT w = int(width) if width else DEFAULT_WIDTH steps = int(num_inference_steps) if num_inference_steps else DEFAULT_STEPS pixels = max(h, 1) * max(w, 1) per_step = 0.5 * (pixels / (1024 * 1024)) seconds = steps * per_step return max(30, min(int(seconds) + 25, 180)) def _friendly_gpu_error(err: Exception) -> str: msg = (str(err) or "").lower() capacity_hints = ("gpu limit", "reached its gpu limit", "gpu quota", "out of quota", "quota", "no gpu", "could not allocate", "busy", "concurrent") if any(h in msg for h in capacity_hints): return "⛔ The shared GPU is currently at capacity. Please wait a minute and retry." if "out of memory" in msg or "oom" in msg or "cuda" in msg: return "💥 Ran out of GPU memory. Try lowering the resolution or inference steps." return f"⚠️ Generation error: {str(err)}" @spaces.GPU(duration=_estimate_duration) def _generate_image_gpu( prompt: str, negative_prompt: str, height: int, width: int, num_inference_steps: int, guidance_scale: float, seed: int, randomize_seed: bool, progress=gr.Progress(track_tqdm=True), ): if not prompt or not prompt.strip(): raise gr.Error("Please provide a prompt.") # Ensure height and width are divisible by 16 height = int(height) // 16 * 16 width = int(width) // 16 * 16 num_inference_steps = int(num_inference_steps) guidance_scale = float(guidance_scale) 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) start_time = time.time() image = pipe( prompt=prompt, negative_prompt=negative_prompt if negative_prompt and negative_prompt.strip() else None, height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, ).images[0] elapsed = time.time() - start_time status_text = f"✨ Generated in {elapsed:.2f}s | Seed: {seed} | Resolution: {width}x{height} | Steps: {num_inference_steps} | CFG: {guidance_scale}" return image, status_text, seed def generate( prompt: str, negative_prompt: str, height: int, width: int, num_inference_steps: int, guidance_scale: float, seed: int, randomize_seed: bool, ): try: return _generate_image_gpu( prompt, negative_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 # Preset resolution lookup def on_aspect_ratio_change(ratio_choice: str) -> Tuple[int, int]: ratios = { "Portrait (832 x 1216) [Recommended]": (832, 1216), "Landscape (1216 x 832)": (1216, 832), "Square (1024 x 1024)": (1024, 1024), "Wide / Cinematic (1360 x 768)": (1360, 768), "Tall / Mobile (768 x 1360)": (768, 1360), } return ratios.get(ratio_choice, (832, 1216)) # Quick prompt tag inserters def add_quality_prefix(current_prompt: str) -> str: prefix = "masterpiece, best quality, score_9, safe, " if current_prompt.startswith(prefix): return current_prompt return prefix + current_prompt # --------------------------------------------------------------------------- # Gradio UI Layout # --------------------------------------------------------------------------- custom_css = """ .container { max-width: 1200px; margin: auto; } .header-badge { display: inline-block; padding: 0.25rem 0.6rem; border-radius: 9999px; font-size: 0.75rem; font-weight: 600; background: #4f46e5; color: #ffffff; margin-right: 0.5rem; } """ with gr.Blocks(title="Anima 2.9B - State of the Art Anime Diffusion") as demo: with gr.Column(elem_classes=["container"]): gr.HTML( """

🎨 Anima 2.9B Demo

2.9 Billion Parameters 40 Transformer Layers Knowledge Cutoff: July 2026

An expanded 40-layer fine-tune of circlestone-labs/Anima trained on 1.7M additional anime/illustration samples with Muon optimization.

""" ) with gr.Row(): with gr.Column(scale=5): prompt_input = gr.Textbox( label="Prompt", placeholder="masterpiece, best quality, score_9, safe, 1girl, solo, silver hair, blue eyes, dynamic battle pose, glowing katana, cherry blossoms, wind, sunset...", lines=4, value="masterpiece, best quality, score_9, safe, 1girl, solo, long silver hair, glowing blue eyes, elegant white kimono with gold embroidery, floating cherry blossom petals, mystical night forest, soft moon lighting, highly detailed background", ) with gr.Row(): btn_add_quality = gr.Button("✨ Prepend Quality Tags", size="sm", variant="secondary") negative_prompt_input = gr.Textbox( label="Negative Prompt", value=DEFAULT_NEGATIVE, lines=2, ) with gr.Row(): aspect_ratio_select = gr.Dropdown( label="Aspect Ratio Preset", choices=[ "Portrait (832 x 1216) [Recommended]", "Landscape (1216 x 832)", "Square (1024 x 1024)", "Wide / Cinematic (1360 x 768)", "Tall / Mobile (768 x 1360)", ], value="Portrait (832 x 1216) [Recommended]", ) with gr.Row(): width_slider = gr.Slider(minimum=512, maximum=1536, step=16, value=DEFAULT_WIDTH, label="Width") height_slider = gr.Slider(minimum=512, maximum=1536, step=16, value=DEFAULT_HEIGHT, label="Height") with gr.Accordion("⚙️ Advanced Generation Settings", open=False): with gr.Row(): steps_slider = gr.Slider(minimum=20, maximum=50, step=1, value=DEFAULT_STEPS, label="Inference Steps") cfg_slider = gr.Slider(minimum=1.0, maximum=10.0, step=0.5, value=DEFAULT_CFG, label="Guidance Scale (CFG)") with gr.Row(): seed_input = gr.Number(value=0, label="Seed", precision=0) randomize_seed_cb = gr.Checkbox(label="Randomize Seed", value=True) generate_btn = gr.Button("🚀 Generate Illustration", variant="primary", size="lg") with gr.Column(scale=5): image_output = gr.Image(label="Generated Image", type="pil", format="png") status_output = gr.Markdown(value="*Click 'Generate Illustration' to create an artwork.*") # Wire resolution selector aspect_ratio_select.change( fn=on_aspect_ratio_change, inputs=[aspect_ratio_select], outputs=[width_slider, height_slider], ) # Wire quality tag button btn_add_quality.click( fn=add_quality_prefix, inputs=[prompt_input], outputs=[prompt_input], ) # Wire generation generate_btn.click( fn=generate, inputs=[ prompt_input, negative_prompt_input, height_slider, width_slider, steps_slider, cfg_slider, seed_input, randomize_seed_cb, ], outputs=[ image_output, status_output, seed_input, ], ) # Curated Examples gr.Markdown("### 💡 Curated Prompt Examples") examples = gr.Examples( examples=[ [ "masterpiece, best quality, score_9, safe, 1girl, solo, long silver hair, glowing blue eyes, elegant white kimono with gold embroidery, floating cherry blossom petals, mystical night forest, soft moon lighting, highly detailed background", DEFAULT_NEGATIVE, 1216, 832, 28, 4.0, 42, False, ], [ "masterpiece, best quality, score_9, safe, 1boy, solo, black hair, spiky hair, crimson eyes, black high-tech combat jacket, holding glowing laser blade, futuristic cyberpunk alley, neon signs, rain reflections, volumetric fog", DEFAULT_NEGATIVE, 1216, 832, 30, 4.5, 1337, False, ], [ "masterpiece, best quality, score_9, safe, scenery, fantasy floating islands in the sky, huge ancient castle with waterfalls, lush green vegetation, glowing crystal formations, golden sunset light, flock of white birds", DEFAULT_NEGATIVE, 832, 1216, 32, 4.0, 2026, False, ], [ "masterpiece, best quality, score_9, safe, 1girl, fern, sousou no frieren, @nnn yryr, long purple hair, purple eyes, black coat over white dress, holding wooden magic staff, gentle smile, library with towering bookshelves, warm sunlight rays", DEFAULT_NEGATIVE, 1216, 832, 30, 4.0, 777, False, ], ], inputs=[ prompt_input, negative_prompt_input, height_slider, width_slider, steps_slider, cfg_slider, seed_input, randomize_seed_cb, ], outputs=[ image_output, status_output, seed_input, ], fn=generate, cache_examples=False, ) # Guide & Documentation Accordion with gr.Accordion("📖 Prompting Tips & Model Details", open=False): gr.Markdown( """ ### 🏷️ Tag Order & Formatting Best Practices Anima is trained on Danbooru/Gelbooru tags and natural language descriptions: 1. **Quality & Meta Tags**: `masterpiece, best quality, score_9, safe, highres` 2. **Character Count**: `1girl`, `1boy`, `2girls`, `solo` 3. **Character Name & Series**: Follow with character tags + franchise name (e.g. `oomuro sakurako, yuru yuri`) 4. **Artist Style Tags**: **Prefix with `@`** (e.g. `@artist_name`). The effect is significantly stronger when using the `@` prefix. 5. **Appearance & Clothing**: `brown hair, blunt bangs, red gloves, winter coat, smile` 6. **Setting & Composition**: `looking at viewer, dramatic lighting, detailed background, snow` ### ⚙️ Recommended Settings - **Resolution**: `832x1216` (Portrait) or `1216x832` (Landscape). Dimensions must be divisible by 16. - **Steps**: `28-50` (28 is fast and high-quality, 50 for fine nuances). - **CFG Guidance**: `3.5 - 5.0`. - **Sampler**: Euler with Flow Match Euler schedule. ### 📜 Model & License Info - **Model**: [Gazingstars123/Anima-2.9B](https://huggingface.co/Gazingstars123/Anima-2.9B) - **Base Architecture**: 40-layer expanded transformer built on [circlestone-labs/Anima](https://huggingface.co/circlestone-labs/Anima) & [nvidia/Cosmos-Predict2-2B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image). - **License**: [CircleStone Labs Non-Commercial License](https://huggingface.co/circlestone-labs/Anima/blob/main/LICENSE.md) (Outputs can be used commercially). """ ) if __name__ == "__main__": demo.launch(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"), css=custom_css)