Fix Diffusers downloads not counting.

Aniimage_2_Interactive_Colab.ipynb

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

README.md

@@ -1,153 +1,153 @@
----
-license: apache-2.0
-tags:
-- anime
-- diffusion
-- text-to-image
-- image-generation
-library_name: diffusers
-pipeline_tag: text-to-image
-language:
-- en
----
-![Aniimage-2 Samples](https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/assets/collage.png)
-
-# Aniimage-2
-
-Aniimage-2 is the second public model in the Aniimage series.
-The model is a 512x512 anime image generation model trained from scratch using a UNet + VAE + CLIP architecture with 441.9 million parameters.
-Aniimage-2 has been trained on 1.2M anime images from [Danbooru](https://danbooru.donmai.us/) and other sources.
-It is not based on any existing models, the UNet is trained from scratch.
-
-## Model Details
-
-| | |
-|---|---|
-| **Resolution** | 512x512 |
-| **Architecture** | Latent Diffusion (UNet + VAE + CLIP) |
-| **Parameters** | 441.9M |
-| **Training Steps** | 70,000 |
-| **Batch Size** | 64 |
-| **Dataset** | ~1.2 million curated anime images from Danbooru |
-| **GPU** | NVIDIA RTX 5060 Ti 16GB |
-| **Scheduler** | DPM++ SDE |
-| **Step count** | 50 using DPM++ SDE|
-
-## Requirements
-
-- **GPU**: ~6 GB VRAM minimum (recommend 8+ GB)
-- **CPU**: ~7 GB RAM. Image generation is extremely slow on CPU. (recommend 10+ GB)
-
-## Prompting
-
-Aniimage uses plain text captions as well as Danbooru tags.
-
-An example of acceptable captions:
-
-"A smiling anime girl with red hair and a school uniform"
-
-"1girl, solo, smile, red_hair, school_uniform, anime_coloring"
-
-Long prompts and enhanced prompts usually result in worse overall results. Try to keep your prompts short.
-
-Although you can use Danbooru tags, you will get the best results with plain text captions
-
-Adding a negative prompt is required for a good output. The one used during inference in the preview image is:
-
-"low quality, ugly, blurry, distorted, deformed, bad anatomy, bad proportions, extra limbs, missing limbs, watermark, text, signature, washed out, flat colors, manga panel, disfigured, poorly drawn, jpeg artifacts, cropped, out of frame"
-
-
-For the best quality, use DPM++ SDE at 50 steps.
-
-**Important:** If you do not include “NSFW” in the negative prompt, the model may generate NSFW images
-
-
-## Run
-
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/#fileId=https://huggingface.co/8BitStudio/Aniimage-2/blob/main/Aniimage_2_Interactive_Colab.ipynb)
-
-[![Download Generator](https://img.shields.io/badge/Download-generate__images.py-blue?style=for-the-badge)](https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/generate_images.py)
-
-after downloading, install the dependencies.
-
-```bash
-pip install torch torchvision diffusers transformers safetensors pillow huggingface_hub
-python generate_hf.py
-```
-
-## Capabilities
-
-Aniimage-2 has significantly higher capabilities than [Aniimage-1](https://huggingface.co/8BitStudio/Aniimage-1). Especially when it comes to backgrounds.
-
-<table width="700">
-  <tr>
-    <td colspan="2">
-      <img
-        src="https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/assets/Softly%20lit%20forest%20at%20midday%2C%20peaceful%20atmosphere%2C%20no%20people.png"
-        width="700"
-        alt="Forest comparison between Aniimage-1 and Aniimage-2"
-      >
-    </td>
-  </tr>
-
-  <tr>
-    <td align="center" width="50%">
-      <b>Aniimage-1</b>
-    </td>
-    <td align="center" width="50%">
-      <b>Aniimage-2</b>
-    </td>
-  </tr>
-
-  <tr>
-    <td colspan="2">
-      <sub>Prompt: Softly lit forest at midday, peaceful atmosphere, no people.</sub>
-    </td>
-  </tr>
-</table>
-
-<table width="700">
-  <tr>
-    <td colspan="2">
-      <img
-        src="https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/assets/anime-boy.png"
-        width="700"
-        alt="Anime boy comparison between Aniimage-1 and Aniimage-2"
-      >
-    </td>
-  </tr>
-
-  <tr>
-    <td align="center" width="50%">
-      <b>Aniimage-1</b>
-    </td>
-    <td align="center" width="50%">
-      <b>Aniimage-2</b>
-    </td>
-  </tr>
-
-Aniimage-2 produces male characters much more reliably than Aniimage-1.
-
-  <tr>
-    <td colspan="2">
-      <sub>Prompt: An anime boy with black hair and green eyes.</sub>
-    </td>
-  </tr>
-</table>
-
-## Limitations
-
-Aniimage-2 is bad at generating hands. It may generate hands but they often melt into surfaces if they are similar colors and usually have more or fewer than five fingers.
-The model generates male characters less reliably than female characters, it may take many images to find a good image like above.
-
-## What's Next
-
-Aniimage-3 is currently in the works and will have a much higher parameter count: 738.1M parameters compared to aniimage-2 and aniimage-1's 441.9M parameters.
-The model will also have significantly more unique images.
-Aniimage-2 used incorrect Min-SNR loss weighting with its v-prediction objective, causing near-pure-noise timesteps to receive excessive emphasis. This has been corrected for Aniimage-3.
-Aniimage-3 also uses a much higher-quality dataset with improved filtering. More than 100,000 low-quality images from the Aniimage-2 dataset were removed.
-
-
-## License
-
+---
+license: apache-2.0
+tags:
+- anime
+- diffusion
+- text-to-image
+- image-generation
+library_name: diffusers
+pipeline_tag: text-to-image
+language:
+- en
+---
+![Aniimage-2 Samples](https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/assets/collage.png)
+
+# Aniimage-2
+
+Aniimage-2 is the second public model in the Aniimage series.
+The model is a 512x512 anime image generation model trained from scratch using a UNet + VAE + CLIP architecture with 441.9 million parameters.
+Aniimage-2 has been trained on 1.2M anime images from [Danbooru](https://danbooru.donmai.us/) and other sources.
+It is not based on any existing models, the UNet is trained from scratch.
+
+## Model Details
+
+| | |
+|---|---|
+| **Resolution** | 512x512 |
+| **Architecture** | Latent Diffusion (UNet + VAE + CLIP) |
+| **Parameters** | 441.9M |
+| **Training Steps** | 70,000 |
+| **Batch Size** | 64 |
+| **Dataset** | ~1.2 million curated anime images from Danbooru |
+| **GPU** | NVIDIA RTX 5060 Ti 16GB |
+| **Scheduler** | DPM++ SDE |
+| **Step count** | 50 using DPM++ SDE|
+
+## Requirements
+
+- **GPU**: ~6 GB VRAM minimum (recommend 8+ GB)
+- **CPU**: ~7 GB RAM. Image generation is extremely slow on CPU. (recommend 10+ GB)
+
+## Prompting
+
+Aniimage uses plain text captions as well as Danbooru tags.
+
+An example of acceptable captions:
+
+"A smiling anime girl with red hair and a school uniform"
+
+"1girl, solo, smile, red_hair, school_uniform, anime_coloring"
+
+Long prompts and enhanced prompts usually result in worse overall results. Try to keep your prompts short.
+
+Although you can use Danbooru tags, you will get the best results with plain text captions
+
+Adding a negative prompt is required for a good output. The one used during inference in the preview image is:
+
+"low quality, ugly, blurry, distorted, deformed, bad anatomy, bad proportions, extra limbs, missing limbs, watermark, text, signature, washed out, flat colors, manga panel, disfigured, poorly drawn, jpeg artifacts, cropped, out of frame"
+
+
+For the best quality, use DPM++ SDE at 50 steps.
+
+**Important:** If you do not include “NSFW” in the negative prompt, the model may generate NSFW images
+
+
+## Run
+
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/#fileId=https://huggingface.co/8BitStudio/Aniimage-2/blob/main/Aniimage_2_Interactive_Colab.ipynb)
+
+[![Download Generator](https://img.shields.io/badge/Download-generate__images.py-blue?style=for-the-badge)](https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/generate_images.py)
+
+after downloading, install the dependencies.
+
+```bash
+pip install torch torchvision diffusers transformers safetensors pillow huggingface_hub
+python generate_images.py
+```
+
+## Capabilities
+
+Aniimage-2 has significantly higher capabilities than [Aniimage-1](https://huggingface.co/8BitStudio/Aniimage-1). Especially when it comes to backgrounds.
+
+<table width="700">
+  <tr>
+    <td colspan="2">
+      <img
+        src="https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/assets/Softly%20lit%20forest%20at%20midday%2C%20peaceful%20atmosphere%2C%20no%20people.png"
+        width="700"
+        alt="Forest comparison between Aniimage-1 and Aniimage-2"
+      >
+    </td>
+  </tr>
+
+  <tr>
+    <td align="center" width="50%">
+      <b>Aniimage-1</b>
+    </td>
+    <td align="center" width="50%">
+      <b>Aniimage-2</b>
+    </td>
+  </tr>
+
+  <tr>
+    <td colspan="2">
+      <sub>Prompt: Softly lit forest at midday, peaceful atmosphere, no people.</sub>
+    </td>
+  </tr>
+</table>
+
+<table width="700">
+  <tr>
+    <td colspan="2">
+      <img
+        src="https://huggingface.co/8BitStudio/Aniimage-2/resolve/main/assets/anime-boy.png"
+        width="700"
+        alt="Anime boy comparison between Aniimage-1 and Aniimage-2"
+      >
+    </td>
+  </tr>
+
+  <tr>
+    <td align="center" width="50%">
+      <b>Aniimage-1</b>
+    </td>
+    <td align="center" width="50%">
+      <b>Aniimage-2</b>
+    </td>
+  </tr>
+
+Aniimage-2 produces male characters much more reliably than Aniimage-1.
+
+  <tr>
+    <td colspan="2">
+      <sub>Prompt: An anime boy with black hair and green eyes.</sub>
+    </td>
+  </tr>
+</table>
+
+## Limitations
+
+Aniimage-2 is bad at generating hands. It may generate hands but they often melt into surfaces if they are similar colors and usually have more or fewer than five fingers.
+The model generates male characters less reliably than female characters, it may take many images to find a good image like above.
+
+## What's Next
+
+Aniimage-3 is currently in the works and will have a much higher parameter count: 738.1M parameters compared to aniimage-2 and aniimage-1's 441.9M parameters.
+The model will also have significantly more unique images.
+Aniimage-2 used incorrect Min-SNR loss weighting with its v-prediction objective, causing near-pure-noise timesteps to receive excessive emphasis. This has been corrected for Aniimage-3.
+Aniimage-3 also uses a much higher-quality dataset with improved filtering. More than 100,000 low-quality images from the Aniimage-2 dataset were removed.
+
+
+## License
+
 Apache 2.0

config.json

@@ -0,0 +1,67 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.37.1",
+  "act_fn": "silu",
+  "addition_embed_type": null,
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": null,
+  "attention_head_dim": 8,
+  "attention_type": "default",
+  "block_out_channels": [
+    256,
+    512,
+    768,
+    1024
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 768,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "use_linear_projection": false
+}

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c298c8d61a380e4941d02372bf2e7de27de15ecd28c2236cd005c59c598105a4
+size 1767752200

model_config.json

@@ -0,0 +1,36 @@
+{
+  "unet": {
+    "sample_size": 64,
+    "in_channels": 4,
+    "out_channels": 4,
+    "layers_per_block": 2,
+    "block_out_channels": [
+      256,
+      512,
+      768,
+      1024
+    ],
+    "down_block_types": [
+      "CrossAttnDownBlock2D",
+      "CrossAttnDownBlock2D",
+      "CrossAttnDownBlock2D",
+      "DownBlock2D"
+    ],
+    "up_block_types": [
+      "UpBlock2D",
+      "CrossAttnUpBlock2D",
+      "CrossAttnUpBlock2D",
+      "CrossAttnUpBlock2D"
+    ],
+    "cross_attention_dim": 768,
+    "attention_head_dim": 8
+  },
+  "prediction_type": "v_prediction",
+  "zero_terminal_snr": true,
+  "beta_schedule": "scaled_linear",
+  "num_train_timesteps": 1000,
+  "clip_penultimate": true,
+  "guidance_rescale": 0.7,
+  "image_size": 512,
+  "vae": "madebyollin/sdxl-vae-fp16-fix"
+}