PhotoMaker_style_colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/PhotoMaker-colab/blob/main/PhotoMaker_style_colab.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjYy0F2gZIPR"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone -b dev https://github.com/camenduru/PhotoMaker\n",
        "%cd /content/PhotoMaker\n",
        "\n",
        "!pip install -q torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2+cu118 torchtext==0.15.2 torchdata==0.6.1 --extra-index-url https://download.pytorch.org/whl/cu118 -U\n",
        "!pip install -q xformers==0.0.20 diffusers accelerate einops onnxruntime-gpu omegaconf\n",
        "\n",
        "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\n",
        "\n",
        "# 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",
        "# https://civitai.com/api/download/models/276923\n",
        "base_model_path = './civitai_models/sdxlUnstableDiffusers_v11.safetensors'\n",
        "photomaker_path = hf_hub_download(repo_id=\"TencentARC/PhotoMaker\", filename=\"photomaker-v1.bin\", repo_type=\"model\")\n",
        "# https://civitai.com/api/download/models/152309?type=Model&format=SafeTensor\n",
        "lora_path = './civitai_models/xl_more_art-full.safetensors'\n",
        "\n",
        "device = \"cuda\"\n",
        "save_path = \"./outputs\"\n",
        "\n",
        "# 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": null,
      "metadata": {},
      "outputs": [],
      "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": null,
      "metadata": {},
      "outputs": [],
      "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": null,
      "metadata": {},
      "outputs": [],
      "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": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}