Upload Aniimage_2_Interactive_Colab.ipynb

Aniimage_2_Interactive_Colab.ipynb

@@ -0,0 +1,296 @@
+{
+  "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
+}