{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import ( \n",
    "    DiffusionPipeline,\n",
    "    AutoencoderKL,\n",
    "    StableDiffusionXLPipeline\n",
    ")\n",
    "from diffusers.utils import make_image_grid\n",
    "import torch\n",
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from diffusers.utils.torch_utils import randn_tensor\n",
    "import glob\n",
    "import random\n",
    "import itertools\n",
    "from IPython.display import clear_output\n",
    "from diffusers.utils.torch_utils import randn_tensor\n",
    "\n",
    "negative_prompt = None\n",
    "model_id = '/data/model/stable-diffusion-xl-base-1.0'\n",
    "vae_id = '/data/model/sdxl-vae-fp16-fix'\n",
    "\n",
    "device = 'cuda:0'\n",
    "vae = AutoencoderKL.from_pretrained(vae_id, torch_dtype=torch.float16)\n",
    "pipe = StableDiffusionXLPipeline.from_pretrained(\n",
    "    model_id,\n",
    "    vae=vae,\n",
    "    torch_dtype=torch.float16, \n",
    "    use_safetensors=True, \n",
    "    variant=\"fp16\"\n",
    "    )\n",
    "\n",
    "pipe = pipe.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bbox_mask(mask_image):\n",
    "    \"\"\"\n",
    "    mask: PIL.Image\n",
    "\n",
    "    return: list\n",
    "    \"\"\"\n",
    "    # Convert the PIL image to a NumPy array\n",
    "    mask = np.array(mask_image)\n",
    "\n",
    "    # Apply a threshold to binarize the mask\n",
    "    threshold_value = 128  # You might need to adjust this value\n",
    "    mask_binary = (mask > threshold_value).astype(np.uint8) * 255\n",
    "    \n",
    "    # Find the bounding box of the non-zero regions in the binary mask\n",
    "    rows = np.any(mask_binary, axis=1)\n",
    "    cols = np.any(mask_binary, axis=0)\n",
    "    ymin, ymax = np.where(rows)[0][[0, -1]]\n",
    "    xmin, xmax = np.where(cols)[0][[0, -1]]\n",
    "\n",
    "    # The bounding box is defined by the coordinates (xmin, ymin) and (xmax, ymax).\n",
    "    bbox = (xmin, ymin, xmax, ymax)\n",
    "\n",
    "    return bbox\n",
    "\n",
    "def image_process(image, mask, margin=32, fixed_scale=None):\n",
    "    \"\"\"\n",
    "    Crop the image and mask, based on the bbox of mask map.\n",
    "    The resize width is 512, and there is an optional mergin.\n",
    "\n",
    "    image: PIL.Image\n",
    "    mask: PIL.Image\n",
    "    \"\"\"\n",
    "    original_size = image.size[0]\n",
    "    # image = image.resize((1024, 1024))\n",
    "    # mask = mask.resize((1024, 1024))\n",
    "\n",
    "    \n",
    "    bbox = bbox_mask(mask)\n",
    "    image = image.crop((bbox[0], bbox[1], bbox[2], bbox[3]))\n",
    "    mask = mask.crop((bbox[0], bbox[1], bbox[2], bbox[3]))\n",
    "    width, height = image.size\n",
    "\n",
    "    if fixed_scale is not None:\n",
    "        rescaling_factor = min((512 - margin) / width, (1024 - margin) / height) * fixed_scale \n",
    "        #(512 - margin) / width * fixed_scale\n",
    "    else:\n",
    "        max_factor = min((512 - margin) / width, (1024 - margin) / height)\n",
    "        # rescaling_factor = np.random.uniform(0.75, max_factor)\n",
    "        rescaling_factor = np.random.uniform(max_factor * 0.75, max_factor)\n",
    "\n",
    "    # rescaling_factor = max_factor\n",
    "    # print(rescaling_factor)\n",
    "    \n",
    "    # new_size = int(original_size * rescaling_factor)\n",
    "\n",
    "    image = image.resize((int(width * rescaling_factor), int(height * rescaling_factor)))\n",
    "    mask = mask.resize((int(width * rescaling_factor), int(height * rescaling_factor)))\n",
    "\n",
    "    return image, mask \n",
    "\n",
    "\n",
    "def image_collage(image_0, mask_0, image_1, mask_1, white=False, margin_max=10, height_sync=False):\n",
    "    rgb_back = 255 if white else 0\n",
    "    background = np.zeros((1024, 1024, 3), dtype=np.uint8)\n",
    "    mask_map = np.zeros((1024, 1024))\n",
    "\n",
    "    if height_sync:\n",
    "        y_max_sync = 1024 - max(image_0.size[1], image_1.size[1])\n",
    "        y_sync = random.randint(0, y_max_sync)\n",
    "        y_pos_0, y_pos_1 = y_sync, y_sync\n",
    "    else:\n",
    "        y_max_0 = 1024 - image_0.size[1]\n",
    "        y_max_1 = 1024 - image_1.size[1]\n",
    "        y_pos_0 = random.randint(0, y_max_0)\n",
    "        y_pos_1 = random.randint(0, y_max_1)\n",
    "\n",
    "    image_0 = np.array(image_0)\n",
    "    binary_mask_0 = np.array(mask_0) > 122\n",
    "    image_0[binary_mask_0 == 0] = 0 # background goes to 0\n",
    "\n",
    "    image_1 = np.array(image_1)\n",
    "    binary_mask_1 = np.array(mask_1) > 122\n",
    "    image_1[binary_mask_1 == 0] = 0 \n",
    "    \n",
    "    # Determin which image goes on top\n",
    "    is_left = True if random.random() < 0.5 else False\n",
    "\n",
    "    # margin for x axis\n",
    "    left_margin = random.randint(0, margin_max)\n",
    "    right_margin = random.randint(0, margin_max)\n",
    "\n",
    "    mask_map[y_pos_0:y_pos_0 + binary_mask_0.shape[0], left_margin:left_margin + binary_mask_0.shape[1]] += binary_mask_0\n",
    "    mask_map[y_pos_1:y_pos_1 + binary_mask_1.shape[0], -right_margin - binary_mask_1.shape[1] -1:-right_margin-1] += binary_mask_1\n",
    "\n",
    "    merged_mask_map = mask_map > 0\n",
    "    union_mask_map = mask_map > 1.5\n",
    "\n",
    "    if is_left:\n",
    "        background[y_pos_0:y_pos_0 + image_0.shape[0], left_margin:left_margin + image_0.shape[1]] = image_0 \n",
    "        background[union_mask_map] = 0\n",
    "        background[y_pos_1:y_pos_1 + image_1.shape[0], -right_margin - image_1.shape[1] -1:-right_margin-1] += image_1 \n",
    "    else:\n",
    "        background[y_pos_1:y_pos_1 + image_1.shape[0], -right_margin - image_1.shape[1] -1:-right_margin-1] = image_1\n",
    "        background[union_mask_map] = 0\n",
    "        background[y_pos_0:y_pos_0 + image_0.shape[0], left_margin:left_margin + image_0.shape[1]] += image_0 \n",
    "    \n",
    "    if white:\n",
    "        background[merged_mask_map == 0] = 255\n",
    "\n",
    "    return background, merged_mask_map\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None):\n",
    "    # get the original timestep using init_timestep\n",
    "    if denoising_start is None:\n",
    "        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)\n",
    "        t_start = max(num_inference_steps - init_timestep, 0)\n",
    "    else:\n",
    "        t_start = 0\n",
    "\n",
    "    timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]\n",
    "\n",
    "    # Strength is irrelevant if we directly request a timestep to start at;\n",
    "    # that is, strength is determined by the denoising_start instead.\n",
    "    if denoising_start is not None:\n",
    "        discrete_timestep_cutoff = int(\n",
    "            round(\n",
    "                self.scheduler.config.num_train_timesteps\n",
    "                - (denoising_start * self.scheduler.config.num_train_timesteps)\n",
    "            )\n",
    "        )\n",
    "\n",
    "        num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()\n",
    "        if self.scheduler.order == 2 and num_inference_steps % 2 == 0:\n",
    "            # if the scheduler is a 2nd order scheduler we might have to do +1\n",
    "            # because `num_inference_steps` might be even given that every timestep\n",
    "            # (except the highest one) is duplicated. If `num_inference_steps` is even it would\n",
    "            # mean that we cut the timesteps in the middle of the denoising step\n",
    "            # (between 1st and 2nd devirative) which leads to incorrect results. By adding 1\n",
    "            # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler\n",
    "            num_inference_steps = num_inference_steps + 1\n",
    "\n",
    "        # because t_n+1 >= t_n, we slice the timesteps starting from the end\n",
    "        timesteps = timesteps[-num_inference_steps:]\n",
    "        return timesteps, num_inference_steps\n",
    "\n",
    "    return timesteps, num_inference_steps - t_start\n",
    "\n",
    "def add_noise(pipe, latents, strength=1.0, noise=None):\n",
    "    # if pipe has no get_timesteps method\n",
    "    if not hasattr(pipe, 'get_timesteps'):\n",
    "        pipe.get_timesteps = partial(get_timesteps, pipe)\n",
    "    \n",
    "    num_inference_steps = 50\n",
    "    device = latents.device\n",
    "    noise = randn_tensor(latents.shape, device=device, dtype=latents.dtype) if noise is None else noise\n",
    "    pipe.scheduler.set_timesteps(num_inference_steps, device=device)\n",
    "    timesteps, _ = pipe.get_timesteps(\n",
    "        num_inference_steps,\n",
    "        strength,\n",
    "        device,\n",
    "        denoising_start=None,\n",
    "    )\n",
    "    latent_timestep = timesteps[:1]\n",
    "    # scaled_latents = latents * pipe.vae.config.scaling_factor\n",
    "    return pipe.scheduler.add_noise(latents, noise, latent_timestep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_0 = Image.open('dataset/category/actionfigure_2_and_dog0/a02.jpg')\n",
    "image_1 = Image.open('dataset/category/actionfigure_2_and_dog0/b01.jpg')\n",
    "\n",
    "mask_0 = Image.open('dataset/category/actionfigure_2_and_dog0/mask_a02.jpg')\n",
    "mask_1 = Image.open('dataset/category/actionfigure_2_and_dog0/mask_b01.jpg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# image collage\n",
    "image_0_, mask_0_ = image_process(image_0, mask_0)\n",
    "image_1_, mask_1_ = image_process(image_1, mask_1)\n",
    "bg, mask = image_collage(image_0_, mask_0_, image_1_, mask_1_, white=True, margin_max=256, height_sync=True)\n",
    "mask_ = Image.fromarray(mask)\n",
    "mask_ = mask_.resize((mask_.size[0] // 8, mask_.size[1] // 8))\n",
    "    \n",
    "mask_ = np.array(mask_)\n",
    "mask_binary = torch.tensor(mask_).to(device='cuda')\n",
    "with torch.no_grad():\n",
    "    out = pipe.vae.encode(pipe.image_processor.preprocess(Image.fromarray(bg)).to('cuda', torch.float16)).latent_dist.sample()\n",
    "init_latent = out * mask_binary.to(out.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 99056 characters]",
      "image/png": "[encoded data, 1250288 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1024x1024>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# visualize\n",
    "with torch.no_grad():\n",
    "    image = pipe.vae.decode(init_latent, return_dict=False)[0]    \n",
    "    image = pipe.image_processor.postprocess(image, output_type='pil')\n",
    "image[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.unload_lora_weights()\n",
    "pipe.load_lora_weights('output/actionfigure_2_and_dog0_name_p0.5_s0_dco/save-1800')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:07<00:00,  6.40it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 115008 characters]",
      "image/png": "[encoded data, 1932184 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1024x1024>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# w/o our Inference\n",
    "prompt = 'olis harry potter toy and hta corgi in a garden full of flowers. high quality'\n",
    "negative_prompt = \"sticker, collage, anime, cartoon, graphic, text, painting, crayon, graphite, abstract glitch, blurry\"\n",
    "seed = 10\n",
    "generator = torch.Generator(device=device).manual_seed(seed)\n",
    "\n",
    "noise = randn_tensor(init_latent.shape, device=pipe.unet.device, dtype=torch.float16, generator=generator)\n",
    "init_latent_ = noise\n",
    "img = pipe(prompt=prompt, negative_prompt=negative_prompt, \n",
    "           latents=init_latent_, cross_attention_kwargs={'scale': 1}).images[0]\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 50/50 [00:08<00:00,  6.24it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 123732 characters]",
      "image/png": "[encoded data, 1869232 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=1024x1024>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# w/ our Inference\n",
    "prompt = 'olis harry potter toy and hta corgi in a garden full of flowers. high quality'\n",
    "negative_prompt = \"sticker, collage, anime, cartoon, graphic, text, painting, crayon, graphite, abstract glitch, blurry\"\n",
    "seed = 10\n",
    "gamma = 1\n",
    "generator = torch.Generator(device=device).manual_seed(seed)\n",
    "\n",
    "noise = randn_tensor(init_latent.shape, device=pipe.unet.device, dtype=torch.float16, generator=generator)\n",
    "init_latent_ = add_noise(pipe, pipe.vae.config.scaling_factor * init_latent * gamma, strength=1., noise=noise)\n",
    "init_latent_ /= pipe.scheduler.init_noise_sigma\n",
    "\n",
    "img = pipe(prompt=prompt, negative_prompt=negative_prompt, \n",
    "           latents=init_latent_, cross_attention_kwargs={'scale': 1}).images[0]\n",
    "img"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mudi",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
