Spaces:
Running on Zero
Running on Zero
| 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: | |
| 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)}" | |
| 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( | |
| """ | |
| <div style="text-align: center; margin-bottom: 1.5rem;"> | |
| <h1 style="font-size: 2.25rem; font-weight: 800; margin-bottom: 0.5rem;"> | |
| 🎨 Anima 2.9B Demo | |
| </h1> | |
| <div style="margin-bottom: 0.75rem;"> | |
| <span class="header-badge">2.9 Billion Parameters</span> | |
| <span class="header-badge" style="background:#059669;">40 Transformer Layers</span> | |
| <span class="header-badge" style="background:#0284c7;">Knowledge Cutoff: July 2026</span> | |
| </div> | |
| <p style="color: #64748b; font-size: 0.95rem; max-width: 800px; margin: auto;"> | |
| An expanded 40-layer fine-tune of <a href="https://huggingface.co/circlestone-labs/Anima" target="_blank" style="color:#4f46e5; text-decoration: underline;">circlestone-labs/Anima</a> | |
| trained on 1.7M additional anime/illustration samples with Muon optimization. | |
| </p> | |
| </div> | |
| """ | |
| ) | |
| 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) | |