{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Aniimage-2 Image generator\n", "\n", "This notebook runs [`8BitStudio/Aniimage-2`](https://huggingface.co/8BitStudio/Aniimage-2) with a Gradio interface. It manually assembles the repository's UNet with the matching VAE and CLIP text encoder because the repository is not packaged as a complete Diffusers pipeline.\n", "\n", "## Before running\n", "\n", "1. In Colab choose **Runtime → Change runtime type → T4 GPU** (or a better GPU).\n", "2. Run every cell in order. The first model load downloads several gigabytes and can take a few minutes.\n", "3. Open the Gradio link printed by the final cell. Generated images are also saved under `/content/aniimage2_outputs`.\n", "\n", "The defaults follow the model card: 512×512, DPM++ SDE Karras, 50 steps, v-prediction, zero-terminal-SNR, CLIP penultimate layer, and CFG rescale 0.7. The default negative prompt includes `NSFW` because omitting it may produce NSFW images.\n", "\n", "If you want to generate NSFW images, remove the `NSFW` text in the negative prompt. \n", "\n", "If your results look \"glitched\", try lowering your CFG to around 7 or 6.5" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip -q install -U \"diffusers>=0.37.1\" \"transformers>=4.46,<5\" accelerate safetensors huggingface_hub \"gradio>=5,<7\" \"click>=8.2\" \"Pillow<11.0.0\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "import random\n", "from pathlib import Path\n", "\n", "import gradio as gr\n", "import numpy as np\n", "import torch\n", "from PIL import Image\n", "from diffusers import (\n", " AutoencoderKL,\n", " DDIMScheduler,\n", " DPMSolverMultistepScheduler,\n", " EulerAncestralDiscreteScheduler,\n", " EulerDiscreteScheduler,\n", " UNet2DConditionModel,\n", ")\n", "from huggingface_hub import hf_hub_download\n", "from transformers import CLIPTextModel, CLIPTokenizer\n", "\n", "REPO_ID = \"8BitStudio/Aniimage-2\"\n", "CLIP_ID = \"openai/clip-vit-large-patch14\"\n", "OUTPUT_DIR = Path(\"/content/aniimage2_outputs\")\n", "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", "if not torch.cuda.is_available():\n", " raise RuntimeError(\"No GPU detected. In Colab select Runtime > Change runtime type > T4 GPU, then run again.\")\n", "\n", "DEVICE = \"cuda\"\n", "DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16\n", "torch.backends.cuda.matmul.allow_tf32 = True\n", "torch.backends.cudnn.allow_tf32 = True\n", "print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n", "print(f\"Model dtype: {DTYPE}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Download the small model metadata first so every component matches training.\n", "config_path = hf_hub_download(REPO_ID, \"Aniimage-2/model_config.json\")\n", "with open(config_path, \"r\", encoding=\"utf-8\") as f:\n", " MODEL_CONFIG = json.load(f)\n", "\n", "VAE_ID = MODEL_CONFIG[\"vae\"]\n", "GUIDANCE_RESCALE = float(MODEL_CONFIG.get(\"guidance_rescale\", 0.7))\n", "\n", "print(\"Loading Aniimage-2 UNet...\")\n", "unet = UNet2DConditionModel.from_pretrained(\n", " REPO_ID,\n", " subfolder=\"Aniimage-2/unet\",\n", " torch_dtype=DTYPE,\n", " low_cpu_mem_usage=True,\n", ").to(DEVICE).eval()\n", "unet.requires_grad_(False)\n", "\n", "print(f\"Loading VAE: {VAE_ID}...\")\n", "vae = AutoencoderKL.from_pretrained(VAE_ID, torch_dtype=DTYPE).to(DEVICE).eval()\n", "vae.requires_grad_(False)\n", "vae.enable_slicing()\n", "\n", "print(f\"Loading text encoder: {CLIP_ID}...\")\n", "tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID)\n", "text_encoder = CLIPTextModel.from_pretrained(CLIP_ID, torch_dtype=DTYPE).to(DEVICE).eval()\n", "text_encoder.requires_grad_(False)\n", "\n", "# Aniimage-2 was trained using CLIP's penultimate transformer layer.\n", "clip_inner = getattr(text_encoder, \"text_model\", text_encoder)\n", "clip_inner.encoder.layers = torch.nn.ModuleList(list(clip_inner.encoder.layers[:-1]))\n", "\n", "print(\"Aniimage-2 is loaded and ready.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def make_scheduler(name):\n", " base = dict(\n", " num_train_timesteps=int(MODEL_CONFIG.get(\"num_train_timesteps\", 1000)),\n", " beta_schedule=MODEL_CONFIG.get(\"beta_schedule\", \"scaled_linear\"),\n", " prediction_type=MODEL_CONFIG.get(\"prediction_type\", \"v_prediction\"),\n", " rescale_betas_zero_snr=bool(MODEL_CONFIG.get(\"zero_terminal_snr\", True)),\n", " timestep_spacing=MODEL_CONFIG.get(\"timestep_spacing\", \"trailing\"),\n", " )\n", " if name == \"DPM++ SDE Karras\":\n", " return DPMSolverMultistepScheduler(\n", " **base, algorithm_type=\"sde-dpmsolver++\", solver_order=2, use_karras_sigmas=True\n", " )\n", " if name == \"DPM++ 2M Karras\":\n", " return DPMSolverMultistepScheduler(\n", " **base, algorithm_type=\"dpmsolver++\", solver_order=2, use_karras_sigmas=True\n", " )\n", " if name == \"Euler a\":\n", " return EulerAncestralDiscreteScheduler(**base)\n", " if name == \"Euler\":\n", " return EulerDiscreteScheduler(**base)\n", " if name == \"DDIM\":\n", " return DDIMScheduler(**base, clip_sample=False, set_alpha_to_one=False)\n", " raise ValueError(f\"Unknown scheduler: {name}\")\n", "\n", "\n", "@torch.inference_mode()\n", "def encode_prompts(prompt, negative_prompt):\n", " tokens = tokenizer(\n", " [negative_prompt or \"\", prompt],\n", " padding=\"max_length\",\n", " max_length=tokenizer.model_max_length,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " )\n", " return text_encoder(tokens.input_ids.to(DEVICE))[0]\n", "\n", "\n", "def rescale_cfg(noise_cfg, noise_text, amount):\n", " dims = tuple(range(1, noise_cfg.ndim))\n", " std_text = noise_text.std(dim=dims, keepdim=True)\n", " std_cfg = noise_cfg.std(dim=dims, keepdim=True).clamp_min(1e-6)\n", " noise_rescaled = noise_cfg * (std_text / std_cfg)\n", " return amount * noise_rescaled + (1.0 - amount) * noise_cfg\n", "\n", "\n", "@torch.inference_mode()\n", "def generate_one(prompt, negative_prompt, scheduler_name, steps, cfg_scale, seed):\n", " scheduler = make_scheduler(scheduler_name)\n", " scheduler.set_timesteps(int(steps), device=DEVICE)\n", " embeddings = encode_prompts(prompt, negative_prompt)\n", "\n", " generator = torch.Generator(device=DEVICE).manual_seed(int(seed))\n", " latent_size = int(MODEL_CONFIG.get(\"image_size\", 512)) // 8\n", " latents = torch.randn(\n", " (1, int(unet.config.in_channels), latent_size, latent_size),\n", " generator=generator, device=DEVICE, dtype=torch.float32,\n", " ) * scheduler.init_noise_sigma\n", "\n", " for timestep in scheduler.timesteps:\n", " latent_input = torch.cat([latents, latents], dim=0)\n", " latent_input = scheduler.scale_model_input(latent_input, timestep)\n", " with torch.autocast(\"cuda\", dtype=DTYPE):\n", " prediction = unet(latent_input, timestep, encoder_hidden_states=embeddings).sample\n", " pred_negative, pred_text = prediction.chunk(2)\n", " prediction = pred_negative + float(cfg_scale) * (pred_text - pred_negative)\n", " prediction = rescale_cfg(prediction, pred_text, GUIDANCE_RESCALE)\n", " latents = scheduler.step(prediction, timestep, latents).prev_sample\n", "\n", " scaled = (latents / vae.config.scaling_factor).to(dtype=DTYPE)\n", " with torch.autocast(\"cuda\", dtype=DTYPE):\n", " image = vae.decode(scaled).sample\n", " image = (image.float() / 2 + 0.5).clamp(0, 1)\n", " array = (image[0].permute(1, 2, 0).cpu().numpy() * 255).round().astype(np.uint8)\n", " return Image.fromarray(array)\n", "\n", "\n", "def generate_gallery(prompt, negative_prompt, scheduler_name, steps, cfg_scale, seed, randomize_seed, image_count, progress=gr.Progress()):\n", " prompt = (prompt or \"\").strip()\n", " if not prompt:\n", " raise gr.Error(\"Enter a prompt first.\")\n", "\n", " count = int(image_count)\n", " base_seed = random.randint(0, 2**31 - 1) if randomize_seed or int(seed) < 0 else int(seed)\n", " images, records = [], []\n", "\n", " for index in range(count):\n", " used_seed = (base_seed + index) % (2**31)\n", " progress(index / count, desc=f\"Generating image {index + 1} of {count}\")\n", " image = generate_one(prompt, negative_prompt, scheduler_name, steps, cfg_scale, used_seed)\n", " path = OUTPUT_DIR / f\"aniimage2_{used_seed}.png\"\n", " image.save(path)\n", " images.append(image)\n", " records.append(f\"- Seed `{used_seed}` — `{path}`\")\n", "\n", " progress(1.0, desc=\"Done\")\n", " details = \"### Results\\n\" + \"\\n\".join(records)\n", " return images, details, base_seed\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "DEFAULT_NEGATIVE = (\n", " \"NSFW, low quality, ugly, blurry, distorted, deformed, bad anatomy, bad proportions, \"\n", " \"extra limbs, missing limbs, watermark, text, signature, washed out, flat colors, \"\n", " \"manga panel, disfigured, poorly drawn, jpeg artifacts, cropped, out of frame\"\n", ")\n", "\n", "with gr.Blocks(title=\"Aniimage-2 Generator\", theme=gr.themes.Soft()) as demo:\n", " gr.Markdown(\n", " \"# Aniimage-2 Interactive Generator\\n\"\n", " \"Use a short plain-English prompt for best results. Images are 512×512.\"\n", " )\n", " with gr.Row():\n", " with gr.Column(scale=2):\n", " prompt = gr.Textbox(\n", " label=\"Prompt\",\n", " value=\"A smiling anime girl with red hair and a school uniform\",\n", " lines=3,\n", " )\n", " negative = gr.Textbox(label=\"Negative prompt\", value=DEFAULT_NEGATIVE, lines=4)\n", " generate_button = gr.Button(\"Generate\", variant=\"primary\")\n", " with gr.Column(scale=1):\n", " scheduler = gr.Dropdown(\n", " [\"DPM++ SDE Karras\", \"DPM++ 2M Karras\", \"Euler a\", \"Euler\", \"DDIM\"],\n", " value=\"DPM++ SDE Karras\", label=\"Scheduler\",\n", " )\n", " steps = gr.Slider(10, 80, value=50, step=1, label=\"Steps\")\n", " cfg = gr.Slider(1.0, 15.0, value=7.5, step=0.1, label=\"CFG scale\")\n", " seed = gr.Number(value=-1, precision=0, label=\"Seed (-1 = random)\")\n", " randomize = gr.Checkbox(value=True, label=\"Randomize seed\")\n", " count = gr.Slider(1, 4, value=1, step=1, label=\"Number of images\")\n", "\n", " gallery = gr.Gallery(label=\"Generated images\", columns=2, object_fit=\"contain\", height=620)\n", " result_info = gr.Markdown()\n", "\n", " generate_button.click(\n", " fn=generate_gallery,\n", " inputs=[prompt, negative, scheduler, steps, cfg, seed, randomize, count],\n", " outputs=[gallery, result_info, seed],\n", " concurrency_limit=1,\n", " )\n", "\n", "demo.queue(max_size=8)\n", "print(\"Interface created. Run the final cell to launch it.\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "demo.launch(share=True, show_error=True)\n" ] } ], "metadata": { "accelerator": "GPU", "colab": { "name": "Aniimage_2_Interactive_Colab.ipynb", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }