SDXL colab.ipynb

sdxl-colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/YuenSzeHong/609639a44b2026f5dbf65db61eaa4523/sdxl-colab.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7wfjqqcpbmOj"
      },
      "outputs": [],
      "source": [
        "! pip install transformers accelerate safetensors diffusers invisible-watermark>=0.2.0 torch"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 载入模型\n",
        "# Define how many steps and what % of steps to be run on each experts (80/20) here\n",
        "# @markdown 每多少步增加一次\n",
        "n_steps = 40 # @param\n",
        "# @markdown 底噪\n",
        "high_noise_frac = 0.8 # @param\n",
        "# @markdown 精细图像？\n",
        "use_refiner = True # @param {type:\"boolean\"}\n",
        "\n",
        "cpu_offload = True # @param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "import torch, gc\n",
        "from diffusers import DiffusionPipeline\n",
        "\n",
        "\n",
        "# torch.no_grad()\n",
        "\n",
        "try:\n",
        "    del base\n",
        "    del refiner\n",
        "except NameError:\n",
        "    pass\n",
        "\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "# load both base & refiner\n",
        "base = DiffusionPipeline.from_pretrained(\n",
        "    \"stabilityai/stable-diffusion-xl-base-1.0\",\n",
        "    torch_dtype=torch.float16,\n",
        "\n",
        "    variant=\"fp16\",\n",
        "    use_safetensors=True\n",
        ")\n",
        "\n",
        "# base.unet = torch.compile(base.unet, mode=\"reduce-overhead\", fullgraph=True)\n",
        "base.to('cuda')\n",
        "\n",
        "if use_refiner:\n",
        "    refiner = DiffusionPipeline.from_pretrained(\n",
        "        \"stabilityai/stable-diffusion-xl-refiner-1.0\",\n",
        "        text_encoder_2=base.text_encoder_2,\n",
        "        vae=base.vae,\n",
        "        torch_dtype=torch.float16,\n",
        "        use_safetensors=True,\n",
        "        variant=\"fp16\",\n",
        "    )\n",
        "\n",
        "    # refiner.unet = torch.compile(refiner.unet, mode=\"reduce-overhead\", fullgraph=True)\n",
        "\n",
        "    if cpu_offload:\n",
        "        refiner.enable_model_cpu_offload()"
      ],
      "metadata": {
        "id": "jNrJ7SrqTrPO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 生成图片 { vertical-output: true }\n",
        "from IPython.display import display, Image\n",
        "import math\n",
        "# @markdown 提示词\n",
        "prompt = \"An anime girl wearing a full-body swimsuit, with long white hair tinged with pink with a cowlick, lying on several big yoga balls of different colors, from a side view, in a play room. \" # @param {type:\"string\"}\n",
        "# @markdown 负面提示词\n",
        "negative_prompt = \"(worst quality, low quality:1.4), (malformed hands:1.4),(poorly drawn hands:1.4),(mutated fingers:1.4),(extra limbs:1.35), (poorly drawn face:1.4), (wide forehead:1.4),\" # @param {type:\"string\"}\n",
        "# @markdown 图片数量：基本只能每次一张\n",
        "no_of_image = 1 # @param {type:\"integer\"}\n",
        "# @markdown 图片比例 (w:h)\n",
        "aspect_ratio = \"9:16\" # @param {type:\"string\"}\n",
        "# @markdown 图片最大总像素数 (最好别动，经过调试)\n",
        "max_pixel_side = 1024 # @param {type:\"slider\", min:256, max:1024, step:1}\n",
        "max_pixel = max_pixel_side * max_pixel_side\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "w, h = (int(e) for e in aspect_ratio.split(':'))\n",
        "width, height = (round(math.sqrt(max_pixel * x / y) / 8) * 8 for x, y in ((w, h), (h, w)))\n",
        "del w\n",
        "del h\n",
        "print(width, height)\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "# run both experts\n",
        "image = base(\n",
        "    prompt=prompt,\n",
        "    num_inference_steps=n_steps,\n",
        "    negative_prompt=negative_prompt,\n",
        "    width=width, height=height,\n",
        "    num_images_per_prompt=no_of_image,\n",
        "    denoising_end=high_noise_frac if use_refiner else 1.0,\n",
        "    output_type=\"latent\" if use_refiner else \"pil\",\n",
        ").images\n",
        "torch.cuda.empty_cache()\n",
        "if use_refiner:\n",
        "    images = refiner(\n",
        "        prompt=prompt,\n",
        "        num_inference_steps=n_steps,\n",
        "        denoising_start=high_noise_frac if use_refiner else 0.0,\n",
        "        image=image,\n",
        "    ).images\n",
        "if not use_refiner:\n",
        "    images = image\n",
        "del width\n",
        "del height\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "for image in images:\n",
        "    display(image)\n",
        "    print(\"\\n\\n\\n\")"
      ],
      "metadata": {
        "id": "P2brCDpdUmUh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 清显存\n",
        "try:\n",
        "    del images\n",
        "    if use_refiner:\n",
        "        del image\n",
        "except NameError:\n",
        "    pass\n",
        "\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "ysrddYi3gFmA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 保存图片\n",
        "import uuid\n",
        "!mkdir -p images\n",
        "images[0].save(f\"images/{str(uuid.uuid4())}.png\")\n",
        "\n",
        "! rm .ipynb_checkpoints\n",
        "\n",
        "del images, image\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "BsZjDsqCNWuN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 打包保存\n",
        "\n",
        "!zip images images\n",
        "from google.colab import files\n",
        "files.download(\"images.zip\")"
      ],
      "metadata": {
        "id": "ckmknKK4QZfk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title 保存Google Drive\n",
        "\n",
        "google_drive_path = \"\" # @param {type:\"string\"}\n",
        "\n",
        "use_shared_drive = False # @param {type:\"boolean\"}\n",
        "\n",
        "shared_drive_name = \"\" # @param {type:\"string\"}\n",
        "\n",
        "if not use_shared_drive:\n",
        "  path = f\"drive/MyDrive{google_drive_path}\"\n",
        "else:\n",
        "  path = f\"drive/Shareddrives/{shared_drive_name}/{google_drive_path}\"\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "!cp -r images $path/SDXL\\ images"
      ],
      "metadata": {
        "id": "X2r2qxSJRlLQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "! pip install pipreqs"
      ],
      "metadata": {
        "id": "u6yawAV9Z1bu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "! pipreqs ."
      ],
      "metadata": {
        "id": "KkxWSw67Z-hA"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4",
      "cell_execution_strategy": "setup",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}