{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-14 13:06:44.852659: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-01-14 13:06:45.609743: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "from PIL import Image\n",
    "\n",
    "from diffusers.utils import load_image\n",
    "from diffusers import DDIMScheduler\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "from photomaker.pipeline import PhotoMakerStableDiffusionXLPipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gloal variable and function\n",
    "def image_grid(imgs, rows, cols, size_after_resize):\n",
    "    assert len(imgs) == rows*cols\n",
    "\n",
    "    w, h = size_after_resize, size_after_resize\n",
    "    \n",
    "    grid = Image.new('RGB', size=(cols*w, rows*h))\n",
    "    grid_w, grid_h = grid.size\n",
    "    \n",
    "    for i, img in enumerate(imgs):\n",
    "        img = img.resize((w,h))\n",
    "        grid.paste(img, box=(i%cols*w, i//cols*h))\n",
    "    return grid\n",
    "\n",
    "base_model_path = './civitai_models/sdxlUnstableDiffusers_v11.safetensors'\n",
    "photomaker_path = './release_model/photomaker-v1.bin'\n",
    "lora_path = './civitai_models/xl_more_art-full.safetensors'\n",
    "\n",
    "device = \"cuda\"\n",
    "save_path = \"./outputs\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading PhotoMaker components [1] id_encoder from [./release_model]...\n",
      "Loading PhotoMaker components [2] lora_weights from [./release_model]\n",
      "Loading lora...\n"
     ]
    }
   ],
   "source": [
    "# Load base model\n",
    "pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(\n",
    "    base_model_path, \n",
    "    torch_dtype=torch.bfloat16, \n",
    "    original_config_file=None,\n",
    ").to(device)\n",
    "\n",
    "# Load PhotoMaker checkpoint\n",
    "pipe.load_photomaker_adapter(\n",
    "    os.path.dirname(photomaker_path),\n",
    "    subfolder=\"\",\n",
    "    weight_name=os.path.basename(photomaker_path),\n",
    "    trigger_word=\"img\"\n",
    ")     \n",
    "\n",
    "pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)\n",
    "print(\"Loading lora...\")\n",
    "pipe.load_lora_weights(os.path.dirname(lora_path), weight_name=os.path.basename(lora_path), adapter_name=\"xl_more_art-full\")\n",
    "pipe.set_adapters([\"photomaker\", \"xl_more_art-full\"], adapter_weights=[1.0, 0.5])\n",
    "pipe.fuse_lora()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input ID images:\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 12476 characters]",
      "image/png": "[encoded data, 113596 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=224x224>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# define and show the input ID images\n",
    "image_path = './examples/scarletthead_woman/scarlett_0.jpg'\n",
    "\n",
    "input_id_images = []\n",
    "input_id_images.append(load_image(image_path))\n",
    "    \n",
    "input_grid = image_grid(input_id_images, 1, 1, size_after_resize=224)\n",
    "print(\"Input ID images:\")\n",
    "input_grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 2048]) torch.Size([1, 2048])\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b17df91cd07a46d082663a8ae37c3036",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Note that the trigger word `img` must follow the class word for personalization\n",
    "prompt = \"A girl img riding dragon over a whimsical castle, 3d CGI, art by Pixar, half-body, screenshot from animation\"\n",
    "negative_prompt = \"realistic, photo-realistic, bad quality, bad anatomy, worst quality, low quality, lowres, extra fingers, blur, blurry, ugly, wrong proportions, watermark, image artifacts, bad eyes, bad hands, bad arms\"\n",
    "generator = torch.Generator(device=device).manual_seed(42)\n",
    "\n",
    "## Parameter setting\n",
    "num_steps = 50\n",
    "style_strength_ratio = 20\n",
    "start_merge_step = int(float(style_strength_ratio) / 100 * num_steps)\n",
    "if start_merge_step > 30:\n",
    "    start_merge_step = 30\n",
    "    \n",
    "images = pipe(\n",
    "    prompt=prompt,\n",
    "    input_id_images=input_id_images,\n",
    "    negative_prompt=negative_prompt,\n",
    "    num_images_per_prompt=4,\n",
    "    num_inference_steps=num_steps,\n",
    "    start_merge_step=start_merge_step,\n",
    "    generator=generator,\n",
    ").images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results:\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 295896 characters]",
      "image/png": "[encoded data, 2789684 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=2048x512>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Show and save the results\n",
    "## Downsample for visualization\n",
    "grid = image_grid(images, 1, 4, size_after_resize=512)\n",
    "\n",
    "os.makedirs(save_path, exist_ok=True)\n",
    "for idx, image in enumerate(images):\n",
    "    image.save(os.path.join(save_path, f\"photomaker_style_{idx:02d}.png\"))\n",
    "    \n",
    "print(\"Results:\")\n",
    "grid"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python(env-novelai)",
   "language": "python",
   "name": "env-novelai"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}