playground.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HyperTile Playground"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I'll presume you have some familiarity with the `diffusers` library.\n",
    "\n",
    "I encourage you to experiment with both the `text2img` and `img2img` variations. Keep in mind that due to the absence of an HD-LoRA model for `diffusers`, the `text2img` results may exhibit suboptimal structures. However, it's essential to recognize that these limitations stem from the initial training of SD at 512x512 resolution, and improvements are anticipated in the future.\n",
    "\n",
    "If you find yourself needing further information, don't hesitate to consult the comprehensive documentation provided by `diffusers`; it offers valuable insights and guidance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the packages we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from tqdm.auto import trange\n",
    "import logging\n",
    "\n",
    "import torch\n",
    "\n",
    "from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, StableDiffusionXLPipeline\n",
    "from diffusers.schedulers import UniPCMultistepScheduler\n",
    "\n",
    "from diffusers.utils import load_image\n",
    "\n",
    "from hyper_tile import split_attention, flush\n",
    "\n",
    "# To log attention-splitting\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select the path to the model you want (*safetensors*)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = torch.float16 # bfloat16 can result of out-of-memory with big images due to interpolation limits, well-document in diffusers library\n",
    "device = torch.device('cuda')\n",
    "\n",
    "model_path = r\"path-to-model.safetensors\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text-to-Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe: StableDiffusionPipeline = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=dtype, local_files_only=True, use_safetensors=True, load_safety_checker=False) # type: ignore\n",
    "pipe.to(device)\n",
    "pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Choose your desired `height` and `width`.\n",
    "\n",
    "2. You have the flexibility to adjust the `tile_size` independently for the VAE and UNet components. For the VAE, a `tile_size` of 128 is optimal without sacrificing performance. However, for the UNet, it's advisable to use a chunk size of 256 or greater. `swap_size` determine how many different tiles per dimension are used, to avoid overlap seams in some cases.\n",
    "\n",
    "3. Modify the `disable` parameter to either True or False to observe the results with or without HyperTile.\n",
    "\n",
    "**Note**: For improved chunk division, consider using dimensions that are multiples of 128. This can enhance the effectiveness of the chunking process. (This is enforced)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try lower value if you dont have 16 Gb VRAM\n",
    "\n",
    "height, width = 2688, 1792\n",
    "\n",
    "height = int(height)//128*128 # enforcing multiples of 128\n",
    "width = int(width)//128*128\n",
    "print(height, width)\n",
    "\n",
    "with split_attention(pipe.vae, height, width, tile_size=128):\n",
    "    # ! Change the tile_size and disable to see their effects\n",
    "    with split_attention(pipe.unet, height, width, tile_size=128, swap_size=2, disable=False):\n",
    "        flush()\n",
    "        img = pipe(\n",
    "            # ! Change the prompt and other parameters\n",
    "            prompt='forest, path, stone, red trees, detailed, buildings', \n",
    "            negative_prompt='blurry, low detail',\n",
    "            num_inference_steps=26, guidance_scale=7.5, \n",
    "            height=height, width=width,\n",
    "        ).images[0]\n",
    "img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image-to-Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe: StableDiffusionImg2ImgPipeline  = StableDiffusionImg2ImgPipeline.from_single_file(model_path, torch_dtype=dtype, local_files_only=True, use_safetensors=True, load_safety_checker=False) # type: ignore\n",
    "pipe.to(device)\n",
    "pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the image that you want."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = load_image(\"image.png\")\n",
    "ar = image.height / image.width"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Define the target `height` and the number of loopbacks required, indicating how many times we perform image-to-image operations. This parameter is essential as we employ simple Lanczos upscaling.\n",
    "\n",
    "2. Adjust the `strength` and `loopback` settings to achieve the optimal outcome. You can experiment with lower strength values paired with more loopbacks or larger strength values with fewer loopbacks.\n",
    "\n",
    "3. Customize the `tile_size` separately for the VAE and UNet components. A `tile_size` of 128 is recommended for the VAE without compromising quality. For the UNet, it's advisable to use a `tile_size` size of 256 or greater. `swap_size` determine how many different tiles per dimension are used, to avoid overlap seams in some cases.\n",
    "\n",
    "4. Toggle the `disable` parameter between True and False to observe the results with or without the use of HyperTile.\n",
    "\n",
    "**Note**: For improved chunk division, consider using dimensions that are multiples of 128. This practice enhances the efficiency of the chunking process.\n",
    "\n",
    "**Note**: The inclusion of loopbacks is essential due to the original training of Stable-Diffusion (SD) on 512x512 images. When we upscale these images 3-4 times or more, the use of Lanczos upscaling introduces blurriness. Loopbacks play a crucial role in mitigating this issue, effectively restoring image clarity and preserving details during the upscaling process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loopback = 2 # Use 1 or 2, depending on how much upscaling we are doing\n",
    "# Try lower value if you dont have 16 Gb VRAM\n",
    "height = 512*3\n",
    "\n",
    "width = height/ar\n",
    "height = int(height)//128*128 # enforcing multiples of 128\n",
    "width = int(width)//128*128\n",
    "print(height, width)\n",
    "\n",
    "# Upscale to the correct resolution\n",
    "img = image.resize((width, height), resample=Image.LANCZOS) if image.size != (width, height) else image\n",
    "\n",
    "with split_attention(pipe.vae, height, width, tile_size=128):\n",
    "    # ! Change the chunk and disable to see their effects\n",
    "    with split_attention(pipe.unet, height, width, tile_size=256, swap_size=2, disable=False):\n",
    "        flush()\n",
    "        for i in trange(loopback):\n",
    "            img = pipe(\n",
    "                prompt='forest, path, stone, red trees, detailed', \n",
    "                negative_prompt='blurry, low detail',\n",
    "                num_inference_steps=28, guidance_scale=7.5, \n",
    "                image=img, strength=0.46, # ! you can also change the strength\n",
    "            ).images[0]\n",
    "img"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}