test_with_controlnet_openpose.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Usage:\n",
    "1. init\n",
    "2. load reference image and get pose\n",
    "3. test a single image with a single prompt (StableIdentity & ControlNet)\n",
    "4. test a single image with prompts (StableIdentity & ControlNet)\n",
    "5. test all images with prompts (StableIdentity & ControlNet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "from transformers import ViTModel, ViTImageProcessor\n",
    "from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, DPMSolverMultistepScheduler\n",
    "from utils import latents_to_images, downsampling\n",
    "from omegaconf import OmegaConf\n",
    "from accelerate.utils import set_seed\n",
    "from tqdm import tqdm\n",
    "from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipelineOutput\n",
    "from PIL import Image\n",
    "from models.celeb_embeddings import embedding_forward\n",
    "from controlnet_aux import OpenposeDetector\n",
    "from diffusers.utils import load_image, make_image_grid\n",
    "import numpy as np\n",
    "\n",
    "seed = 42\n",
    "set_seed(seed)\n",
    "torch.cuda.set_device(0)\n",
    "\n",
    "'''set the paths of pretrained models'''\n",
    "# vit face recognition\n",
    "vit_face_path = \"/home/user/.cache/huggingface/hub/vit-base-patch16-224-in21k-face-recognition\"\n",
    "# sd2.1\n",
    "model_path = \"/home/user/.cache/huggingface/hub/models--stabilityai--stable-diffusion-2-1/snapshots/5cae40e6a2745ae2b01ad92ae5043f95f23644d6\"\n",
    "# controlnet\n",
    "controlnet_path = \"/home/user/.cache/huggingface/hub/controlnetv11_21_diffuser\" # https://huggingface.co/thibaud/controlnet-sd21-openpose-diffusers\n",
    "# openpose\n",
    "openpose_path = \"models/openpose_models/\"\n",
    "\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "\n",
    "'''load the pretrained models'''\n",
    "# vit face recognition\n",
    "vit_face_recog_processor = ViTImageProcessor.from_pretrained(vit_face_path)    \n",
    "vit_face_recognition_model = ViTModel.from_pretrained(vit_face_path).to(device)    \n",
    "\n",
    "# openpose\n",
    "open_pose = OpenposeDetector.from_pretrained(openpose_path)       \n",
    "\n",
    "\n",
    "# sd2.1\n",
    "# model_id = \"stabilityai/stable-diffusion-2-1\"\n",
    "# pipe = StableDiffusionPipeline.from_pretrained(model_path)  \n",
    "pipe = StableDiffusionControlNetPipeline.from_pretrained(\n",
    "    model_path,\n",
    "    controlnet=ControlNetModel.from_pretrained(controlnet_path)\n",
    ")\n",
    "pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)\n",
    "pipe = pipe.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. load reference image and get pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pose_img_path = \"datasets_face/pose_input/1.png\"\n",
    "ref_image = load_image(pose_img_path)\n",
    "print(\"load_image\", pose_img_path)         \n",
    "pose_image = open_pose(ref_image, detect_resolution=512, image_resolution=512)\n",
    "pose_image = np.array(pose_image)[:, :, ::-1]           \n",
    "pose_image = Image.fromarray(np.uint8(pose_image))\n",
    "pose_image\n",
    "# pose_image.save('use_pose.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. test a single image with a single prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_emb_path = \"experiments512/save_00059/00059-450.pt\"\n",
    "test_emb = torch.load(test_emb_path).cuda()\n",
    "v1_emb = test_emb[:, 0]\n",
    "v2_emb = test_emb[:, 1]\n",
    "\n",
    "\n",
    "tokens = [\"v1*\", \"v2*\"]\n",
    "embeddings = [v1_emb, v2_emb]\n",
    "# add tokens and get ids\n",
    "pipe.tokenizer.add_tokens(tokens)\n",
    "token_ids = pipe.tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "# resize token embeddings and set new embeddings\n",
    "pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer), pad_to_multiple_of = 8)\n",
    "for token_id, embedding in zip(token_ids, embeddings):\n",
    "    pipe.text_encoder.get_input_embeddings().weight.data[token_id] = embedding\n",
    "\n",
    "prompt = \"v1* v2* wearing a Superman outfit, facing to camera, best quality, ultra high res\"\n",
    "image = pipe(prompt, image=pose_image, guidance_scale = 8.5).images[0]  # negative_prompt=\"ugly, blurry, bad, deformed, bad anatomy\"\n",
    "# image.save(prompt + \".png\")\n",
    "image\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. test a single image with prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = \"00059\"\n",
    "save_dir = \"results/\" + index + \"/with_controlnet\"\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "test_emb_path = f\"experiments512/save_{index}/{index}-450.pt\"\n",
    "test_emb = torch.load(test_emb_path).cuda()\n",
    "v1_emb = test_emb[:, 0]\n",
    "v2_emb = test_emb[:, 1]\n",
    "\n",
    "'''insert into tokenizer & embedding layer'''\n",
    "tokens = [\"v1*\", \"v2*\"]\n",
    "embeddings = [v1_emb, v2_emb]\n",
    "# add tokens and get ids\n",
    "pipe.tokenizer.add_tokens(tokens)\n",
    "token_ids = pipe.tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "# resize token embeddings and set new embeddings\n",
    "pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer), pad_to_multiple_of = 8)\n",
    "for token_id, embedding in zip(token_ids, embeddings):\n",
    "    pipe.text_encoder.get_input_embeddings().weight.data[token_id] = embedding\n",
    "\n",
    "prompts_list = [\"a photo of v1* v2*, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a Superman outfit, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a spacesuit, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a red sweater, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a purple wizard outfit, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a blue hoodie, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing headphones, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* with red hair, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing headphones with red hair, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a Christmas hat, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing sunglasses, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing sunglasses and necklace, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a blue cap, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a doctoral cap, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* with white hair, wearing glasses, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* in a helmet and vest riding a motorcycle, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* holding a bottle of red wine, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* driving a bus in the desert, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* playing basketball, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* playing the violin, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* piloting a spaceship, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* riding a horse, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* coding in front of a computer, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* laughing on the lawn, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* frowning at the camera, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* happily smiling, looking at the camera, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* crying disappointedly, with tears flowing, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing sunglasses, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* playing the guitar in the view of left side, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* holding a bottle of red wine, upper body, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing sunglasses and necklace, close-up, in the view of right side, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* riding a horse, in the view of the top, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* wearing a doctoral cap, upper body, with the left side of the face facing the camera, best quality, ultra high res\",\n",
    "\"v1* v2* crying disappointedly, with tears flowing, with left side of the face facing the camera, best quality, ultra high res\",\n",
    "\"v1* v2* sitting in front of the camera, with a beautiful purple sunset at the beach in the background, best quality, ultra high res\",\n",
    "\"v1* v2* swimming in the pool, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* climbing a mountain, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* skiing on the snowy mountain, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* in the snow, facing to camera, best quality, ultra high res\",\n",
    "\"v1* v2* in space wearing a spacesuit, facing to camera, best quality, ultra high res\",]\n",
    "\n",
    "for prompt in prompts_list:\n",
    "    image = pipe(prompt, image=pose_image, guidance_scale = 8.5).images[0]\n",
    "    image.save(os.path.join(save_dir, prompt.replace(\"v1* v2*\", \"a person\") + '.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. test all images with prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folders = \"datasets_face/test_data_demo\"\n",
    "folder_names = os.listdir(folders)\n",
    "\n",
    "for img_name in folder_names:\n",
    "    index = img_name[:-4]\n",
    "    '''test single image with prompts_list'''\n",
    "    save_dir = \"results/\" + index + \"/with_controlnet\"\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    test_emb_path = f\"experiments512/save_{index}/{index}-450.pt\"\n",
    "    test_emb = torch.load(test_emb_path).cuda()\n",
    "    v1_emb = test_emb[:, 0]\n",
    "    v2_emb = test_emb[:, 1]\n",
    "\n",
    "    '''insert into tokenizer & embedding layer'''\n",
    "    tokens = [\"v1*\", \"v2*\"]\n",
    "    embeddings = [v1_emb, v2_emb]\n",
    "    # add tokens and get ids\n",
    "    pipe.tokenizer.add_tokens(tokens)\n",
    "    token_ids = pipe.tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "    # resize token embeddings and set new embeddings\n",
    "    pipe.text_encoder.resize_token_embeddings(len(pipe.tokenizer), pad_to_multiple_of = 8)\n",
    "    for token_id, embedding in zip(token_ids, embeddings):\n",
    "        pipe.text_encoder.get_input_embeddings().weight.data[token_id] = embedding\n",
    "\n",
    "    prompts_list = [\"a photo of v1* v2*, facing to camera, best quality, ultra high res\",\n",
    "        \"v1* v2* wearing a Superman outfit, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a spacesuit, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a red sweater, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a purple wizard outfit, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a blue hoodie, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing headphones, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* with red hair, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing headphones with red hair, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a Christmas hat, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing sunglasses, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing sunglasses and necklace, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a blue cap, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a doctoral cap, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* with white hair, wearing glasses, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* in a helmet and vest riding a motorcycle, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* holding a bottle of red wine, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* driving a bus in the desert, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* playing basketball, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* playing the violin, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* piloting a spaceship, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* riding a horse, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* coding in front of a computer, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* laughing on the lawn, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* frowning at the camera, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* happily smiling, looking at the camera, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* crying disappointedly, with tears flowing, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing sunglasses, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* playing the guitar in the view of left side, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* holding a bottle of red wine, upper body, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing sunglasses and necklace, close-up, in the view of right side, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* riding a horse, in the view of the top, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* wearing a doctoral cap, upper body, with the left side of the face facing the camera, best quality, ultra high res\",\n",
    "    \"v1* v2* crying disappointedly, with tears flowing, with left side of the face facing the camera, best quality, ultra high res\",\n",
    "    \"v1* v2* sitting in front of the camera, with a beautiful purple sunset at the beach in the background, best quality, ultra high res\",\n",
    "    \"v1* v2* swimming in the pool, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* climbing a mountain, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* skiing on the snowy mountain, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* in the snow, facing to camera, best quality, ultra high res\",\n",
    "    \"v1* v2* in space wearing a spacesuit, facing to camera, best quality, ultra high res\",\n",
    "    ]\n",
    "\n",
    "    for prompt in prompts_list:\n",
    "        image = pipe(prompt, image=pose_image, guidance_scale = 8.5).images[0]\n",
    "        image.save(os.path.join(save_dir, prompt.replace(\"v1* v2*\", \"a person\") + '.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ldm",
   "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.8.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}