progressive-distillation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from diffusers import AutoencoderKL, UNet2DModel, DDIMPipeline, DDIMScheduler, DDPMPipeline, DDPMScheduler\n",
    "from progressive_distillation import DistillationPipelineOld\n",
    "from diffusers.optimization import get_scheduler\n",
    "from diffusers.training_utils import EMAModel\n",
    "import math\n",
    "import requests\n",
    "from torchvision.transforms import (\n",
    "    CenterCrop,\n",
    "    Compose,\n",
    "    InterpolationMode,\n",
    "    Normalize,\n",
    "    RandomHorizontalFlip,\n",
    "    Resize,\n",
    "    ToTensor,\n",
    "    ToPILImage\n",
    ")\n",
    "from torch.utils.data import Dataset\n",
    "from accelerate import Accelerator\n",
    "from progressive_distillation import utils\n",
    "from tqdm import tqdm\n",
    "import torch.nn.functional as F\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['CUDA_LAUNCH_BLOCKING']=\"1\"\n",
    "os.environ['TORCH_USE_CUDA_DSA']=\"1\"\n",
    "\n",
    "# # Enable/Disable CUDA\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.cuda.is_available = lambda : True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(0)\n",
    "training_config = utils.DiffusionTrainingArgs()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "UnidentifiedImageError",
     "evalue": "cannot identify image file <_io.BytesIO object at 0x7f329fa17650>",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mUnidentifiedImageError\u001b[0m                    Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [4], line 4\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Load an image of my dog for this example\u001b[39;00m\n\u001b[1;32m      3\u001b[0m image_url \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhttps://i.imgur.com/IJcs4Aa.jpeg\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 4\u001b[0m image \u001b[38;5;241m=\u001b[39m \u001b[43mImage\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrequests\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimage_url\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstream\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraw\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/work/pi_ahoumansadr_umass_edu/svkulkarni/conda_envs/distill-sd/lib/python3.10/site-packages/PIL/Image.py:3280\u001b[0m, in \u001b[0;36mopen\u001b[0;34m(fp, mode, formats)\u001b[0m\n\u001b[1;32m   3278\u001b[0m     warnings\u001b[38;5;241m.\u001b[39mwarn(message)\n\u001b[1;32m   3279\u001b[0m msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcannot identify image file \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;241m%\u001b[39m (filename \u001b[38;5;28;01mif\u001b[39;00m filename \u001b[38;5;28;01melse\u001b[39;00m fp)\n\u001b[0;32m-> 3280\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m UnidentifiedImageError(msg)\n",
      "\u001b[0;31mUnidentifiedImageError\u001b[0m: cannot identify image file <_io.BytesIO object at 0x7f329fa17650>"
     ]
    }
   ],
   "source": [
    "# Load an image of my dog for this example\n",
    "\n",
    "image_url = \"https://i.imgur.com/IJcs4Aa.jpeg\"\n",
    "image = Image.open(requests.get(image_url, stream=True).raw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the transforms to apply to the image for training\n",
    "augmentations = utils.get_train_transforms(training_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SingleImageDataset(Dataset):\n",
    "    def __init__(self, image, batch_size):\n",
    "        self.image = image\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.batch_size\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_image = augmentations(image.convert(\"RGB\"))\n",
    "train_dataset = SingleImageDataset(train_image, training_config.batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "teacher = UNet2DModel.from_pretrained(\"bglick13/minnie-diffusion\")\n",
    "distiller = DistillationPipelineOld()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1000\n",
    "generator = torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Distill step 0 from 1000 -> 500\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'distiller' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [5], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m distill_step \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m2\u001b[39m):\n\u001b[1;32m      5\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDistill step \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mdistill_step\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m from \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mN\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m -> \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mN \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m \u001b[38;5;241m2\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 6\u001b[0m     teacher, distilled_ema, distill_accelrator \u001b[38;5;241m=\u001b[39m \u001b[43mdistiller\u001b[49m(teacher, N, train_dataset, epochs\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m300\u001b[39m, batch_size\u001b[38;5;241m=\u001b[39mtraining_config\u001b[38;5;241m.\u001b[39mbatch_size)\n\u001b[1;32m      7\u001b[0m     N \u001b[38;5;241m=\u001b[39m N \u001b[38;5;241m/\u001b[39m\u001b[38;5;241m/\u001b[39m \u001b[38;5;241m2\u001b[39m\n\u001b[1;32m      8\u001b[0m     new_scheduler \u001b[38;5;241m=\u001b[39m DDPMScheduler(num_train_timesteps\u001b[38;5;241m=\u001b[39mN, beta_schedule\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msquaredcos_cap_v2\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'distiller' is not defined"
     ]
    }
   ],
   "source": [
    "teacher = UNet2DModel.from_pretrained(\"bglick13/minnie-diffusion\")\n",
    "N = 1000\n",
    "distilled_images = []\n",
    "for distill_step in range(2):\n",
    "    print(f\"Distill step {distill_step} from {N} -> {N // 2}\")\n",
    "    teacher, distilled_ema, distill_accelrator = distiller(teacher, N, train_dataset, epochs=300, batch_size=training_config.batch_size)\n",
    "    N = N // 2\n",
    "    new_scheduler = DDPMScheduler(num_train_timesteps=N, beta_schedule=\"squaredcos_cap_v2\")\n",
    "    pipeline = DDPMPipeline(\n",
    "        unet=distill_accelrator.unwrap_model(distilled_ema.averaged_model if training_config.use_ema else teacher),\n",
    "        scheduler=new_scheduler,\n",
    "    )\n",
    "\n",
    "    # run pipeline in inference (sample random noise and denoise)\n",
    "    images = pipeline(generator=generator, batch_size=training_config.batch_size, output_type=\"numpy\").images\n",
    "\n",
    "    # denormalize the images and save to tensorboard\n",
    "    images_processed = (images * 255).round().astype(\"uint8\")\n",
    "    distilled_images.append(images_processed[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train_image' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [6], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Display train image for reference\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m train_image_display \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_image\u001b[49m \u001b[38;5;241m*\u001b[39m \u001b[38;5;241m0.5\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m0.5\u001b[39m\n\u001b[1;32m      3\u001b[0m train_image_display \u001b[38;5;241m=\u001b[39m ToPILImage()(train_image_display)\n\u001b[1;32m      4\u001b[0m display(train_image_display)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'train_image' is not defined"
     ]
    }
   ],
   "source": [
    "# Display train image for reference\n",
    "train_image_display = train_image * 0.5 + 0.5\n",
    "train_image_display = ToPILImage()(train_image_display)\n",
    "display(train_image_display)\n",
    "\n",
    "for i, image in enumerate(distilled_images):\n",
    "    print(f\"Distilled image {i}\")\n",
    "    display(Image.fromarray(image))\n",
    "    Image.fromarray(image).save(f\"distilled_{i}.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Image.fromarray(images_processed[0]))\n",
    "display(Image.fromarray(images_processed[1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA error: device-side assert triggered\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m/work/pi_ahoumansadr_umass_edu/svkulkarni/Projects/distill-sd/progressive-distillation.ipynb Cell 17\u001b[0m in \u001b[0;36m2\n\u001b[1;32m     <a href='vscode-notebook-cell://ood.unity.rc.umass.edu/work/pi_ahoumansadr_umass_edu/svkulkarni/Projects/distill-sd/progressive-distillation.ipynb#X22sdnNjb2RlLXJlbW90ZQ%3D%3D?line=24'>25</a>\u001b[0m \u001b[39m# Below code will run on gpu, please pass cpu everywhere as the device and set 'dtype' to torch.float32 for cpu inference.\u001b[39;00m\n\u001b[1;32m     <a href='vscode-notebook-cell://ood.unity.rc.umass.edu/work/pi_ahoumansadr_umass_edu/svkulkarni/Projects/distill-sd/progressive-distillation.ipynb#X22sdnNjb2RlLXJlbW90ZQ%3D%3D?line=25'>26</a>\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39minference_mode():\n\u001b[0;32m---> <a href='vscode-notebook-cell://ood.unity.rc.umass.edu/work/pi_ahoumansadr_umass_edu/svkulkarni/Projects/distill-sd/progressive-distillation.ipynb#X22sdnNjb2RlLXJlbW90ZQ%3D%3D?line=26'>27</a>\u001b[0m     gen \u001b[39m=\u001b[39m Generator(\u001b[39m\"\u001b[39;49m\u001b[39mcuda\u001b[39;49m\u001b[39m\"\u001b[39;49m)\n\u001b[1;32m     <a href='vscode-notebook-cell://ood.unity.rc.umass.edu/work/pi_ahoumansadr_umass_edu/svkulkarni/Projects/distill-sd/progressive-distillation.ipynb#X22sdnNjb2RlLXJlbW90ZQ%3D%3D?line=27'>28</a>\u001b[0m     gen\u001b[39m.\u001b[39mmanual_seed(\u001b[39m1674753452\u001b[39m)\n\u001b[1;32m     <a href='vscode-notebook-cell://ood.unity.rc.umass.edu/work/pi_ahoumansadr_umass_edu/svkulkarni/Projects/distill-sd/progressive-distillation.ipynb#X22sdnNjb2RlLXJlbW90ZQ%3D%3D?line=28'>29</a>\u001b[0m     pipe \u001b[39m=\u001b[39m DiffusionPipeline\u001b[39m.\u001b[39mfrom_pretrained(path, torch_dtype\u001b[39m=\u001b[39mtorch\u001b[39m.\u001b[39mfloat16, safety_checker\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m, requires_safety_checker\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m)\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: device-side assert triggered\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n"
     ]
    }
   ],
   "source": [
    "from diffusers import DiffusionPipeline\n",
    "import torch\n",
    "\n",
    "import torch\n",
    "from diffusers import DiffusionPipeline\n",
    "from diffusers import DPMSolverMultistepScheduler\n",
    "from torch import Generator\n",
    "\n",
    "\n",
    "# path = 'segmind/portrait-finetuned' # Path to the appropriate model-type\n",
    "path = 'CompVis/stable-diffusion-v1-4'\n",
    "# path = 'segmind/small-sd'\n",
    "# Insert your prompt below.\n",
    "# prompt = \"Faceshot Portrait of pretty young (18-year-old) Caucasian wearing a high neck sweater, (masterpiece, extremely detailed skin, photorealistic, heavy shadow, dramatic and cinematic lighting, key light, fill light), sharp focus, BREAK epicrealism\"\n",
    "# # Insert negative prompt below. We recommend using this negative prompt for best results.\n",
    "# negative_prompt = \"(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck\" \n",
    "# prompt = \"a woman in a red and gold costume with feathers on her head\"\n",
    "# extra_prompt = \", facing the camera, photograph, highly detailed face, depth of field, moody light, style by Yasmin Albatoul, Harry Fayt, centered, extremely detailed, Nikon D850, award winning photography\"\n",
    "# negative_prompt = \"cartoon, anime, ugly, (aged, white beard, black skin, wrinkle:1.1), (bad proportions, unnatural feature, incongruous feature:1.4), (blurry, un-sharp, fuzzy, un-detailed skin:1.2), (facial contortion, poorly drawn face, deformed iris, deformed pupils:1.3), (mutated hands and fingers:1.5), disconnected hands, disconnected limbs\"\n",
    "prompt = \"Living in the light with a businesswoman\"\n",
    "\n",
    "torch.set_grad_enabled(False)\n",
    "torch.backends.cudnn.benchmark = True\n",
    "\n",
    "# Below code will run on gpu, please pass cpu everywhere as the device and set 'dtype' to torch.float32 for cpu inference.\n",
    "with torch.inference_mode():\n",
    "    gen = Generator(\"cuda\")\n",
    "    gen.manual_seed(1674753452)\n",
    "    pipe = DiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False)\n",
    "    pipe.to('cuda')\n",
    "    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)\n",
    "    pipe.unet.to(device='cuda', dtype=torch.float16, memory_format=torch.channels_last)\n",
    "\n",
    "    img = pipe(prompt=prompt, width=512, height=512, num_inference_steps=25, guidance_scale = 7, num_images_per_prompt=1, generator = gen).images[0]\n",
    "    img.save(\"businesswoman.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Distill-SD",
   "language": "python",
   "name": "distill-sd"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}