HighQuality_Text2Image_Stable_Diffusion_ls.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/HemantKArya/HqStableDiffusionColab/blob/main/HighQuality_Text2Image_Stable_Diffusion_ls.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "upLwll2nS9QL"
      },
      "source": [
        "#**High Quality Text to Image Generation using Stable Diffusion,Real-ESR and Swin IR**\n",
        "[![github](https://img.shields.io/badge/GitHub-100000?style=for-the-badge&logo=github&logoColor=white)](https://github.com/HemantKArya/HqStableDiffusionColab)\n",
        "\n",
        "Generate 4K and FULL HD Images and Artworks for Free Using Stable Diffusion. \n",
        "\n",
        "Visit Logical Spot for Video Help:-\n",
        "\n",
        " [![YouTube](https://img.shields.io/badge/YouTube-%23FF0000.svg?style=for-the-badge&logo=YouTube&logoColor=white)](https://www.youtube.com/c/LogicalSpot)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gd-vX3cavOCt"
      },
      "source": [
        "# **Stable Diffusion** 🎨\n",
        "*...using `🧨diffusers`*\n",
        "\n",
        "Stable Diffusion is a text-to-image latent diffusion model created by the researchers and engineers from [CompVis](https://github.com/CompVis), [Stability AI](https://stability.ai/) and [LAION](https://laion.ai/). It's trained on 512x512 images from a subset of the [LAION-5B](https://laion.ai/blog/laion-5b/) database. This model uses a frozen CLIP ViT-L/14 text encoder to condition the model on text prompts. With its 860M UNet and 123M text encoder, the model is relatively lightweight and runs on a GPU with at least 10GB VRAM.\n",
        "See the [model card](https://huggingface.co/CompVis/stable-diffusion) for more information.\n",
        "\n",
        "This Colab notebook shows how to use Stable Diffusion with the 🤗 Hugging Face [🧨 Diffusers library](https://github.com/huggingface/diffusers).\n",
        "https://github.com/CompVis/stable-diffusion\n",
        "\n",
        "orignal-link to colab https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb\n",
        "\n",
        "Let's get started!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zsMRD1FeQVvM"
      },
      "source": [
        "# **Real-ESRGAN**\n",
        "[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2107.10833)\n",
        "[![GitHub Stars](https://img.shields.io/github/stars/xinntao/Real-ESRGAN?style=social)](https://github.com/xinntao/Real-ESRGAN)\n",
        "[![download](https://img.shields.io/github/downloads/xinntao/Real-ESRGAN/total.svg)](https://github.com/xinntao/Real-ESRGAN/releases)\n",
        "\n",
        "Real-ESRGAN aims at developing Practical Algorithms for General Image/Video Restoration.We extend the powerful ESRGAN to a practical restoration application (namely, Real-ESRGAN), which is trained with pure synthetic data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ykPIKg_tRzG8"
      },
      "source": [
        "# **SwinIR**\n",
        "[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2108.10257)\n",
        "[![GitHub Stars](https://img.shields.io/github/stars/JingyunLiang/SwinIR?style=social)](https://github.com/JingyunLiang/SwinIR)\n",
        "[![download](https://img.shields.io/github/downloads/JingyunLiang/SwinIR/total.svg)](https://github.com/JingyunLiang/SwinIR/releases)\n",
        "\n",
        "SwinIR achieves state-of-the-art performance on six tasks: image super-resolution (including classical, lightweight and real-world image super-resolution), image denoising (including grayscale and color image denoising) and JPEG compression artifact reduction. See our [paper](https://arxiv.org/abs/2108.10257) and [project page](https://github.com/JingyunLiang/SwinIR) for detailed results."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uFq1WRZFdayX"
      },
      "source": [
        "### (No colorization; No CUDA extensions required)\n",
        "\n",
        "[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2101.04061)\n",
        "[![GitHub Stars](https://img.shields.io/github/stars/TencentARC/GFPGAN?style=social)](https://github.com/TencentARC/GFPGAN)\n",
        "[![download](https://img.shields.io/github/downloads/TencentARC/GFPGAN/total.svg)](https://github.com/TencentARC/GFPGAN/releases)\n",
        "\n",
        "## GFPGAN - Towards Real-World Blind Face Restoration with Generative Facial Prior\n",
        "\n",
        "GFPGAN is a blind face restoration algorithm towards real-world face images. <br>\n",
        "It leverages the generative face prior in a pre-trained GAN (*e.g.*, StyleGAN2) to restore realistic faces while precerving fidelity. <br>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYOlvQ1nQL7c"
      },
      "source": [
        "### Setup\n",
        "\n",
        "First, please make sure you are using a GPU runtime to run this notebook, so inference is much faster. If the following command fails, use the `Runtime` menu above and select `Change runtime type`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zHkHsdtnry57"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "paJt_cx5QgVz"
      },
      "source": [
        "## Next, you should install `diffusers==0.3.0` as well `scipy`, `ftfy` and `transformers`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aIrgth7sqFML"
      },
      "outputs": [],
      "source": [
        "!pip install diffusers==0.10.0\n",
        "!pip install transformers scipy ftfy accelerate\n",
        "!pip install \"ipywidgets>=7,<8\"\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Fcyyt0daU4e"
      },
      "source": [
        "You also need to accept the model license before downloading or using the weights. In this post we'll use model version `v1-4`, so you'll need to  visit [its card](https://huggingface.co/CompVis/stable-diffusion-v1-4), read the license and tick the checkbox if you agree.\n",
        "\n",
        "You have to be a registered user in 🤗 Hugging Face Hub, and you'll also need to use an access token for the code to work. For more information on access tokens, please refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ou0Ijygormum"
      },
      "source": [
        "As google colab has disabled external widgtes, we need to enable it explicitly. Run the following cell to be able to use `notebook_login`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3NnPOMAqAABv"
      },
      "source": [
        "# Stable Diffusion Pipeline\n",
        "\n",
        "`StableDiffusionPipeline` is an end-to-end inference pipeline that you can use to generate images from text with just a few lines of code.\n",
        "\n",
        "First, we load the pre-trained weights of all components of the model.\n",
        "\n",
        "In addition to the model id [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4), we're also passing a specific `revision`, `torch_dtype` and `use_auth_token` to the `from_pretrained` method.\n",
        "`use_auth_token` is necessary to verify that you have indeed accepted the model's license.\n",
        "\n",
        "We want to ensure that every free Google Colab can run Stable Diffusion, hence we're loading the weights from the half-precision branch [`fp16`](https://huggingface.co/CompVis/stable-diffusion-v1-4/tree/fp16) and also tell `diffusers` to expect the weights in float16 precision by passing `torch_dtype=torch.float16`.\n",
        "\n",
        "If you want to ensure the highest possible precision, please make sure to remove `revision=\"fp16\"` and `torch_dtype=torch.float16` at the cost of a higher memory usage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xSKWBKFPArKS"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import glob\n",
        "os.environ['PYTORCH_CUDA_ALLOC_CONF'] = \"max_split_size_mb:100\"\n",
        "import torch\n",
        "import numpy as np\n",
        "import gc\n",
        "from diffusers import StableDiffusionPipeline\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9CS0e09d9WoJ"
      },
      "outputs": [],
      "source": [
        "# make sure you're logged in with `huggingface-cli login`\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"fp16\", torch_dtype=torch.float16)  \n",
        "pipe = pipe.to(\"cuda\")\n",
        "# img2imgpipe = StableDiffusionImg2ImgPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\",revision=\"fp16\", torch_dtype=torch.float16,use_auth_token=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setting  Up Real-ESR and SwinIR"
      ],
      "metadata": {
        "id": "fD3ZVMLQEnHR"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ldl84rPaKCrb"
      },
      "outputs": [],
      "source": [
        "\n",
        "!git clone https://github.com/TencentARC/GFPGAN.git\n",
        "os.chdir(\"GFPGAN\")\n",
        "\n",
        "\n",
        "!pip install basicsr\n",
        "\n",
        "!pip install facexlib\n",
        "!pip install -r requirements.txt\n",
        "!python setup.py develop\n",
        "os.chdir(\"..\")\n",
        "\n",
        "#realesr\n",
        "!git clone https://github.com/xinntao/Real-ESRGAN.git\n",
        "# !pip install realesrgan\n",
        "!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /content/Real-ESRGAN/experiments/pretrained_models\n",
        "os.chdir(\"Real-ESRGAN\")\n",
        "!pip install gfpgan\n",
        "!pip install -r requirements.txt\n",
        "!python setup.py develop\n",
        "!python setup.py install\n",
        "os.chdir(\"..\")\n",
        "\n",
        "\n",
        "!wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P /content/GFPGAN/experiments/pretrained_models\n",
        "!git clone https://github.com/JingyunLiang/SwinIR.git\n",
        "!pip install timm\n",
        "!wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth -P experiments/pretrained_models\n",
        "!wget https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth -P experiments/pretrained_models\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_DM2JSo6v-1J"
      },
      "outputs": [],
      "source": [
        "os.chdir(\"Real-ESRGAN\")\n",
        "from realesrgan import RealESRGANer\n",
        "from basicsr.archs.rrdbnet_arch import RRDBNet\n",
        "esrdevice='cuda'\n",
        "with torch.no_grad():\n",
        "  RealESRUpScale = RealESRGANer(model_path=\"./experiments/pretrained_models/RealESRGAN_x4plus.pth\",scale=4,device=esrdevice,model= RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4))\n",
        "os.chdir(\"..\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8MgNzTxwbASv"
      },
      "source": [
        "Next, let's move the pipeline to GPU to have faster inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HVC-aLwl8BBX"
      },
      "outputs": [],
      "source": [
        "\n",
        "from google.colab import files\n",
        "import shutil\n",
        "\n",
        "upload_folder = 'upload'\n",
        "result_folder = 'results'\n",
        "final_folder = 'final'\n",
        "forswin_folder = 'forswin'\n",
        "\n",
        "if os.path.isdir(upload_folder):\n",
        "    shutil.rmtree(upload_folder)\n",
        "if os.path.isdir(result_folder):\n",
        "    shutil.rmtree(result_folder)\n",
        "if os.path.isdir(final_folder):\n",
        "    shutil.rmtree(final_folder)\n",
        "if os.path.isdir(forswin_folder):\n",
        "    shutil.rmtree(forswin_folder)\n",
        "os.mkdir(upload_folder)\n",
        "os.mkdir(result_folder)\n",
        "os.mkdir(final_folder)\n",
        "os.mkdir(forswin_folder)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ZcgsflpBoEM"
      },
      "source": [
        "\n",
        "\n",
        "Let's first write a helper function to display a grid of images. Just run the following cell to create the `image_grid` function, or disclose the code if you are interested in how it's done."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "REF_yuHprSa1"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "\n",
        "\n",
        "def imread(img_path):\n",
        "  img = cv2.imread(img_path)\n",
        "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "  return img\n",
        "\n",
        "def display_com(img_paths,gen_paths,output_name=\"Result\"):\n",
        "  fig,axs = plt.subplots(len(img_paths),2,figsize=(20,30))\n",
        "  axs[0,0].set_title(\"Input\", fontsize=16)\n",
        "  axs[0,1].set_title(output_name, fontsize=16)\n",
        "  for i,j,k in zip(range(len(img_paths)),img_paths,gen_paths):\n",
        "    axs[i,0].imshow(imread(j))\n",
        "    axs[i,1].imshow(imread(k))\n",
        "    axs[i,0].axis(\"off\")\n",
        "    axs[i,1].axis(\"off\")\n",
        "  plt.subplots_adjust(wspace=0, hspace=0, left=0, right=1, bottom=0, top=1)\n",
        "  plt.show()\n",
        "\n",
        "\n",
        "def image_grid(imgs, rows, cols):\n",
        "    assert len(imgs) == rows*cols\n",
        "\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new('RGB', size=(cols*w, rows*h))\n",
        "    grid_w, grid_h = grid.size\n",
        "\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i%cols*w, i//cols*h))\n",
        "    return grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AcHccTDWbQRU"
      },
      "source": [
        "# Now, we can generate a grid image once having run the pipeline with a list of 4 prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DIzYXn35bC_G"
      },
      "outputs": [],
      "source": [
        "y_prompt = [\"A Ring made of Cherry Blossom flower\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-YAFLvWWrSdM"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "from torch import autocast\n",
        "\n",
        "dir = './upload'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "num_images = 3\n",
        "prompt = y_prompt * num_images\n",
        "rn = random.randint(0,10000000)\n",
        "\n",
        "# rn=1624589\n",
        "\n",
        "print(rn)\n",
        "with torch.no_grad():\n",
        "  generator = torch.Generator(\"cuda\").manual_seed(rn)\n",
        "  with autocast(\"cuda\"):\n",
        "    images = pipe(prompt,generator=generator, num_inference_steps=50).images\n",
        "    for i,j in enumerate(images):\n",
        "      j.save(\"./upload/img{}-{}\".format(i,rn)+'.png')\n",
        "images_2up = './upload'\n",
        "torch.cuda.empty_cache()\n",
        "grid = image_grid(images, rows=1, cols=3)\n",
        "grid\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SJk5KfEdDBYu"
      },
      "source": [
        "# Refine Same Images(Generated in Prev step) with 100 steps. You Can Directly jump to the next Cell (Optional: if results are weird)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jBX2VVz9BeJZ"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "generator = torch.Generator(\"cuda\").manual_seed(rn)\n",
        "\n",
        "dir = './upload'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "with autocast(\"cuda\"):\n",
        "  images = pipe(prompt,generator=generator, num_inference_steps=100).images\n",
        "  for i,j in enumerate(images):\n",
        "    j.save(\"./upload/img{}-{}\".format(i,rn)+'.png')\n",
        "images_2up = './upload'\n",
        "torch.cuda.empty_cache()\n",
        "grid = image_grid(images, rows=1, cols=3)\n",
        "grid\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ch8UY6Q_DO2z"
      },
      "source": [
        "# Refine Same Images(Generated in Prev step) with 200 steps. (Optional: if results are weird)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "w1SE6qi5BtxR"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "generator = torch.Generator(\"cuda\").manual_seed(rn)\n",
        "\n",
        "dir = './upload'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "\n",
        "with autocast(\"cuda\"):\n",
        "  images = pipe(prompt,generator=generator, num_inference_steps=200).images\n",
        "  for i,j in enumerate(images):\n",
        "    j.save(\"./upload/img{}-{}\".format(i,rn)+'.png')\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "images_2up = './upload'\n",
        "grid = image_grid(images, rows=1, cols=3)\n",
        "grid\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qfnEYcfAQwgn"
      },
      "source": [
        "# **Run This Cell Only if Face is not Restored Completely.** (Optional: if faces are destorted) "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YiCQvTuxQEWu"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "# os.chdir(\"/content/GFPGAN\")\n",
        "dir = './results/restored_imgs'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "!python GFPGAN/inference_gfpgan.py -i /content/upload -o /content/results -s 2\n",
        "# os.chdir(\"/content\")\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "# display each image in the upload folder\n",
        "\n",
        "images_2up = './results/restored_imgs'\n",
        "\n",
        "input_folder = './upload'\n",
        "result_folder = './results/restored_imgs'\n",
        "input_list = sorted(glob.glob(os.path.join(input_folder, '*.png')))\n",
        "output_list = sorted(glob.glob(os.path.join(result_folder, '*.png')))\n",
        "print(len(input_list))\n",
        "print(len(output_list))\n",
        "display_com(input_list,output_list,output_name=\"GFPGAN\")\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5zQLbDEV84wZ"
      },
      "source": [
        "# **Real-ESRGAN Upscale**\n",
        "Make Sure to download all result from **final folder**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "Sme9aQnIJiC9"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "dir = './final'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "# images_2up = \"results/restored_imgs\"\n",
        "imgs_path = sorted(glob.glob(os.path.join(images_2up, '*.png')))\n",
        "imgs = []\n",
        "for path in imgs_path:\n",
        "  imgs.append(Image.open(path).convert('RGB'))\n",
        "\n",
        "\n",
        "for i,j in enumerate(imgs):\n",
        "  with torch.no_grad():\n",
        "    output, _ = RealESRUpScale.enhance(np.array(j), outscale=8)\n",
        "  output = Image.fromarray(output)\n",
        "  output.save(\"final/fin_img_upscaled_{}_{}.jpg\".format(rn,i),quality=98)\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "input_folder = images_2up\n",
        "result_folder = './final'\n",
        "input_list = sorted(glob.glob(os.path.join(input_folder, '*.*')))\n",
        "output_list = sorted(glob.glob(os.path.join(result_folder, '*.*')))\n",
        "display_com(input_list,output_list,output_name=\"Real-ESRGAN\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Reiik2GxFfRP"
      },
      "source": [
        "# **(Optional) Denoise Using SwinIR and Upscale**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "JDZSZ66VW1cr"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "\n",
        "\n",
        "dir = './forswin'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "dir = './results/swinir_real_sr_x4'\n",
        "if os.path.isdir(dir):\n",
        "  for f in os.listdir(dir):\n",
        "    try:\n",
        "      os.remove(os.path.join(dir, f))\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "toDown_list = sorted(glob.glob(os.path.join(images_2up, '*.png')))\n",
        "for path in toDown_list:\n",
        "  img = Image.open(path)\n",
        "  img = img.resize((512,512))   #comment out this if not getting supreme quality\n",
        "  img.save('./forswin/{}'.format(os.path.basename(path)))\n",
        "\n",
        "# os.chdir(\"/content/SwinIR\")\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "\n",
        "\n",
        "!python SwinIR/main_test_swinir.py --task real_sr --model_path ./experiments/pretrained_models/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.pth --folder_lq forswin --scale 4\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "gc.collect()\n",
        "# display each image in the upload folder\n",
        "\n",
        "input_folder = images_2up\n",
        "result_folder = './results/swinir_real_sr_x4'\n",
        "input_list = sorted(glob.glob(os.path.join(input_folder, '*.png')))\n",
        "output_list = sorted(glob.glob(os.path.join(result_folder, '*.png')))\n",
        "display_com(input_list,output_list,output_name=\"Swin\")\n",
        "print(len(input_list))\n",
        "print(len(output_list))\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "esIeu0OXIPJQ"
      },
      "source": [
        "# **(Optional) Compare ALL Generated images**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WiHUZ9QAGr0l"
      },
      "outputs": [],
      "source": [
        "def imread(img_path):\n",
        "  img = cv2.imread(img_path)\n",
        "  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "  return img\n",
        "def compare_all(img_input,img_swin,img_esr):\n",
        "  fig,axs = plt.subplots(len(img_input),3,figsize=(35,35))\n",
        "  axs[0,0].set_title(\"Input\", fontsize=16)\n",
        "  axs[0,1].set_title(\"Swin\", fontsize=16)\n",
        "  axs[0,2].set_title(\"ESRGAN\", fontsize=16)\n",
        "  for i in range(len(img_input)):\n",
        "    axs[i,0].imshow(imread(img_input[i]))\n",
        "    axs[i,1].imshow(imread(img_swin[i]))\n",
        "    axs[i,2].imshow(imread(img_esr[i]))\n",
        "  for x in axs.flatten():\n",
        "    x.axis(\"off\")\n",
        "  plt.subplots_adjust(wspace=0, hspace=0, left=0, right=1, bottom=0, top=1)\n",
        "  plt.show()\n",
        "\n",
        "input_folder = './upload'\n",
        "swin_folder = './results/swinir_real_sr_x4'\n",
        "esr_folder = './final'\n",
        "input_list = sorted(glob.glob(os.path.join(input_folder, '*.png')))\n",
        "swin_list = sorted(glob.glob(os.path.join(swin_folder, '*.png')))\n",
        "esr_list = sorted(glob.glob(os.path.join(esr_folder, '*.png')))\n",
        "print(len(input_list))\n",
        "print(len(swin_list))\n",
        "print(len(esr_list))\n",
        "\n",
        "compare_all(input_list,swin_list,esr_list)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c1IZIMhtR-vV"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "paJt_cx5QgVz",
        "3NnPOMAqAABv",
        "fD3ZVMLQEnHR",
        "SJk5KfEdDBYu",
        "Ch8UY6Q_DO2z",
        "qfnEYcfAQwgn",
        "Reiik2GxFfRP",
        "esIeu0OXIPJQ"
      ],
      "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
}