txt2image_widgets.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c9a6d15c-509e-43c6-bbbf-f9515d5e3684",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e380c236a83e43c68f0d4fed0b9e7b96",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "StableDiffusionPipeline {\n",
       "  \"_class_name\": \"StableDiffusionPipeline\",\n",
       "  \"_diffusers_version\": \"0.20.0.dev0\",\n",
       "  \"_name_or_path\": \"runwayml/stable-diffusion-v1-5\",\n",
       "  \"feature_extractor\": [\n",
       "    null,\n",
       "    null\n",
       "  ],\n",
       "  \"requires_safety_checker\": false,\n",
       "  \"safety_checker\": [\n",
       "    null,\n",
       "    null\n",
       "  ],\n",
       "  \"scheduler\": [\n",
       "    \"diffusers\",\n",
       "    \"PNDMScheduler\"\n",
       "  ],\n",
       "  \"text_encoder\": [\n",
       "    \"transformers\",\n",
       "    \"CLIPTextModel\"\n",
       "  ],\n",
       "  \"tokenizer\": [\n",
       "    \"transformers\",\n",
       "    \"CLIPTokenizer\"\n",
       "  ],\n",
       "  \"unet\": [\n",
       "    \"diffusers\",\n",
       "    \"UNet2DConditionModel\"\n",
       "  ],\n",
       "  \"vae\": [\n",
       "    \"diffusers\",\n",
       "    \"AutoencoderKL\"\n",
       "  ]\n",
       "}"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from diffusers import StableDiffusionPipeline\n",
    "import torch\n",
    "from torchvision import models\n",
    "from torchvision import transforms\n",
    "#from decimal import Decimal\n",
    "\n",
    "device = \"cuda\"\n",
    "\n",
    "# load model\n",
    "model_path = \"./trained_weights/als_10_rings/\" #pytorch_lora_weights.bin\"\n",
    "pipe = StableDiffusionPipeline.from_pretrained(\n",
    "    \"runwayml/stable-diffusion-v1-5\",\n",
    "    torch_dtype=torch.float16,\n",
    "    safety_checker=None,\n",
    "    feature_extractor=None,\n",
    "    requires_safety_checker=False\n",
    ")\n",
    "\n",
    "# load lora weights\n",
    "pipe.unet.load_attn_procs(model_path)\n",
    "# set to use GPU for inference\n",
    "pipe.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "72ca53b1-1e56-4bf5-b6b7-26e20c87ab91",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(prompt, out_name='im'):\n",
    "    out_dir = '40k_generated/test'\n",
    "    if not os.path.exists(f'./{out_dir}'):\n",
    "        os.makedirs(f'./{out_dir}')\n",
    "    image = pipe(prompt, num_inference_steps=30).images[0]\n",
    "    image.save(f\"./{out_dir}/{out_name}.jpg\")\n",
    "    # print(outputs.images)\n",
    "    # for idx, image in enumerate(outputs.images):\n",
    "    #     image.save(f\"./{out_dir}/{out_name}_{idx}.jpg\")\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "abe03ca9-b627-4ada-bb95-9e9f238e796a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def classifier(model, image, threshold=0.5) -> bool:\n",
    "    model = model.to(device)\n",
    "    image = image.to(device).unsqueeze(0)\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        output = model(image)\n",
    "    prob = torch.nn.Softmax(dim=1)(output)\n",
    "    #print(prob[:,1] > threshold)\n",
    "    return prob[:,1] > threshold, prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bdb453d5-b0fc-4e05-b11e-3d50eb1eb355",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aec138c9946e4bcebabf5c6f6dd4e770",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Text(value='', description='Prompt:', placeholder='Enter a prompt'), Dropdown(description='Clas…"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display, Image, clear_output\n",
    "from PIL import Image as PIL_Image\n",
    "import math\n",
    "\n",
    "text_input = widgets.Text(\n",
    "    value='',\n",
    "    placeholder='Enter a prompt',\n",
    "    description='Prompt:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "# threshold_input = widgets.Text(\n",
    "#     value='',\n",
    "#     placeholder='Enter a probability [0-1], defualt 0.5',\n",
    "#     description='Threshold:',\n",
    "#     disabled=False\n",
    "# )\n",
    "\n",
    "threshold_input = widgets.FloatSlider(\n",
    "    value=0.5,           # Initial value\n",
    "    min=0,               # Minimum value\n",
    "    max=1,               # Maximum value\n",
    "    step=0.01,           # Step size\n",
    "    description='Threshold:',\n",
    "    orientation='horizontal'  # Slider orientation (horizontal or vertical)\n",
    ")\n",
    "\n",
    "\n",
    "batch_size = widgets.IntSlider(\n",
    "    value=1,           # Initial value\n",
    "    min=1,               # Minimum value\n",
    "    max=48,               # Maximum value\n",
    "    step=1,           # Step size\n",
    "    description='Batchsize:',\n",
    "    orientation='horizontal'  # Slider orientation (horizontal or vertical)\n",
    ")\n",
    "\n",
    "\n",
    "max_epoch = widgets.IntSlider(\n",
    "    value=5,           # Initial value\n",
    "    min=1,               # Minimum value\n",
    "    max=100,               # Maximum value\n",
    "    step=1,           # Step size\n",
    "    description='N_attempts:',\n",
    "    orientation='horizontal'  # Slider orientation (horizontal or vertical)\n",
    ")\n",
    "\n",
    "display_button = widgets.Button(\n",
    "    description=\"Generate Image\",\n",
    "    button_style=\"info\"\n",
    ")\n",
    "\n",
    "\n",
    "pretrained_classfiers = {\n",
    "    'ResNet50': 'resnet_100epochs.pth',\n",
    "    'ViT-16x16': 'vit_p16_50epochs.pth',\n",
    "    'ViT-32x32': 'vit_p32_50epochs.pth',\n",
    "}\n",
    "\n",
    "dropdown = widgets.Dropdown(\n",
    "    options=pretrained_classfiers,\n",
    "    description='Classifier:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "\n",
    "# Create an output widget to display the image\n",
    "output_image = widgets.Output()\n",
    "filename='/lovelace/zhuowen/diffusers/als/40k_generated/test/im.jpg'\n",
    "\n",
    "# Define a function to display the image when the button is clicked\n",
    "def display_image(button):\n",
    "    flag = False\n",
    "    epoch = 0\n",
    "    prob = 0\n",
    "    \n",
    "    while not flag and epoch < max_epoch.value:\n",
    "        with output_image:\n",
    "            clear_output()  # Clear any previous output\n",
    "            image_prompt = text_input.value.strip()\n",
    "            #threshold = threshold_input.value.strip()\n",
    "            cls_dir = './discriminator_data/checkpoints/'\n",
    "            cls_model = torch.load(cls_dir+dropdown.value)\n",
    "            \n",
    "            epoch += 1\n",
    "            if image_prompt:\n",
    "                print(f\"Generating {batch_size.value} images, {epoch}th attempt.\")\n",
    "                print_prob = prob[0][1].to('cpu').item() if isinstance(prob, torch.Tensor) else 0 \n",
    "                print(f'Previous probability: {round(print_prob, 2)}')\n",
    "                generate(image_prompt)\n",
    "                try:\n",
    "                    image = Image(filename=filename)\n",
    "                except Exception as e:\n",
    "                    print(f\"Error: {e}\")\n",
    "    \n",
    "                transform=transforms.Compose([transforms.Resize(224),\n",
    "                                              transforms.CenterCrop(224),\n",
    "                                              transforms.ToTensor(),\n",
    "                                              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "                                             ])\n",
    "                image_tensor = transform(PIL_Image.open(filename))\n",
    "                b, prob= classifier(cls_model, image_tensor, threshold=threshold_input.value)\n",
    "                flag = b\n",
    "                print_prob = prob[0][1].to('cpu').item() if isinstance(prob, torch.Tensor) else 0\n",
    "                print(f'Current probability: {round(print_prob, 2)}')\n",
    "\n",
    "                if flag:\n",
    "                    display(image)\n",
    "\n",
    "\n",
    "\n",
    "# Attach the function to the button's click event\n",
    "display_button.on_click(display_image)\n",
    "\n",
    "# Display the widgets\n",
    "widgets.VBox([text_input, dropdown, threshold_input, batch_size, max_epoch, display_button, output_image])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2fbcb69-67bc-4a9d-bdd9-71696146556c",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}