stable_diffusion_2_0.ipynb

stable_diffusion_2_0.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "private_outputs": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/thx-pw/645c587a08ff317a5be6874c10eb36d5/stable_diffusion_2_0.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gfKvWAVnz8OB"
      },
      "source": [
        "参考元: \n",
        "\n",
        "https://colab.research.google.com/drive/1eE4mqlN9pkEo93Mtbtxz1Lk5mKdhUPX0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2q_EUvA-m4BP"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!pip install -q transformers==4.19.2 diffusers invisible-watermark omegaconf einops pytorch_lightning open_clip_torch gradio\n",
        "!git clone https://github.com/Stability-AI/stablediffusion.git\n",
        "%cd stablediffusion\n",
        "\n",
        "!sed -i -e 's|map_location=\"cpu\"|map_location=\"cuda\"|g' scripts/txt2img.py\n",
        "!sed -i -e 's|map_location=\"cpu\"|map_location=\"cuda\"|g' scripts/img2img.py\n",
        "\n",
        "!pip install -q -e .\n",
        "\n",
        "!pip install -q -U --pre triton\n",
        "!pip install -q https://github.com/metrolobo/xformers_wheels/releases/download/1d31a3ac_various_6/xformers-0.0.14.dev0-cp37-cp37m-linux_x86_64.whl"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget https://huggingface.co/stabilityai/stable-diffusion-2/resolve/main/768-v-ema.ckpt -P /content"
      ],
      "metadata": {
        "id": "OXtMVwFOMSlc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import argparse, os\n",
        "import cv2\n",
        "import torch\n",
        "import numpy as np\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from tqdm import tqdm, trange\n",
        "from itertools import islice\n",
        "from einops import rearrange\n",
        "from torchvision.utils import make_grid\n",
        "from pytorch_lightning import seed_everything\n",
        "from torch import autocast\n",
        "from contextlib import nullcontext\n",
        "from imwatermark import WatermarkEncoder\n",
        "\n",
        "from ldm.util import instantiate_from_config\n",
        "from ldm.models.diffusion.ddim import DDIMSampler\n",
        "from ldm.models.diffusion.plms import PLMSSampler\n",
        "from ldm.models.diffusion.dpm_solver import DPMSolverSampler"
      ],
      "metadata": {
        "id": "r_or7-5awumO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.set_grad_enabled(False)\n",
        "\n",
        "def chunk(it, size):\n",
        "    it = iter(it)\n",
        "    return iter(lambda: tuple(islice(it, size)), ())\n",
        "\n",
        "\n",
        "def load_model_from_config(config, ckpt, verbose=False):\n",
        "    model = instantiate_from_config(config.model)\n",
        "    \n",
        "    print(f\"Loading model from {ckpt}\")\n",
        "    pl_sd = torch.load(ckpt, map_location=\"cuda\")\n",
        "    if \"global_step\" in pl_sd:\n",
        "        print(f\"Global Step: {pl_sd['global_step']}\")\n",
        "    sd = pl_sd[\"state_dict\"]\n",
        "    m, u = model.load_state_dict(sd, strict=False)\n",
        "    if len(m) > 0 and verbose:\n",
        "        print(\"missing keys:\")\n",
        "        print(m)\n",
        "    if len(u) > 0 and verbose:\n",
        "        print(\"unexpected keys:\")\n",
        "        print(u)\n",
        "\n",
        "    model.cuda()\n",
        "    model.eval()\n",
        "    return model"
      ],
      "metadata": {
        "id": "-g9xz0ezmthH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "seed_everything(123)\n",
        "\n",
        "config = OmegaConf.load('configs/stable-diffusion/v2-inference-v.yaml')\n",
        "model = load_model_from_config(config, \"/content/768-v-ema.ckpt\")\n",
        "\n",
        "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "model = model.to(device)\n",
        "\n",
        "# if opt.plms:\n",
        "#     sampler = PLMSSampler(model)\n",
        "# elif opt.dpm:\n",
        "#     sampler = DPMSolverSampler(model)\n",
        "# else:\n",
        "# sampler = PLMSSampler(model)\n",
        "sampler = DDIMSampler(model)"
      ],
      "metadata": {
        "id": "HrOaD1z13X18"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "precision_scope = autocast\n",
        "prompt = \"anime illustration, highly detailed 4K, concept art of aqualium underwater, by Nikon SIGMA 30mm F1.4\"\n",
        "n_samples = 1\n",
        "start_code = torch.randn([n_samples, 4, 768 // 8, 768 // 8], device=device)\n",
        "with torch.no_grad(), \\\n",
        "    precision_scope(\"cuda\"), \\\n",
        "    model.ema_scope():\n",
        "        uc = None\n",
        "        # if opt.scale != 1.0:\n",
        "        #     uc = model.get_learned_conditioning(1 * [\"\"])\n",
        "        # if isinstance(prompt, tuple):\n",
        "        #     prompts = list(prompts)\n",
        "        uc = model.get_learned_conditioning(1 * [\"\"])\n",
        "        c = model.get_learned_conditioning([prompt])\n",
        "        shape = [4, 768 // 8, 768 // 8]\n",
        "        samples, _ = sampler.sample(S=50,\n",
        "                                    conditioning=c,\n",
        "                                    batch_size=1,\n",
        "                                    shape=shape,\n",
        "                                    verbose=False,\n",
        "                                    unconditional_guidance_scale=7.0,\n",
        "                                    unconditional_conditioning=uc,\n",
        "                                    eta=0.0,\n",
        "                                    x_T=start_code\n",
        "        )\n",
        "\n",
        "        x_samples = model.decode_first_stage(samples)\n",
        "        x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)\n",
        "\n",
        "        for x_sample in x_samples:\n",
        "            x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')\n",
        "            img = Image.fromarray(x_sample.astype(np.uint8))\n",
        "img"
      ],
      "metadata": {
        "id": "dR55lEic71b_"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}