Stable_Diffusion_with_diffusers_(text2img,_img2img).ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "# @title Install libraries\n",
        "# @markdown This cell will take a while because you have to download multiple libraries\n",
        " \n",
        "!pip install transformers scipy ftfy \"ipywidgets>=7,<8\"\n",
        "!git clone https://github.com/huggingface/diffusers.git\n",
        "!pip install git+https://github.com/huggingface/diffusers.git\n",
        "cd diffusers"
      ],
      "metadata": {
        "id": "K5ELsv-g0FPb",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Login in Hugging Face\n",
        "# @markdown Login in hugginggace and take in account setting access token https://huggingface.co/ .\n",
        "\n",
        "# @markdown And in this link https://huggingface.co/CompVis/stable-diffusion-v1-4 you must accept the agreement on the use of this model (check the box).\n",
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ],
      "metadata": {
        "id": "w8lDIyRB0e-l",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Import libraries and define things\n",
        "import os\n",
        "import torch\n",
        "import shutil\n",
        "from PIL import Image\n",
        "from torch import autocast\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.image as mpimg\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from examples.inference.image_to_image import StableDiffusionImg2ImgPipeline, preprocess\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"fp16\", torch_dtype=torch.float16, use_auth_token=True).to(\"cuda\")\n",
        "pipeimg = StableDiffusionImg2ImgPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"fp16\", torch_dtype=torch.float16, use_auth_token=True).to(\"cuda\")"
      ],
      "metadata": {
        "id": "fwAguL3o0j8K",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title (Run after from here to make changes)\n",
        "# @markdown Your prompt and size image\n",
        "PROMPT = \"Test\" #@param {type:\"string\"}\n",
        "# @markdown i recommend use small size of image because Free Colab don't have more GPU\n",
        "HEIGHT = 512 #@param {type:\"integer\"}\n",
        "WIDTH = 512 #@param {type:\"integer\"}\n",
        "# @markdown Number of example\n",
        "EXAMPLES = 1 #@param {type:\"integer\"}\n",
        "# @markdown Number of step (recommend 50-100)\n",
        "STEPS = 50 #@param {type:\"integer\"}\n",
        "GUIDANCE_SCALE = 7.5 #@param {type:\"number\"}\n",
        "# @markdown for image2image generating\n",
        "STRENGTH = 0.75 #@param {type:\"number\"}\n",
        "INIT_IMAGE = \"\" #@param {type:\"string\"}\n",
        "#SAVE_DIR = \"\"\n",
        "try:\n",
        "  if not PROMPT.strip(): #если строка промпта пустая\n",
        "    print(\"Please write your prompt\")\n",
        "    pass\n",
        "  else: #если не пустая\n",
        "    for i in range(EXAMPLES): #кол-во примеров\n",
        "      if not INIT_IMAGE.strip(): #если изображение пусто\n",
        "        try:\n",
        "          shutil.rmtree(f'{PROMPT}') #Удаляет папку если она есть\n",
        "        except FileNotFoundError:\n",
        "          pass\n",
        "        os.mkdir(f'{PROMPT}')\n",
        "        with autocast(\"cuda\"):\n",
        "          plt.figure(figsize=(8,8))\n",
        "          image = pipe(PROMPT, num_inference_steps=STEPS, height=HEIGHT, width=WIDTH, guidance_scale=GUIDANCE_SCALE)[\"sample\"][0]\n",
        "          image.save(f'{PROMPT}/test{i}.png')\n",
        "          imgi = mpimg.imread(f'{PROMPT}/test{i}.png')\n",
        "          imgplot = plt.imshow(imgi)\n",
        "          plt.axis('off')\n",
        "          plt.show()\n",
        "      else: #если изображение есть\n",
        "        init_image = Image.open(INIT_IMAGE).convert(\"RGB\")\n",
        "        init_image = init_image.resize((WIDTH, HEIGHT))\n",
        "        init_image = preprocess(init_image)\n",
        "        prompt_save = \"init image__\" + prompt\n",
        "        try:\n",
        "          shutil.rmtree(f'{prompt_save}') \n",
        "        except FileNotFoundError:\n",
        "          pass\n",
        "        os.mkdir(f'{prompt_save}')\n",
        "        with autocast(\"cuda\"):\n",
        "          plt.figure(figsize=(8,8))\n",
        "          image = pipeimg(PROMPT, init_image=init_image, strength=STRENGTH, num_inference_steps=STEPS, guidance_scale=GUIDANCE_SCALE)[\"sample\"][0]\n",
        "          image.save(f'{prompt_save}/test{i}.png')\n",
        "          imgi = mpimg.imread(f'{prompt_save}/test{i}.png')\n",
        "          imgplot = plt.imshow(imgi)\n",
        "          plt.axis('off')\n",
        "          plt.show()\n",
        "except:\n",
        "  pass"
      ],
      "metadata": {
        "id": "YjWTUqOG0zRM",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}