001_stable_diffusion_gui_basic.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/karaage0703/stable-diffusion-colab-tools/blob/main/001_stable_diffusion_gui_basic.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w-0mHhP4tpTT"
      },
      "source": [
        "# Stable Diffusion GUI Basic\n",
        "Stable Diffusion easy and useful GUI basic tool\n",
        "\n",
        "GitHub repository is below:\n",
        "[stable-diffusion-colab-tools](https://github.com/karaage0703/stable-diffusion-colab-tools)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9jUMAVmhVedQ",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title **Hugging Face Login**\n",
        "#@markdown　You need access token of Hugging Face.\n",
        "!pip -qq install diffusers==0.7.2\n",
        "\n",
        "!pip -qq install transformers scipy ftfy gradio\n",
        "!pip -qq install \"ipywidgets>=7,<8\"\n",
        "\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()\n",
        "\n",
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Launch App**\n",
        "#@markdown　Execute and click URL ex: `Running on public URL: https://xxxx.gradio.app` \n",
        "\n",
        "import gradio as gr\n",
        "import torch\n",
        "from diffusers import StableDiffusionPipeline, DDIMScheduler\n",
        "\n",
        "\n",
        "#@markdown　Select model \n",
        "\n",
        "device = \"cuda\"\n",
        "model_id = \"CompVis/stable-diffusion-v1-4\" #@param [\"CompVis/stable-diffusion-v1-4\", \"hakurei/waifu-diffusion\", \"naclbit/trinart_stable_diffusion_v2\"] {allow-input: true}\n",
        "\n",
        "if model_id == \"CompVis/stable-diffusion-v1-4\":\n",
        "    pipe = StableDiffusionPipeline.from_pretrained(\n",
        "        model_id,\n",
        "        revision=\"fp16\",\n",
        "        torch_dtype=torch.float16,\n",
        "        use_auth_token=True,\n",
        "    ).to(device)\n",
        "\n",
        "if model_id == \"hakurei/waifu-diffusion\":\n",
        "    pipe = StableDiffusionPipeline.from_pretrained(\n",
        "        model_id,\n",
        "        torch_dtype=torch.float16,\n",
        "        revision=\"fp16\",\n",
        "        scheduler=DDIMScheduler(\n",
        "            beta_start=0.00085,\n",
        "            beta_end=0.012,\n",
        "            beta_schedule=\"scaled_linear\",\n",
        "            clip_sample=False,\n",
        "            set_alpha_to_one=False,\n",
        "        ),\n",
        "        use_auth_token=True,\n",
        "    ).to(device)\n",
        "\n",
        "if model_id == \"naclbit/trinart_stable_diffusion_v2\":\n",
        "    pipe = StableDiffusionPipeline.from_pretrained(\n",
        "        model_id,\n",
        "        torch_dtype=torch.float16,\n",
        "        revision=\"diffusers-60k\",\n",
        "        use_auth_token=True,\n",
        "    ).to(device)\n",
        "\n",
        "# Full model\n",
        "#pipe = StableDiffusionPipeline.from_pretrained(\n",
        "#    model_id,\n",
        "#    use_auth_token=True,\n",
        "#).to(device)\n",
        "\n",
        "\n",
        "def infer(prompt, num_images, num_inference_steps, guidance_scale_value, width_images, height_images, seed_number):\n",
        "    generator = torch.Generator(device=device)\n",
        "    latents = None\n",
        "    seeds = []\n",
        "\n",
        "    width = int(width_images)\n",
        "    height = int(height_images)\n",
        "    num_images = int(num_images)\n",
        "    num_inference_steps = int(num_inference_steps)\n",
        "    seed_number = int(seed_number)\n",
        "\n",
        "    images = []\n",
        "\n",
        "    for _ in range(num_images):\n",
        "        # Get a new random seed, store it and use it as the generator state\n",
        "        if seed_number < 0:\n",
        "            seed = generator.seed()\n",
        "        else:\n",
        "            seed = seed_number\n",
        "\n",
        "        print('seed=' + str(seed))\n",
        "        seeds.append(seed)\n",
        "        generator = generator.manual_seed(seed)\n",
        "    \n",
        "        image_latents = torch.randn(\n",
        "            (1, pipe.unet.in_channels, height // 8, width // 8),\n",
        "            generator = generator,\n",
        "            device = device\n",
        "        )\n",
        "        latents = image_latents if latents is None else torch.cat((latents, image_latents))\n",
        "\n",
        "    for latent in latents:\n",
        "        with torch.autocast('cuda'):\n",
        "            image = pipe(\n",
        "                [prompt],\n",
        "                width=width,\n",
        "                height=height,\n",
        "                guidance_scale=guidance_scale_value,\n",
        "                num_inference_steps=num_inference_steps,\n",
        "                latents = latent.unsqueeze(dim=0)\n",
        "            ).images[0]\n",
        "        images.append(image)\n",
        "\n",
        "    return images\n",
        "\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "block = gr.Blocks(css=\".container { max-width: 800px; margin: auto; }\")\n",
        "\n",
        "with block as demo:\n",
        "    gr.Markdown(\"<h1><center>Stable Diffusion Tool</center></h1>\")\n",
        "    gr.Markdown(\n",
        "        'Stable Diffusion useful web tool'\n",
        "    )\n",
        "    with gr.Group():\n",
        "        with gr.Box():\n",
        "            gr.Markdown(\n",
        "                'Enter prompt and Run!!'\n",
        "            )\n",
        "            with gr.Row().style(mobile_collapse=False, equal_height=True):\n",
        "\n",
        "                text = gr.Textbox(\n",
        "                    label='Enter prompt', show_label=False, max_lines=1\n",
        "                ).style(\n",
        "                    border=(True, False, True, True),\n",
        "                    rounded=(True, False, False, True),\n",
        "                    container=False,\n",
        "                )\n",
        "                btn = gr.Button(\"Run\").style(\n",
        "                    margin=False,\n",
        "                    rounded=(False, True, True, False),\n",
        "                )\n",
        "\n",
        "        num_images = gr.Number(\n",
        "                    label='Number of images', value=1\n",
        "                )\n",
        "\n",
        "        seed_number = gr.Number(\n",
        "                    label='Seed(-1 is random)', value=-1\n",
        "                )\n",
        "\n",
        "        num_inference_steps = gr.Slider(\n",
        "                    label='Number of inference steps', minimum=1, maximum=200, value=50\n",
        "                )\n",
        "\n",
        "        guidance_scale_value = gr.Slider(\n",
        "                    label='Guidance scale', minimum=1, maximum=20, value=7.5, step=0.1\n",
        "                )\n",
        "\n",
        "        width_images = gr.Slider(\n",
        "                    label='Width of images', minimum=64, maximum=640, value=512, step=64\n",
        "                )\n",
        "\n",
        "        height_images = gr.Slider(\n",
        "                    label='Height of images', minimum=64, maximum=640, value=512, step=64\n",
        "                )\n",
        "\n",
        "        gallery = gr.Gallery(label=\"Generated images\", show_label=False).style(\n",
        "            grid=[2], height=\"auto\"\n",
        "        )\n",
        "\n",
        "        btn.click(infer,\n",
        "                 inputs=[text, num_images, num_inference_steps, guidance_scale_value, width_images, height_images, seed_number], outputs=gallery)\n",
        "\n",
        "    gr.Markdown(\n",
        "        \"\"\"___\n",
        "   <p style='text-align: center'>\n",
        "   Created by CompVis and Stability AI\n",
        "   <br/>\n",
        "   </p>\"\"\"\n",
        "    )\n",
        "\n",
        "clear_output()\n",
        "demo.launch(debug=True)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "3O3bN_cE6LCQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Reference\n",
        "Special Thanks\n",
        "- https://note.com/npaka/n/ndd549d2ce556\n",
        "- http://cedro3.com/ai/image2image/\n",
        "- https://colab.research.google.com/github/pcuenca/diffusers-examples/blob/main/notebooks/stable-diffusion-seeds.ipynb\n",
        "- https://huggingface.co/hakurei/waifu-diffusion\n",
        "- https://huggingface.co/naclbit/trinart_stable_diffusion_v2"
      ],
      "metadata": {
        "id": "8_ieawWDdpjt"
      }
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}