{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "colab": { "provenance": [], "gpuType": "T4" }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Aniimage-1 Colab Generator\n", "\n", "Hugging Face's default `DiffusionPipeline.from_pretrained(...)` snippet does not work for this repo because `8BitStudio/Aniimage-1` is not packaged as a full Diffusers pipeline. It ships the Aniimage UNet weights/config, so this notebook manually loads:\n", "\n", "- Aniimage-1 UNet from `8BitStudio/Aniimage-1`\n", "- SD VAE from `stabilityai/sd-vae-ft-mse`\n", "- CLIP text encoder/tokenizer from `openai/clip-vit-large-patch14`\n", "\n", "Use Colab GPU runtime: `Runtime > Change runtime type > T4 GPU` or better.\n", "If you have any errors or issues, open a discussion on the Hugging Face repo." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip -q install -U diffusers transformers accelerate safetensors huggingface_hub pillow" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "from pathlib import Path\n", "\n", "import numpy as np\n", "import torch\n", "import torch.nn.functional as F\n", "from PIL import Image, ImageEnhance, ImageFilter\n", "from IPython.display import display\n", "from huggingface_hub import hf_hub_download\n", "from safetensors.torch import load_file\n", "from diffusers import (\n", " AutoencoderKL,\n", " DDIMScheduler,\n", " DPMSolverMultistepScheduler,\n", " EulerAncestralDiscreteScheduler,\n", " EulerDiscreteScheduler,\n", " UNet2DConditionModel,\n", ")\n", "from transformers import CLIPTextModel, CLIPTokenizer\n", "\n", "REPO_ID = \"8BitStudio/Aniimage-1\"\n", "VAE_ID = \"stabilityai/sd-vae-ft-mse\"\n", "CLIP_ID = \"openai/clip-vit-large-patch14\"\n", "MODEL_DIR = Path(\"/content/aniimage_1\")\n", "\n", "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "autocast_dtype = torch.float16 if device == \"cuda\" else torch.float32\n", "\n", "if device == \"cuda\":\n", " torch.backends.cuda.matmul.allow_tf32 = True\n", " torch.backends.cudnn.allow_tf32 = True\n", "\n", "print(\"Device:\", device)\n", "if device == \"cuda\":\n", " print(\"GPU:\", torch.cuda.get_device_name(0))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def make_scheduler(name=\"DPM++ 2M Karras\"):\n", " base = dict(\n", " num_train_timesteps=1000,\n", " beta_schedule=\"scaled_linear\",\n", " prediction_type=\"epsilon\",\n", " )\n", " if name == \"DPM++ 2M Karras\":\n", " return DPMSolverMultistepScheduler(\n", " **base,\n", " algorithm_type=\"dpmsolver++\",\n", " solver_order=2,\n", " use_karras_sigmas=True,\n", " )\n", " if name == \"DPM++ SDE Karras\":\n", " return DPMSolverMultistepScheduler(\n", " **base,\n", " algorithm_type=\"sde-dpmsolver++\",\n", " solver_order=2,\n", " 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", "def load_aniimage1():\n", " MODEL_DIR.mkdir(parents=True, exist_ok=True)\n", "\n", " print(\"Downloading Aniimage-1 UNet config/weights if needed...\")\n", " config_path = hf_hub_download(\n", " repo_id=REPO_ID,\n", " filename=\"config.json\",\n", " local_dir=str(MODEL_DIR),\n", " )\n", " weights_path = hf_hub_download(\n", " repo_id=REPO_ID,\n", " filename=\"diffusion_pytorch_model.safetensors\",\n", " local_dir=str(MODEL_DIR),\n", " )\n", "\n", " print(\"Loading VAE...\")\n", " vae = AutoencoderKL.from_pretrained(VAE_ID).to(device).eval()\n", " vae.requires_grad_(False)\n", "\n", " print(\"Loading CLIP tokenizer/text encoder...\")\n", " tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID)\n", " text_encoder = CLIPTextModel.from_pretrained(CLIP_ID).to(device).eval()\n", " text_encoder.requires_grad_(False)\n", "\n", " print(\"Loading Aniimage-1 UNet...\")\n", " with open(config_path, \"r\", encoding=\"utf-8\") as f:\n", " unet_config = json.load(f)\n", " unet = UNet2DConditionModel.from_config(unet_config)\n", " state = load_file(weights_path, device=\"cpu\")\n", " unet.load_state_dict(state)\n", " unet.to(device).eval()\n", " unet.requires_grad_(False)\n", "\n", " print(\"Aniimage-1 ready.\")\n", " return vae, tokenizer, text_encoder, unet\n", "\n", "\n", "vae, tokenizer, text_encoder, unet = load_aniimage1()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def encode_prompt(prompt, negative_prompt=\"\"):\n", " prompts = [negative_prompt or \"\", prompt]\n", " tokens = tokenizer(\n", " prompts,\n", " padding=\"max_length\",\n", " max_length=tokenizer.model_max_length,\n", " truncation=True,\n", " return_tensors=\"pt\",\n", " )\n", " with torch.no_grad():\n", " return text_encoder(tokens.input_ids.to(device))[0]\n", "\n", "\n", "def sharpen_latents(latents, amount=0.08):\n", " blurred = F.avg_pool2d(latents, kernel_size=3, stride=1, padding=1)\n", " return latents + amount * (latents - blurred)\n", "\n", "\n", "def post_process(img):\n", " img = img.filter(ImageFilter.UnsharpMask(radius=1.5, percent=40, threshold=2))\n", " img = ImageEnhance.Contrast(img).enhance(1.06)\n", " img = ImageEnhance.Color(img).enhance(1.10)\n", " return img\n", "\n", "\n", "@torch.no_grad()\n", "def decode_latents(latents, polish=True):\n", " scaled = latents / vae.config.scaling_factor\n", " # Keep decode in fp32. It is a little slower but avoids washed/banded VAE output.\n", " image = vae.decode(scaled.float()).sample\n", " image = (image.float() / 2 + 0.5).clamp(0, 1)\n", " image = image.cpu().permute(0, 2, 3, 1).numpy()[0]\n", " image = (image * 255).round().astype(\"uint8\")\n", " img = Image.fromarray(image)\n", " return post_process(img) if polish else img\n", "\n", "\n", "@torch.no_grad()\n", "def generate_aniimage1(\n", " prompt,\n", " negative_prompt=\"\",\n", " scheduler_name=\"DPM++ 2M Karras\",\n", " steps=25,\n", " guidance_scale=7.5,\n", " seed=-1,\n", " polish=True,\n", "):\n", " if seed is None or int(seed) < 0:\n", " seed = int(torch.randint(0, 2**31 - 1, (1,)).item())\n", " else:\n", " seed = int(seed)\n", "\n", " generator = torch.Generator(device=device).manual_seed(seed)\n", " text_embeddings = encode_prompt(prompt, negative_prompt)\n", "\n", " scheduler = make_scheduler(scheduler_name)\n", " scheduler.set_timesteps(int(steps), device=device)\n", "\n", " latents = torch.randn(\n", " (1, 4, 32, 32),\n", " generator=generator,\n", " device=device,\n", " dtype=torch.float32,\n", " )\n", " latents = latents * scheduler.init_noise_sigma\n", "\n", " for t in scheduler.timesteps:\n", " latent_input = torch.cat([latents, latents], dim=0)\n", " latent_input = scheduler.scale_model_input(latent_input, t)\n", "\n", " with torch.autocast(\n", " device_type=\"cuda\",\n", " dtype=autocast_dtype,\n", " enabled=(device == \"cuda\"),\n", " ):\n", " noise_pred = unet(\n", " latent_input,\n", " t,\n", " encoder_hidden_states=text_embeddings,\n", " ).sample\n", "\n", " noise_uncond, noise_text = noise_pred.chunk(2)\n", " noise_pred = noise_uncond + float(guidance_scale) * (noise_text - noise_uncond)\n", " latents = scheduler.step(noise_pred, t, latents).prev_sample\n", "\n", " if polish:\n", " latents = sharpen_latents(latents)\n", " return decode_latents(latents, polish=polish), seed" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "prompt = \"an anime girl with long blue hair\" #@param {type:\"string\"}\n", "negative_prompt = \"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\" #@param {type:\"string\"}\n", "scheduler_name = \"DPM++ 2M Karras\" #@param [\"DPM++ 2M Karras\", \"DPM++ SDE Karras\", \"Euler a\", \"Euler\", \"DDIM\"]\n", "steps = 25 #@param {type:\"slider\", min:10, max:80, step:1}\n", "guidance_scale = 7.5 #@param {type:\"slider\", min:1, max:15, step:0.1}\n", "seed = -1 #@param {type:\"integer\"}\n", "polish = False #@param {type:\"boolean\"}\n", "\n", "image, used_seed = generate_aniimage1(\n", " prompt=prompt,\n", " negative_prompt=negative_prompt,\n", " scheduler_name=scheduler_name,\n", " steps=steps,\n", " guidance_scale=guidance_scale,\n", " seed=seed,\n", " polish=polish,\n", ")\n", "\n", "out_path = \"/content/aniimage1_output.png\"\n", "image.save(out_path)\n", "print(\"Seed:\", used_seed)\n", "print(\"Saved:\", out_path)\n", "display(image)" ] } ] }