{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir(\"..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from diffusers.pipelines import FluxPipeline\n",
    "from PIL import Image\n",
    "\n",
    "from omini.pipeline.flux_omini import Condition, generate, seed_everything"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = FluxPipeline.from_pretrained(\n",
    "    \"black-forest-labs/FLUX.1-dev\", torch_dtype=torch.bfloat16\n",
    ")\n",
    "pipe = pipe.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe.unload_lora_weights()\n",
    "\n",
    "for style_type in [\"ghibli\", \"irasutoya\", \"simpsons\", \"snoopy\"]:\n",
    "    pipe.load_lora_weights(\n",
    "        \"Yuanshi/OminiControlArt\",\n",
    "        weight_name=f\"v0/{style_type}.safetensors\",\n",
    "        adapter_name=style_type,\n",
    "    )\n",
    "\n",
    "pipe.set_adapters([\"ghibli\", \"irasutoya\", \"simpsons\", \"snoopy\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def resize(img, factor=16):\n",
    "    # Resize the image to be divisible by the factor\n",
    "    w, h = img.size\n",
    "    new_w, new_h = w // factor * factor, h // factor * factor\n",
    "    padding_w, padding_h = (w - new_w) // 2, (h - new_h) // 2\n",
    "    img = img.crop((padding_w, padding_h, new_w + padding_w, new_h + padding_h))\n",
    "    return img\n",
    "\n",
    "\n",
    "def bound_image(image):\n",
    "    factor = 512 / max(image.size)\n",
    "    image = resize(\n",
    "        image.resize(\n",
    "            (int(image.size[0] * factor), int(image.size[1] * factor)),\n",
    "            Image.LANCZOS,\n",
    "        )\n",
    "    )\n",
    "    delta = (0, -image.size[0] // 16)\n",
    "    return image, delta\n",
    "\n",
    "sizes = {\n",
    "    \"2:3\": (640, 960),\n",
    "    \"1:1\": (640, 640),\n",
    "    \"3:2\": (960, 640),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = Image.open(\"assets/ominicontrol_art/DistractedBoyfriend.webp\").convert(\"RGB\")\n",
    "image, delta = bound_image(image)\n",
    "condition = Condition(image, \"ghibli\", position_delta=delta)\n",
    "\n",
    "seed_everything()\n",
    "\n",
    "size = sizes[\"3:2\"]\n",
    "\n",
    "result_img = generate(\n",
    "    pipe,\n",
    "    prompt=\"\",\n",
    "    conditions=[condition],\n",
    "    max_sequence_length=32,\n",
    "    width=size[0],\n",
    "    height=size[1],\n",
    "    image_guidance_scale=1.5,\n",
    ").images[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = Image.open(\"assets/ominicontrol_art/oiiai.png\").convert(\"RGB\")\n",
    "image, delta = bound_image(image)\n",
    "condition = Condition(image, \"irasutoya\", position_delta=delta)\n",
    "\n",
    "seed_everything()\n",
    "\n",
    "size = sizes[\"1:1\"]\n",
    "\n",
    "result_img = generate(\n",
    "    pipe,\n",
    "    prompt=\"\",\n",
    "    conditions=[condition],\n",
    "    max_sequence_length=32,\n",
    "    width=size[0],\n",
    "    height=size[1],\n",
    "    image_guidance_scale=1.5,\n",
    ").images[0]\n",
    "\n",
    "result_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = Image.open(\"assets/ominicontrol_art/breakingbad.jpg\").convert(\"RGB\")\n",
    "image, delta = bound_image(image)\n",
    "condition = Condition(image, \"simpsons\", position_delta=delta)\n",
    "\n",
    "seed_everything()\n",
    "\n",
    "size = sizes[\"3:2\"]\n",
    "\n",
    "result_img = generate(\n",
    "    pipe,\n",
    "    prompt=\"\",\n",
    "    conditions=[condition],\n",
    "    max_sequence_length=32,\n",
    "    width=size[0],\n",
    "    height=size[1],\n",
    "    image_guidance_scale=1.5,\n",
    ").images[0]\n",
    "\n",
    "result_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = Image.open(\"assets/ominicontrol_art/PulpFiction.jpg\").convert(\"RGB\")\n",
    "image, delta = bound_image(image)\n",
    "condition = Condition(image, \"snoopy\", position_delta=delta)\n",
    "\n",
    "seed_everything()\n",
    "\n",
    "size = sizes[\"3:2\"]\n",
    "\n",
    "result_img = generate(\n",
    "    pipe,\n",
    "    prompt=\"\",\n",
    "    conditions=[condition],\n",
    "    max_sequence_length=32,\n",
    "    width=size[0],\n",
    "    height=size[1],\n",
    "    image_guidance_scale=1.5,\n",
    ").images[0]\n",
    "\n",
    "result_img"
   ]
  }
 ],
 "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.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}