Kohya_SDXL_Inference.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/Kohya_SDXL_Inference.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r1d1SH3SsMs0"
      },
      "source": [
        "[![visitor][visitor-badge]][visitor-stats]\n",
        "<!-- [![ko-fi][ko-fi-badge]][ko-fi-link] -->\n",
        "\n",
        "# **Kohya Imagen XL**\n",
        "Colab Notebook for SDXL inference\n",
        "\n",
        "[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=Kohya%20Imagen%20XL&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
        "[visitor-stats]: https://visitorbadge.io/status?path=Kohya%20Imagen%20XL\n",
        "[ko-fi-badge]: https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat\n",
        "[ko-fi-link]: https://ko-fi.com/linaqruf\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VjHXkMhMJfZr"
      },
      "outputs": [],
      "source": [
        "# @title ## **1. Install Environment**\n",
        "import os\n",
        "import zipfile\n",
        "import shutil\n",
        "import time\n",
        "import requests\n",
        "import torch\n",
        "from subprocess import getoutput\n",
        "from IPython.utils import capture\n",
        "from google.colab import drive\n",
        "\n",
        "%store -r\n",
        "\n",
        "# root_dir\n",
        "root_dir          = \"/content\"\n",
        "drive_dir         = os.path.join(root_dir, \"drive/MyDrive\")\n",
        "deps_dir          = os.path.join(root_dir, \"deps\")\n",
        "repo_dir          = os.path.join(root_dir, \"kohya-trainer\")\n",
        "training_dir      = os.path.join(root_dir, \"LoRA\")\n",
        "models_dir        = os.path.join(root_dir, \"models\")\n",
        "vae_dir           = os.path.join(root_dir, \"vae\")\n",
        "lora_dir          = os.path.join(root_dir, \"network_weight\")\n",
        "repositories_dir  = os.path.join(root_dir, \"repositories\")\n",
        "config_dir        = os.path.join(training_dir, \"config\")\n",
        "tools_dir         = os.path.join(repo_dir, \"tools\")\n",
        "finetune_dir      = os.path.join(repo_dir, \"finetune\")\n",
        "accelerate_config = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
        "\n",
        "for store in [\"root_dir\", \"repo_dir\", \"training_dir\", \"models_dir\", \"vae_dir\", \"repositories_dir\", \"accelerate_config\", \"tools_dir\", \"finetune_dir\", \"config_dir\"]:\n",
        "    with capture.capture_output() as cap:\n",
        "        %store {store}\n",
        "        del cap\n",
        "\n",
        "repo_dict = {\n",
        "    \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" : \"https://github.com/qaneel/kohya-trainer\",\n",
        "    \"kohya-ss/sd-scripts (original repo, latest update)\"                    : \"https://github.com/kohya-ss/sd-scripts\",\n",
        "}\n",
        "\n",
        "repository        = \"kohya-ss/sd-scripts (original repo, latest update)\" #@param [\"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\", \"kohya-ss/sd-scripts (original repo, latest update)\"] {allow-input: true}\n",
        "repo_url          = repo_dict[repository]\n",
        "branch            = \"sdxl\"  # @param {type: \"string\"}\n",
        "output_to_drive   = False  # @param {type: \"boolean\"}\n",
        "\n",
        "def clone_repo(url, dir, branch):\n",
        "    if not os.path.exists(dir):\n",
        "       !git clone -b {branch} {url} {dir}\n",
        "\n",
        "def mount_drive(dir):\n",
        "    output_dir      = os.path.join(training_dir, \"output\")\n",
        "\n",
        "    if output_to_drive:\n",
        "        if not os.path.exists(drive_dir):\n",
        "            drive.mount(os.path.dirname(drive_dir))\n",
        "        output_dir  = os.path.join(drive_dir, \"kohya-trainer\", \"output\")\n",
        "\n",
        "    return output_dir\n",
        "\n",
        "def setup_directories():\n",
        "    global output_dir\n",
        "\n",
        "    output_dir      = mount_drive(drive_dir)\n",
        "\n",
        "    for dir in [training_dir, config_dir, models_dir, vae_dir, repositories_dir, output_dir]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "def install_dependencies():\n",
        "    gpu_info          = getoutput('nvidia-smi')\n",
        "    t4_xformers_wheel = \"https://github.com/Linaqruf/colab-xformers/releases/download/0.0.20/xformers-0.0.20+1d635e1.d20230519-cp310-cp310-linux_x86_64.whl\"\n",
        "\n",
        "    !apt install aria2 -yqq\n",
        "    !pip install -q --upgrade diffusers[torch]==0.18.2 transformers==4.30.2 einops==0.6.0 open-clip-torch==2.20.0 -e .\n",
        "\n",
        "    if '2.0.1+cu118' in torch.__version__:\n",
        "        if 'T4' in gpu_info:\n",
        "            !pip install -q {t4_xformers_wheel}\n",
        "        else:\n",
        "            !pip install -q xformers==0.0.20\n",
        "    else:\n",
        "        !pip install -q torch==2.0.0+cu118 torchvision==0.15.1+cu118 torchaudio==2.0.1+cu118 torchtext==0.15.1 torchdata==0.6.0 --extra-index-url https://download.pytorch.org/whl/cu118 -U\n",
        "        !pip install -q xformers==0.0.19 triton==2.0.0 -U\n",
        "\n",
        "    from accelerate.utils import write_basic_config\n",
        "\n",
        "    if not os.path.exists(accelerate_config):\n",
        "        write_basic_config(save_location=accelerate_config)\n",
        "\n",
        "def prepare_environment():\n",
        "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
        "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
        "    os.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n",
        "\n",
        "def main():\n",
        "    os.chdir(root_dir)\n",
        "    clone_repo(repo_url, repo_dir, branch)\n",
        "    os.chdir(repo_dir)\n",
        "    setup_directories()\n",
        "    install_dependencies()\n",
        "    prepare_environment()\n",
        "\n",
        "main()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NW48qR65KAm6"
      },
      "outputs": [],
      "source": [
        "# @title ## **2. Download SDXL**\n",
        "import os\n",
        "import re\n",
        "import json\n",
        "import glob\n",
        "import gdown\n",
        "import requests\n",
        "import subprocess\n",
        "from IPython.utils import capture\n",
        "from urllib.parse import urlparse, unquote\n",
        "from pathlib import Path\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Place your Huggingface [Read Token](https://huggingface.co/settings/tokens) Here. Get your SDXL access [here](https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9).\n",
        "\n",
        "HUGGINGFACE_TOKEN = \"\"#@param {type: \"string\"}\n",
        "SDXL_MODEL_URL = \"https://huggingface.co/Linaqruf/sdxl_finetune/resolve/main/hitokomoru-xl-2500.safetensors\" #@param {type: \"string\"}\n",
        "SDXL_VAE_URL = \"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors\" #@param {type: \"string\"}\n",
        "SDXL_LORA_URL = \"\" #@param {type: \"string\"}\n",
        "\n",
        "def get_supported_extensions():\n",
        "    return tuple([\".ckpt\", \".safetensors\", \".pt\", \".pth\"])\n",
        "\n",
        "def get_filename(url, quiet=True):\n",
        "    extensions = get_supported_extensions()\n",
        "\n",
        "    if url.startswith(drive_dir) or url.endswith(tuple(extensions)):\n",
        "        filename = os.path.basename(url)\n",
        "    else:\n",
        "        response = requests.get(url, stream=True)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        if 'content-disposition' in response.headers:\n",
        "            content_disposition = response.headers['content-disposition']\n",
        "            filename = re.findall('filename=\"?([^\"]+)\"?', content_disposition)[0]\n",
        "        else:\n",
        "            url_path = urlparse(url).path\n",
        "            filename = unquote(os.path.basename(url_path))\n",
        "\n",
        "    if filename.endswith(tuple(get_supported_extensions())):\n",
        "        return filename\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "def get_most_recent_file(directory):\n",
        "    files = glob.glob(os.path.join(directory, \"*\"))\n",
        "    if not files:\n",
        "        return None\n",
        "    most_recent_file = max(files, key=os.path.getmtime)\n",
        "    basename = os.path.basename(most_recent_file)\n",
        "\n",
        "    return most_recent_file\n",
        "\n",
        "def parse_args(config):\n",
        "    args = []\n",
        "\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args.append(f\"{v}\")\n",
        "        elif isinstance(v, str) and v is not None:\n",
        "            args.append(f'--{k}={v}')\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args.append(f\"--{k}\")\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "\n",
        "    return args\n",
        "\n",
        "def aria2_download(dir, filename, url):\n",
        "    # hf_token    = \"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\" if not HUGGINGFACE_TOKEN else HUGGINGFACE_TOKEN\n",
        "    user_header = f\"Authorization: Bearer {HUGGINGFACE_TOKEN}\"\n",
        "\n",
        "    aria2_config = {\n",
        "        \"console-log-level\"         : \"error\",\n",
        "        \"summary-interval\"          : 10,\n",
        "        \"header\"                    : user_header if \"huggingface.co\" in url else None,\n",
        "        \"continue\"                  : True,\n",
        "        \"max-connection-per-server\" : 16,\n",
        "        \"min-split-size\"            : \"1M\",\n",
        "        \"split\"                     : 16,\n",
        "        \"dir\"                       : dir,\n",
        "        \"out\"                       : filename,\n",
        "        \"_url\"                      : url,\n",
        "    }\n",
        "    aria2_args = parse_args(aria2_config)\n",
        "    subprocess.run([\"aria2c\", *aria2_args])\n",
        "\n",
        "def gdown_download(url, dst, filepath):\n",
        "    if \"/uc?id/\" in url:\n",
        "        return gdown.download(url, filepath, quiet=False)\n",
        "    elif \"/file/d/\" in url:\n",
        "        return gdown.download(url=url, output=filepath, quiet=False, fuzzy=True)\n",
        "    elif \"/drive/folders/\" in url:\n",
        "        os.chdir(dst)\n",
        "        return gdown.download_folder(url, quiet=True, use_cookies=False)\n",
        "\n",
        "def download(url, dst):\n",
        "    filename = get_filename(url, quiet=False)\n",
        "    filepath = os.path.join(dst, filename)\n",
        "\n",
        "    if \"drive.google.com\" in url:\n",
        "        gdown = gdown_download(url, dst, filepath)\n",
        "    elif url.startswith(\"/content/drive/MyDrive/\"):\n",
        "        # Path(filepath).write_bytes(Path(url).read_bytes())\n",
        "        return url\n",
        "    else:\n",
        "        if \"huggingface.co\" in url:\n",
        "            if \"/blob/\" in url:\n",
        "                url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        aria2_download(dst, filename, url)\n",
        "\n",
        "def get_filepath(url, dst):\n",
        "    extensions = get_supported_extensions()\n",
        "    filename = get_filename(url)\n",
        "\n",
        "    if not filename.endswith(extensions):\n",
        "        most_recent_file = get_most_recent_file(dst)\n",
        "        filename = os.path.basename(most_recent_file)\n",
        "\n",
        "    filepath = os.path.join(dst, filename)\n",
        "\n",
        "    return filepath\n",
        "\n",
        "def main():\n",
        "    global model_path, vae_path, lora_path\n",
        "\n",
        "    model_path = vae_path = lora_path = None\n",
        "\n",
        "    download_targets = {\n",
        "        \"model\" : (SDXL_MODEL_URL, models_dir),\n",
        "        \"vae\"   : (SDXL_VAE_URL, vae_dir),\n",
        "        \"lora\"   : (SDXL_LORA_URL, lora_dir),\n",
        "    }\n",
        "    selected_files = {}\n",
        "\n",
        "    for target, (url, dst) in download_targets.items():\n",
        "        if url and f\"PASTE {target.upper()} URL OR GDRIVE PATH HERE\" not in url:\n",
        "            downloader = download(url, dst)\n",
        "            selected_files[target] = get_filepath(url, dst)\n",
        "\n",
        "            if target == \"model\":\n",
        "                model_path = selected_files[\"model\"] if not downloader else downloader\n",
        "            elif target == \"vae\":\n",
        "                vae_path = selected_files[\"vae\"] if not downloader else downloader\n",
        "            elif target == \"lora\":\n",
        "                lora_path = selected_files[\"lora\"] if not downloader else downloader\n",
        "\n",
        "    for category, path in {\n",
        "        \"model\": model_path,\n",
        "        \"vae\": vae_path,\n",
        "        \"lora\": lora_path,\n",
        "    }.items():\n",
        "        if path is not None and os.path.exists(path):\n",
        "            print(f\"Selected {category}: {path}\")\n",
        "\n",
        "main()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iYALQ05ONbei"
      },
      "outputs": [],
      "source": [
        "# @title ## **3. Run to update list. Then select model**\n",
        "import ipywidgets as widgets\n",
        "\n",
        "model_list = os.listdir(models_dir)\n",
        "selected_model = widgets.Dropdown(options=model_list)\n",
        "\n",
        "print(\"Select Model: \")\n",
        "selected_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Z-nzEoOuNcmh"
      },
      "outputs": [],
      "source": [
        "# @title ## **4. Run to update list. Then select LoRA**\n",
        "import ipywidgets as widgets\n",
        "\n",
        "lora_list = os.listdir(lora_dir)\n",
        "lora_list.append(\"\")\n",
        "selected_lora = widgets.Dropdown(options=lora_list)\n",
        "\n",
        "print(\"Select LoRA: \")\n",
        "selected_lora\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kHlv-yzsJY8l"
      },
      "outputs": [],
      "source": [
        "#@title ## **5. Inference**\n",
        "import math\n",
        "from PIL import Image, ImageOps\n",
        "import os\n",
        "from IPython.display import display\n",
        "%store -r\n",
        "\n",
        "# @markdown ### Model Config\n",
        "if \"selected_model\" in globals():\n",
        "    ckpt_path = os.path.join(models_dir, selected_model.value)\n",
        "else:\n",
        "    ckpt_path = model_path\n",
        "\n",
        "if \"selected_lora\" in globals():\n",
        "    network_weights = os.path.join(lora_dir, selected_lora.value)\n",
        "else:\n",
        "    network_weights = \"\"\n",
        "\n",
        "network_module = \"networks.lora\"  # @param [\"networks.lora\", \"networks.dylora\"]\n",
        "network_mul = 1 #@param {type:\"slider\", min:-1, max:2, step:0.1}\n",
        "interactive_mode = True # @param {type:\"boolean\"}\n",
        "# @markdown ### Prompt Config\n",
        "prompt = \"masterpiece, best quality, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, watercolor, night, turtleneck\" #@param {type:'string'}\n",
        "negative_prompt = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry\" #@param {type:'string'}\n",
        "output_path = \"/content/tmp/\" #@param {type:'string'}\n",
        "resolution = \"1024, 1024\" # @param {type: \"string\"}\n",
        "optimization = \"xformers\" # @param [\"memory efficient attention\", \"xformers\", \"scaled dot-product attention\"]\n",
        "conditional_resolution = \"1024, 1024\" # @param {type: \"string\"}\n",
        "crop_top_and_left = \"0, 0\" # @param {type: \"string\"}\n",
        "steps = 28 # @param {type: \"number\"}\n",
        "sampler = \"euler_a\"  # @param [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]\n",
        "scale = 7 # @param {type: \"number\"}\n",
        "seed = -1 # @param {type: \"number\"}\n",
        "images_per_prompt = 1 # @param {type: \"number\"}\n",
        "batch_size = 1 # @param {type: \"number\"}\n",
        "clip_skip = 2 # @param {type: \"number\"}\n",
        "\n",
        "# if interactive_mode:\n",
        "\n",
        "os.makedirs(output_path, exist_ok=True)\n",
        "\n",
        "separators = [\"*\", \"x\", \",\"]\n",
        "\n",
        "for separator in separators:\n",
        "    if separator in resolution:\n",
        "        width, height = [value.strip() for value in resolution.split(separator)]\n",
        "        break\n",
        "\n",
        "for separator in separators:\n",
        "    if separator in conditional_resolution:\n",
        "        original_width, original_height = [value.strip() for value in conditional_resolution.split(separator)]\n",
        "        break\n",
        "\n",
        "for separator in separators:\n",
        "    if separator in crop_top_and_left:\n",
        "        crop_top, crop_left = [value.strip() for value in crop_top_and_left.split(separator)]\n",
        "        break\n",
        "\n",
        "network_config = {\n",
        "    \"network_module\": network_module,\n",
        "    \"network_weights\": network_weights,\n",
        "    \"network_show_meta\": True,\n",
        "}\n",
        "\n",
        "config = {\n",
        "    \"interactive\": interactive_mode,\n",
        "    \"no_preview\": interactive_mode,\n",
        "    \"prompt\": prompt + \" --n \" + negative_prompt if not interactive_mode else None,\n",
        "    \"images_per_prompt\": images_per_prompt,\n",
        "    \"outdir\": output_path,\n",
        "    \"W\": width,\n",
        "    \"H\": height,\n",
        "    \"original_width\": original_width,\n",
        "    \"original_height\": original_height,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"vae_batch_size\": 1,\n",
        "    \"no_half_vae\": True,\n",
        "    \"steps\": steps,\n",
        "    \"sampler\": sampler,\n",
        "    \"scale\": scale,\n",
        "    \"ckpt\": ckpt_path,\n",
        "    \"vae\": vae_path,\n",
        "    \"seed\": seed if seed > 0 else None,\n",
        "    \"fp16\": True,\n",
        "    \"sdpa\": True if optimization == \"scaled dot-product attention\" else False,\n",
        "    \"xformers\": True if optimization == \"xformers\" else False,\n",
        "    \"mem_eff_attn\": True if optimization == \"mem_eff_attn\" else False,\n",
        "    \"opt_channels_last\": True,\n",
        "    \"clip_skip\": clip_skip,\n",
        "    \"max_embeddings_multiples\": 3,\n",
        "}\n",
        "\n",
        "interactive_config = {\n",
        "    \"n\" : negative_prompt,\n",
        "    \"w\" : width,\n",
        "    \"h\" : height,\n",
        "    \"ow\" : original_width,\n",
        "    \"oh\" : original_height,\n",
        "    \"ct\" : crop_top,\n",
        "    \"cl\" : crop_left,\n",
        "    \"s\" : steps,\n",
        "    \"d\" : seed,\n",
        "    \"l\" : scale,\n",
        "    \"am\" : network_mul if network_weights.endswith(\".safetensors\") else None,\n",
        "}\n",
        "\n",
        "def display_results(count):\n",
        "    samples = os.listdir(output_path)\n",
        "    samples.sort(reverse=True)\n",
        "    samples = samples[:count]\n",
        "\n",
        "    for sample in samples:\n",
        "        if sample.endswith((\".png\", \".jpg\")):\n",
        "            image_path = os.path.join(output_path, sample)\n",
        "\n",
        "            if os.path.exists(image_path):\n",
        "                img = Image.open(image_path)\n",
        "                img = img.resize((512, 512))  # Resize the image to 512x512 pixels\n",
        "                display(img)\n",
        "\n",
        "\n",
        "def parse_args(config, to_logs=False):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if to_logs:\n",
        "            if k.startswith(\"_\"):\n",
        "                args += f'{v} '\n",
        "            elif isinstance(v, str):\n",
        "                args += f'--{k} {v} '\n",
        "            elif isinstance(v, bool) and v:\n",
        "                args += f\"--{k} \"\n",
        "            elif isinstance(v, float) and not isinstance(v, bool):\n",
        "                args += f\"--{k} {v} \"\n",
        "            elif isinstance(v, int) and not isinstance(v, bool):\n",
        "                args += f\"--{k} {v} \"\n",
        "        else:\n",
        "            if k.startswith(\"_\"):\n",
        "                args += f'\"{v}\" '\n",
        "            elif isinstance(v, str):\n",
        "                args += f'--{k}=\"{v}\" '\n",
        "            elif isinstance(v, bool) and v:\n",
        "                args += f\"--{k} \"\n",
        "            elif isinstance(v, float) and not isinstance(v, bool):\n",
        "                args += f\"--{k}={v} \"\n",
        "            elif isinstance(v, int) and not isinstance(v, bool):\n",
        "                args += f\"--{k}={v} \"\n",
        "\n",
        "    return args.strip()\n",
        "\n",
        "if network_weights.endswith(\".safetensors\"):\n",
        "    config.update(network_config)\n",
        "\n",
        "gen_args = parse_args(config)\n",
        "interactive_args = parse_args(interactive_config, to_logs=True)\n",
        "\n",
        "final_args = f\"python sdxl_gen_img.py {gen_args}\"\n",
        "\n",
        "if interactive_mode:\n",
        "    print(\"Append this parameter to your prompt in interactive mode:\")\n",
        "    display(f\"{prompt} {interactive_args}\")\n",
        "    print(\"\\n\")\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}\n",
        "\n",
        "display_results(batch_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8APXceQ5SIkC"
      },
      "outputs": [],
      "source": [
        "# @title ## **6. Display your result**\n",
        "import math\n",
        "from PIL import Image, ImageOps\n",
        "\n",
        "How_many_previous_results_to_show = 2 #@param {type:\"integer\"}\n",
        "# @markdown _The preview are resized to 512x512_\n",
        "resize_images = True #@param {type:'boolean'}\n",
        "\n",
        "def display_results(count, resize_images=True):\n",
        "    samples = os.listdir(output_path)\n",
        "    samples.sort(reverse=True)\n",
        "    samples = samples[:count]\n",
        "\n",
        "    if len(samples) == 1 or not resize_images and len(samples) > 1:\n",
        "        sample = samples[0]\n",
        "        if sample.endswith((\".png\", \".jpg\")):\n",
        "            image_path = os.path.join(output_path, sample)\n",
        "\n",
        "            if os.path.exists(image_path):\n",
        "                img = Image.open(image_path)\n",
        "                if resize_images:\n",
        "                    img = img.resize((512, 512))\n",
        "                display(img)\n",
        "    else:\n",
        "        num_cols = math.ceil(math.sqrt(count))\n",
        "        num_rows = math.ceil(count / num_cols)\n",
        "        collage_width = num_cols * 512\n",
        "        collage_height = num_rows * 512\n",
        "        collage = Image.new('RGB', (collage_width, collage_height))\n",
        "\n",
        "        for i, sample in enumerate(samples):\n",
        "            if sample.endswith((\".png\", \".jpg\")):\n",
        "                image_path = os.path.join(output_path, sample)\n",
        "\n",
        "                if os.path.exists(image_path):\n",
        "                    img = Image.open(image_path)\n",
        "                    img = img.resize((512, 512))\n",
        "\n",
        "                    row = i // num_cols\n",
        "                    col = i % num_cols\n",
        "\n",
        "                    x = col * 512\n",
        "                    y = row * 512\n",
        "\n",
        "                    collage.paste(img, (x, y))\n",
        "\n",
        "        display(collage)\n",
        "\n",
        "display_results(How_many_previous_results_to_show, resize_images=resize_images)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "authorship_tag": "ABX9TyP/Dc0VuuVY4xCBV+NSYuoj",
      "gpuType": "T4",
      "include_colab_link": true,
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}