stablediffusion_infinity.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "55fb968d-4565-4e09-b81b-684386047fb5",
   "metadata": {},
   "source": [
    "# stablediffusion-infinity\n",
    "\n",
    "https://github.com/lkwq007/stablediffusion-infinity\n",
    "\n",
    "Outpainting with Stable Diffusion on an infinite canvas"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bb49640-b3dd-424b-909a-83efc7125b83",
   "metadata": {},
   "source": [
    "## import libs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88a92062-8b51-4de0-853b-2747ef223c5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp -nrf PyPatchMatch/csrc .\n",
    "!cp -nf PyPatchMatch/Makefile .\n",
    "!cp -nf PyPatchMatch/travis.sh .\n",
    "!cp -nf PyPatchMatch/patch_match.py ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8eeadd62-c339-4c5c-9107-df0b4331b9d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "from ipycanvas import Canvas, hold_canvas, MultiCanvas\n",
    "import ipywidgets as widgets\n",
    "\n",
    "from utils import *\n",
    "from canvas import InfCanvas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3da5c0a-3b45-4ac6-91cb-1811f867fda9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import autocast\n",
    "from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16ddabe9-6652-4994-8a2e-497ce6f7dbd4",
   "metadata": {},
   "source": [
    "## setup StableDiffusionInpaintPipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c799c58-4f49-41eb-b5c9-66e45ef32964",
   "metadata": {},
   "outputs": [],
   "source": [
    "text2img = StableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"fp16\", torch_dtype=torch.float16, use_auth_token=True).to(\"cuda\")\n",
    "inpaint=StableDiffusionInpaintPipeline(\n",
    "            vae=text2img.vae,\n",
    "            text_encoder=text2img.text_encoder,\n",
    "            tokenizer=text2img.tokenizer,\n",
    "            unet=text2img.unet,\n",
    "            scheduler=text2img.scheduler,\n",
    "            safety_checker=text2img.safety_checker,\n",
    "            feature_extractor=text2img.feature_extractor\n",
    ").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "584319f4-d662-4732-94b2-ce7636496660",
   "metadata": {},
   "source": [
    "## setup UI control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6a47549-1bfc-4378-bf7f-1fbc31d61bd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "strength_slider = widgets.FloatSlider(\n",
    "    value=0.75,\n",
    "    min=0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"Strength:\",\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation=\"horizontal\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\",\n",
    "    # layout=widgets.Layout(width='100px')\n",
    ")\n",
    "\n",
    "step_input = widgets.BoundedIntText(\n",
    "    value=50,\n",
    "    min=1,\n",
    "    max=1000,\n",
    "    step=1,\n",
    "    description=\"Steps:\",\n",
    "    disabled=False,\n",
    "    layout=widgets.Layout(width=\"180px\"),\n",
    ")\n",
    "guidance_input = widgets.FloatText(\n",
    "    value=7.5,\n",
    "    step=0.1,\n",
    "    description=\"Guidance:\",\n",
    "    disabled=False,\n",
    "    layout=widgets.Layout(width=\"180px\"),\n",
    ")\n",
    "resize_check = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description=\"Resize SD input to 512x512\",\n",
    "    disabled=False,\n",
    "    indent=False,\n",
    "    layout=widgets.Layout(width=\"180px\"),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d45ade9b-50df-43c0-8e26-d081567d38a3",
   "metadata": {},
   "source": [
    "## setup InfCanvas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6905335-6ffb-4650-9596-58a07c8bca1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "base = InfCanvas(800, 600, selection_size=256)\n",
    "base.setup_widgets()\n",
    "base.setup_mouse()\n",
    "\n",
    "\n",
    "def run_outpaint(btn):\n",
    "    with base.output:\n",
    "        base.output.clear_output()\n",
    "        base.read_selection_from_buffer()\n",
    "        img = base.sel_buffer[:, :, 0:3]\n",
    "        mask = base.sel_buffer[:, :, -1]\n",
    "        process_size = 512 if resize_check.value else base.selection_size\n",
    "        if mask.sum() > 0:\n",
    "            img, mask = functbl[base.fill_button.value](img, mask)\n",
    "            init_image = Image.fromarray(img)\n",
    "            mask = 255 - mask\n",
    "            # mask=skimage.measure.block_reduce(mask,(8,8),np.max)\n",
    "            # mask=mask.repeat(8, axis=0).repeat(8, axis=1)\n",
    "            mask_image = Image.fromarray(mask)\n",
    "            # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))\n",
    "            with autocast(\"cuda\"):\n",
    "                images = inpaint(\n",
    "                    prompt=base.text_input.value,\n",
    "                    init_image=init_image.resize(\n",
    "                        (process_size, process_size), resample=Image.Resampling.LANCZOS\n",
    "                    ),\n",
    "                    mask_image=mask_image.resize((process_size, process_size)),\n",
    "                    strength=strength_slider.value,\n",
    "                    num_inference_steps=step_input.value,\n",
    "                    guidance_scale=guidance_input.value,\n",
    "                )[\"sample\"]\n",
    "        else:\n",
    "            with autocast(\"cuda\"):\n",
    "                images = text2img(\n",
    "                    prompt=base.text_input.value,\n",
    "                    height=process_size,\n",
    "                    width=process_size,\n",
    "                )[\"sample\"]\n",
    "        out = base.sel_buffer.copy()\n",
    "        out[:, :, 0:3] = np.array(\n",
    "            images[0].resize(\n",
    "                (base.selection_size, base.selection_size),\n",
    "                resample=Image.Resampling.LANCZOS,\n",
    "            )\n",
    "        )\n",
    "        out[:, :, -1] = 255\n",
    "        base.fill_selection(out)\n",
    "        with hold_canvas():\n",
    "            base.draw_selection_box()\n",
    "\n",
    "\n",
    "def export_button_clicked(btn):\n",
    "    with base.output:\n",
    "        base.output.clear_output()\n",
    "        img = base.export()\n",
    "        pil = Image.fromarray(img)\n",
    "        time_str = time.strftime(\"%Y%m%d_%H%M%S\")\n",
    "        pil.save(f\"outpaint_{time_str}.png\")\n",
    "        print(f\"Canvas saved to outpaint_{time_str}.png\")\n",
    "        display(pil)\n",
    "\n",
    "\n",
    "base.run_button.on_click(run_outpaint)\n",
    "base.export_button.on_click(export_button_clicked)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fffd7a88-3bb1-4484-8d2a-94f88653091a",
   "metadata": {
    "tags": []
   },
   "source": [
    "## upload an image?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3491395-5554-4e96-a71f-b093b9989c0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "from PIL import ImageOps\n",
    "uploader_label = widgets.Label(\"[Optional] Upload an image? (will be resized to fit into canvas)\")\n",
    "uploader = widgets.FileUpload(\n",
    "    accept=\"image/*\",\n",
    "    multiple=False\n",
    ")\n",
    "uploader_output = widgets.Output()\n",
    "upload_button = widgets.Button(description='Confirm')\n",
    "def start_func(btn):\n",
    "    with uploader_output:\n",
    "        uploader_output.clear_output()\n",
    "        if len(uploader.value)>0:\n",
    "            keys=list(uploader.value.keys())\n",
    "            val=uploader.value[keys[-1]][\"content\"]\n",
    "            pil=Image.open(io.BytesIO(val))\n",
    "            print(f\"Will use {keys[-1]} as the base image for outpainting\")\n",
    "            w,h=pil.size\n",
    "            if w>base.width-100 or h>base.height-100:\n",
    "                pil=ImageOps.contain(pil, (base.width-100,base.height-100))\n",
    "            base.buffer_dirty=True\n",
    "            w,h=pil.size\n",
    "            arr=np.array(pil)\n",
    "            yo=(base.height-h)//2\n",
    "            xo=(base.width-w)//2\n",
    "            base.buffer[yo:yo+h,xo:xo+w,0:3]=arr[:,:,0:3]\n",
    "            base.buffer[yo:yo+h,xo:xo+w,-1]=255\n",
    "        else:\n",
    "            print(\"No image uploaded\")\n",
    "\n",
    "upload_button.on_click(start_func)\n",
    "lst=[uploader_label,widgets.HBox([uploader,upload_button]),uploader_output]\n",
    "for item in lst:\n",
    "    display(item)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81fac64a-2472-49c9-8aed-8d653ec3fd45",
   "metadata": {},
   "source": [
    "## have fun here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f75607c-0704-4a0a-a5c1-edffdc656b12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# disable upload button\n",
    "upload_button.disabled=True\n",
    "display_lst=base.display()\n",
    "display_lst.insert(-1,widgets.HBox([resize_check,step_input,guidance_input,strength_slider]))\n",
    "for item in display_lst:\n",
    "    display(item)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ed2249-5c6e-41a5-a63c-387281734162",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35b6c670-33f8-4a63-ae13-2d3c1c8d6807",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}