train_1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d9587e5d-7ed2-4280-af92-10b4e6894f91",
   "metadata": {},
   "source": [
    "## Comments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afeac06a-3f21-4c85-9220-3bee88b8a95b",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "244e6704-e54c-4000-a5ad-0a51de4fdcc4",
   "metadata": {},
   "source": [
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "542b151d-5fb7-4e4c-bb70-1261ecfd6037",
   "metadata": {},
   "source": [
    "# T2I Train scripts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9a4a3b3-f5ae-4433-9b33-1015995a7b9d",
   "metadata": {},
   "source": [
    "Install necessary libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79fa5b7a-8365-40d6-89cb-163890603551",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install torch==2.0.1 --force-reinstall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01ec97ac-d5e2-4a5c-8f71-e1bb10311c95",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install pytorch_lightning==1.06"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48effa63-bfcd-4a1f-87a9-e241eb639f7f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install opencv-python basicsr omegaconf datasets diffusers transformers einops xformers==0.0.22 open-clip-torch==2.24.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58889edf-40c7-443d-b5cd-0bf7b5416c76",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/cu117"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b548d2c4-a9e6-4416-a385-37eb458054fd",
   "metadata": {},
   "source": [
    "Environment tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "855f28ec-1439-4ae3-921c-4fac48c400b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad90649c-f28b-4076-973b-418d303767ac",
   "metadata": {},
   "source": [
    "Test if modifying amount of memory for Torch with 'set_per_process_memory_fraction(0.5, 0)' works"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "236bb310-7934-4dc8-9db5-f60dc55af038",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# https://github.com/pytorch/pytorch/blob/47aa2536328afc51876b2e04384c0cfe71ee1f06/test/test_cuda.py#L394-L421\n",
    "import torch \n",
    "\n",
    "tensor = torch.zeros(1024, device='cuda')\n",
    "torch.cuda.empty_cache()\n",
    "total_memory = torch.cuda.get_device_properties(0).total_memory\n",
    "torch.cuda.set_per_process_memory_fraction(0.5, 0)\n",
    "\n",
    "# test 0.499 allocation is ok.\n",
    "application = int(total_memory * 0.499) - torch.cuda.max_memory_reserved()\n",
    "tmp_tensor = torch.empty(application, dtype=torch.int8, device='cuda')\n",
    "del tmp_tensor\n",
    "torch.cuda.empty_cache()\n",
    "\n",
    "application = int(total_memory * 0.5)\n",
    "# it will get OOM when try to allocate more than half memory.\n",
    "with self.assertRaisesRegex(RuntimeError, \"out of memory\"):\n",
    "    torch.empty(application, dtype=torch.int8, device='cuda')\n",
    "\n",
    "# ensure out of memory error doesn't disturb subsequent kernel\n",
    "tensor.fill_(1)\n",
    "self.assertTrue((tensor == 1).all())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a0ba89-4c4b-4878-9e33-43bf26832f54",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.cuda.is_available(), torch.cuda.device_count(), torch.cuda.current_device(), torch.cuda.device(0), torch.cuda.get_device_name(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96ba32bb-c891-43ce-966d-375979acd20a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = 1\n",
    "t = torch.cuda.get_device_properties(device).total_memory\n",
    "r = torch.cuda.memory_reserved(device)\n",
    "a = torch.cuda.memory_allocated(device)\n",
    "f = r-a  # free inside reserved\n",
    "t, r, a, f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "047e801e-1feb-4a9f-865d-3a0592fd9360",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import gc\n",
    "\n",
    "gc.collect()\n",
    "\n",
    "import torch\n",
    "torch.cuda.empty_cache() \n",
    "del gc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94cd40ae-23c4-4399-940b-74684ef51715",
   "metadata": {},
   "source": [
    "## Create Dataset (LLAVA + SHIFT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffca7eb3-19e8-4ab1-91e8-ed324943b840",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/sample_images_train.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85416d36-e8ec-4507-aabf-6dbfb1000873",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/llava_dataset_creation.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27c8607d-4a07-4f1b-85bc-502243ee1ec4",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 1 Condition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ba16019-b438-46bf-8c75-b6797264a1c7",
   "metadata": {},
   "source": [
    "### train_sketch.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "474369c6-27ec-4f3a-a1a1-8207148882d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!accelerate launch T2I-Adapter/train_sketch.py \n",
    "#!python -m torch.distributed.launch T2I-Adapter/train_sketch.py --bsize 1 --epochs 1 --num_workers 1 --n_samples 1 --scale 1 --gpus 1\n",
    "# H = Height pixels / 4\n",
    "# W = Width pixels / 4\n",
    "!torchrun T2I-Adapter/train_sketch.py --local_rank 1 --bsize 1 --epochs 5 --num_workers 1 --n_samples 1 --scale 1 --gpus 1 --H 256 --W 256 "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "627d9b60-4c22-4704-ad24-9bfcec6af2d2",
   "metadata": {},
   "source": [
    "* Samples are created in  T2I-Adapter/experiments/train-sketch/visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82648d29-aff0-48e4-ac5d-bb2422daaa44",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Running it in notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1ee90b-39cf-4e09-a1ee-edc27895653e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "import os.path as osp\n",
    "import time\n",
    "os.chdir(\"T2I-Adapter\")\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from basicsr.utils import (\n",
    "    get_env_info,\n",
    "    get_root_logger,\n",
    "    get_time_str,\n",
    "    img2tensor,\n",
    "    scandir,\n",
    "    tensor2img,\n",
    ")\n",
    "from basicsr.utils.options import copy_opt_file, dict2str\n",
    "from omegaconf import OmegaConf\n",
    "from PIL import Image\n",
    "\n",
    "# from ldm.data.dataset_coco import dataset_coco_mask_color\n",
    "from ldm.data.dataset_fill import dataset_fill_mask\n",
    "\n",
    "from dist_util import get_bare_model, get_dist_info, init_dist, master_only\n",
    "from ldm.models.diffusion.ddim import DDIMSampler\n",
    "from ldm.models.diffusion.dpm_solver import DPMSolverSampler\n",
    "from ldm.models.diffusion.plms import PLMSSampler\n",
    "from ldm.modules.encoders.adapter import Adapter\n",
    "from ldm.util import instantiate_from_config\n",
    "from ldm.modules.extra_condition.model_edge import pidinet\n",
    "\n",
    "\n",
    "\n",
    "def load_model_from_config(config, ckpt, verbose=False):\n",
    "    print(f\"Loading model from {ckpt}\")\n",
    "    pl_sd = torch.load(ckpt, map_location=\"cpu\")\n",
    "    if \"global_step\" in pl_sd:\n",
    "        print(f\"Global Step: {pl_sd['global_step']}\")\n",
    "    sd = pl_sd[\"state_dict\"]\n",
    "    model = instantiate_from_config(config.model)\n",
    "    m, u = model.load_state_dict(sd, strict=False)\n",
    "    if len(m) > 0 and verbose:\n",
    "        print(\"missing keys:\")\n",
    "        print(m)\n",
    "    if len(u) > 0 and verbose:\n",
    "        print(\"unexpected keys:\")\n",
    "        print(u)\n",
    "\n",
    "    model.cuda()\n",
    "    model.eval()\n",
    "    return model\n",
    "\n",
    "\n",
    "@master_only\n",
    "def mkdir_and_rename(path):\n",
    "    \"\"\"mkdirs. If path exists, rename it with timestamp and create a new one.\n",
    "\n",
    "    Args:\n",
    "        path (str): Folder path.\n",
    "    \"\"\"\n",
    "    if osp.exists(path):\n",
    "        new_name = path + \"_archived_\" + get_time_str()\n",
    "        print(f\"Path already exists. Rename it to {new_name}\", flush=True)\n",
    "        os.rename(path, new_name)\n",
    "    os.makedirs(path, exist_ok=True)\n",
    "    os.makedirs(osp.join(experiments_root, \"models\"))\n",
    "    os.makedirs(osp.join(experiments_root, \"training_states\"))\n",
    "    os.makedirs(osp.join(experiments_root, \"visualization\"))\n",
    "\n",
    "\n",
    "def load_resume_state(opt):\n",
    "    resume_state_path = None\n",
    "    if opt.auto_resume:\n",
    "        state_path = osp.join(\"experiments\", opt.name, \"training_states\")\n",
    "        if osp.isdir(state_path):\n",
    "            states = list(\n",
    "                scandir(state_path, suffix=\"state\", recursive=False, full_path=False)\n",
    "            )\n",
    "            if len(states) != 0:\n",
    "                states = [float(v.split(\".state\")[0]) for v in states]\n",
    "                resume_state_path = osp.join(state_path, f\"{max(states):.0f}.state\")\n",
    "                opt.resume_state_path = resume_state_path\n",
    "    # else:\n",
    "    #     if opt['path'].get('resume_state'):\n",
    "    #         resume_state_path = opt['path']['resume_state']\n",
    "    if resume_state_path is None:\n",
    "        resume_state = None\n",
    "    else:\n",
    "        device_id = torch.cuda.current_device()\n",
    "        resume_state = torch.load(\n",
    "            resume_state_path, map_location=lambda storage, loc: storage.cuda(device_id)\n",
    "        )\n",
    "        # check_resume(opt, resume_state['iter'])\n",
    "    return resume_state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a491b594-628a-4d07-b805-558e790f7d16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Opt:\n",
    "    def __init__(self):\n",
    "        # Default attribute values\n",
    "        self.bsize = 8\n",
    "        self.epochs = 10000\n",
    "        self.num_workers = 8\n",
    "        self.use_shuffle = True\n",
    "        self.dpm_solver = False\n",
    "        self.plms = False\n",
    "        self.auto_resume = False\n",
    "        self.ckpt = \"T2I-Adapter/models/sd-v1-4.ckpt\"\n",
    "        self.config = \"configs/stable-diffusion/train_sketch.yaml\"\n",
    "        self.print_fq = 100  # TODO:\n",
    "        self.H = 512\n",
    "        self.W = 512\n",
    "        self.C = 4\n",
    "        self.f = 8\n",
    "        self.ddim_steps = 50\n",
    "        self.n_samples = 1\n",
    "        self.ddim_eta = 0.0\n",
    "        self.scale = 7.5\n",
    "        self.gpus = [0, 1, 2, 3]\n",
    "        self.local_rank = 0\n",
    "        self.launcher = \"pytorch\"\n",
    "        self.l_cond = 4\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a542b1b-03af-470e-9430-1de83298d652",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "opt = Opt()\n",
    "\n",
    "print(\"-> Configuration from\", opt.config)\n",
    "config = OmegaConf.load(f\"{opt.config}\")\n",
    "opt.name = config[\"name\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b21c715-4113-4f1f-96b5-ae342330898b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# distributed setting\n",
    "print(\"-> Distributed setting\")\n",
    "init_dist(opt.launcher)\n",
    "torch.backends.cudnn.benchmark = True\n",
    "device = \"cuda\"\n",
    "torch.cuda.set_device(opt.local_rank)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "834e4bfc-b4b3-4dbb-a56a-ddafc7d1291c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset\n",
    "print(\"-> Dataset\")\n",
    "# FIXME: previous dataset coco2017\n",
    "# path_json_train = \"coco_stuff/mask/annotations/captions_train2017.json\"\n",
    "# path_json_val = \"coco_stuff/mask/annotations/captions_val2017.json\"\n",
    "# train_dataset = dataset_coco_mask_color(\n",
    "#     path_json_train,\n",
    "#     root_path_im=\"coco/train2017\",\n",
    "#     root_path_mask=\"coco_stuff/mask/train2017_color\",\n",
    "#     image_size=512,\n",
    "# )\n",
    "# train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)\n",
    "# val_dataset = dataset_coco_mask_color(\n",
    "#     path_json_val,\n",
    "#     root_path_im=\"coco/val2017\",\n",
    "#     root_path_mask=\"coco_stuff/mask/val2017_color\",\n",
    "#     image_size=512,\n",
    "# )\n",
    "# train_dataloader = torch.utils.data.DataLoader(\n",
    "#     train_dataset,\n",
    "#     batch_size=opt.bsize,\n",
    "#     shuffle=(train_sampler is None),\n",
    "#     num_workers=opt.num_workers,\n",
    "#     pin_memory=True,\n",
    "#     sampler=train_sampler,\n",
    "# )\n",
    "# val_dataloader = torch.utils.data.DataLoader(\n",
    "#     val_dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=False\n",
    "# )\n",
    "\n",
    "# Add data and validation loader with very small dataset\n",
    "print(\"-> Add data and validation loader with very small dataset\")\n",
    "train_dataset = dataset_fill_mask(split=\"train\")\n",
    "train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)\n",
    "val_dataset = dataset_fill_mask(split=\"train\")\n",
    "train_dataloader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=opt.bsize,\n",
    "    shuffle=(train_sampler is None),\n",
    "    num_workers=opt.num_workers,\n",
    "    pin_memory=True,\n",
    "    sampler=train_sampler,\n",
    ")\n",
    "val_dataloader = torch.utils.data.DataLoader(\n",
    "    val_dataset, batch_size=1, shuffle=False, num_workers=1, pin_memory=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37a613eb-6d12-4f85-a8e0-06e1440846d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# edge_generator\n",
    "os.chdir(\"..\")\n",
    "print(\"-> Edge generator\")\n",
    "net_G = pidinet()\n",
    "ckp = torch.load(\"T2I-Adapter/models/table5_pidinet.pth\", map_location=\"cpu\")[\n",
    "    \"state_dict\"\n",
    "]\n",
    "net_G.load_state_dict({k.replace(\"module.\", \"\"): v for k, v in ckp.items()})\n",
    "net_G.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91f56c50-4e1a-4c60-a712-bca7ff796cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# stable diffusion\n",
    "print(\"-> Stable diffusion\")\n",
    "model = load_model_from_config(config, f\"{opt.ckpt}\").to(device)\n",
    "\n",
    "# sketch encoder\n",
    "print(\"-> Sketch encoder\")\n",
    "model_ad = Adapter(\n",
    "    channels=[320, 640, 1280, 1280][:4], nums_rb=2, ksize=1, sk=True, use_conv=False\n",
    ").to(device)\n",
    "\n",
    "# # to gpus\n",
    "print(\"-> to gpus\")\n",
    "model_ad = torch.nn.parallel.DistributedDataParallel(\n",
    "    model_ad, device_ids=[opt.local_rank], output_device=opt.local_rank\n",
    ")\n",
    "model = torch.nn.parallel.DistributedDataParallel(\n",
    "    model, device_ids=[opt.local_rank], output_device=opt.local_rank\n",
    ")\n",
    "# device_ids=[torch.cuda.current_device()])\n",
    "net_G = torch.nn.parallel.DistributedDataParallel(\n",
    "    net_G, device_ids=[opt.local_rank], output_device=opt.local_rank\n",
    ")\n",
    "# device_ids=[torch.cuda.current_device()])\n",
    "\n",
    "# optimizer\n",
    "print(\"-> optimizer\")\n",
    "params = list(model_ad.parameters())\n",
    "optimizer = torch.optim.AdamW(params, lr=config[\"training\"][\"lr\"])\n",
    "\n",
    "experiments_root = osp.join(\"T2I-Adapter/experiments\", opt.name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "082ed694-6f7c-457f-845a-4ce944cb54b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# resume state\n",
    "print(\"-> resume state\")\n",
    "resume_state = load_resume_state(opt)\n",
    "if resume_state is None:\n",
    "    mkdir_and_rename(experiments_root)\n",
    "    start_epoch = 0\n",
    "    current_iter = 0\n",
    "    # WARNING: should not use get_root_logger in the above codes, including the called functions\n",
    "    # Otherwise the logger will not be properly initialized\n",
    "    log_file = osp.join(experiments_root, f\"train_{opt.name}_{get_time_str()}.log\")\n",
    "    logger = get_root_logger(\n",
    "        logger_name=\"basicsr\", log_level=logging.INFO, log_file=log_file\n",
    "    )\n",
    "    logger.info(get_env_info())\n",
    "    logger.info(dict2str(config))\n",
    "else:\n",
    "    # WARNING: should not use get_root_logger in the above codes, including the called functions\n",
    "    # Otherwise the logger will not be properly initialized\n",
    "    log_file = osp.join(experiments_root, f\"train_{opt.name}_{get_time_str()}.log\")\n",
    "    logger = get_root_logger(\n",
    "        logger_name=\"basicsr\", log_level=logging.INFO, log_file=log_file\n",
    "    )\n",
    "    logger.info(get_env_info())\n",
    "    logger.info(dict2str(config))\n",
    "    resume_optimizers = resume_state[\"optimizers\"]\n",
    "    optimizer.load_state_dict(resume_optimizers)\n",
    "    logger.info(\n",
    "        f\"Resuming training from epoch: {resume_state['epoch']}, \"\n",
    "        f\"iter: {resume_state['iter']}.\"\n",
    "    )\n",
    "    start_epoch = resume_state[\"epoch\"]\n",
    "    current_iter = resume_state[\"iter\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5d7c58d-5754-4f0c-9c31-886c7a0858ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# copy the yml file to the experiment root\n",
    "print(\"-> copy the yml file to the experiment root\")\n",
    "os.chdir(\"T2I-Adapter\")\n",
    "# copy_opt_file(opt.config, experiments_root) # TODO: for now we copy yaml file by hand to experiments/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8406c37-6268-4262-b450-ef07ffe39169",
   "metadata": {},
   "outputs": [],
   "source": [
    "# training\n",
    "print(\"-> Training & validation\")\n",
    "logger.info(f\"Start training from epoch: {start_epoch}, iter: {current_iter}\")\n",
    "for epoch in range(start_epoch, opt.epochs):\n",
    "    logger.info(f\"Epoch: {epoch}\")\n",
    "    train_dataloader.sampler.set_epoch(epoch)\n",
    "    # train\n",
    "    for _, data in enumerate(train_dataloader):\n",
    "        current_iter += 1\n",
    "        with torch.no_grad():\n",
    "            edge = net_G(data[\"im\"].cuda(non_blocking=True))[-1]\n",
    "            edge = edge > 0.5\n",
    "            edge = edge.float()\n",
    "            c = model.module.get_learned_conditioning(data[\"sentence\"]) \n",
    "            z = model.module.encode_first_stage( \n",
    "                (data[\"im\"] * 2 - 1.0).cuda(non_blocking=True)\n",
    "            )\n",
    "            z = model.module.get_first_stage_encoding(z) \n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        model.zero_grad()\n",
    "        features_adapter = model_ad(edge)\n",
    "        l_pixel, loss_dict = model(z, c=c, features_adapter=features_adapter)\n",
    "        # l_pixel.backward() #FIXME: takes too much space\n",
    "        optimizer.step()\n",
    "\n",
    "        if (current_iter + 1) % opt.print_fq == 0:\n",
    "            logger.info(loss_dict)\n",
    "\n",
    "        # save checkpoint\n",
    "        rank, _ = get_dist_info()\n",
    "        if (rank == 0) and (\n",
    "            (current_iter + 1) % float(config[\"training\"][\"save_freq\"]) == 0\n",
    "        ):\n",
    "            save_filename = f\"model_ad_{current_iter+1}.pth\"\n",
    "            save_path = os.path.join(experiments_root, \"models\", save_filename)\n",
    "            save_dict = {}\n",
    "            model_ad_bare = get_bare_model(model_ad)\n",
    "            state_dict = model_ad_bare.state_dict()\n",
    "            for key, param in state_dict.items():\n",
    "                if key.startswith(\"module.\"):  # remove unnecessary 'module.'\n",
    "                    key = key[7:]\n",
    "                save_dict[key] = param.cpu()\n",
    "            torch.save(save_dict, save_path)\n",
    "            # save state\n",
    "            state = {\n",
    "                \"epoch\": epoch,\n",
    "                \"iter\": current_iter + 1,\n",
    "                \"optimizers\": optimizer.state_dict(),\n",
    "            }\n",
    "            save_filename = f\"{current_iter+1}.state\"\n",
    "            save_path = os.path.join(\n",
    "                experiments_root, \"training_states\", save_filename\n",
    "            )\n",
    "            torch.save(state, save_path)\n",
    "\n",
    "    # val\n",
    "    rank, _ = get_dist_info()\n",
    "    if rank == 0:\n",
    "        for data in val_dataloader:\n",
    "            with torch.no_grad():\n",
    "                if opt.dpm_solver:\n",
    "                    sampler = DPMSolverSampler(model.module)\n",
    "                elif opt.plms:\n",
    "                    sampler = PLMSSampler(model.module)\n",
    "                else:\n",
    "                    sampler = DDIMSampler(model.module)\n",
    "                print(data[\"im\"].shape)\n",
    "                c = model.module.get_learned_conditioning(data[\"sentence\"])\n",
    "                edge = net_G(data[\"im\"].cuda(non_blocking=True))[-1]\n",
    "                edge = edge > 0.5\n",
    "                edge = edge.float()\n",
    "                im_edge = tensor2img(edge)\n",
    "                cv2.imwrite(\n",
    "                    os.path.join(\n",
    "                        experiments_root, \"visualization\", \"edge_%04d.png\" % epoch\n",
    "                    ),\n",
    "                    im_edge,\n",
    "                )\n",
    "                print(\n",
    "                    \"-> Edge saved in \",\n",
    "                    os.path.join(\n",
    "                        experiments_root, \"visualization\", \"edge_%04d.png\" % epoch\n",
    "                    ),\n",
    "                )\n",
    "                features_adapter = model_ad(edge)\n",
    "                shape = [opt.C, opt.H // opt.f, opt.W // opt.f]\n",
    "                samples_ddim, _ = sampler.sample(\n",
    "                    S=opt.ddim_steps,\n",
    "                    conditioning=c,\n",
    "                    batch_size=opt.n_samples,\n",
    "                    shape=shape,\n",
    "                    verbose=False,\n",
    "                    unconditional_guidance_scale=opt.scale,\n",
    "                    unconditional_conditioning=model.module.get_learned_conditioning(\n",
    "                        opt.n_samples * [\"\"]\n",
    "                    ),\n",
    "                    eta=opt.ddim_eta,\n",
    "                    x_T=None,\n",
    "                    features_adapter=features_adapter,\n",
    "                )\n",
    "                x_samples_ddim = model.module.decode_first_stage(samples_ddim)\n",
    "                x_samples_ddim = torch.clamp(\n",
    "                    (x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0\n",
    "                )\n",
    "                x_samples_ddim = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy()\n",
    "                for id_sample, x_sample in enumerate(x_samples_ddim):\n",
    "                    x_sample = 255.0 * x_sample\n",
    "                    img = x_sample.astype(np.uint8)\n",
    "                    img = cv2.putText(\n",
    "                        img.copy(),\n",
    "                        data[\"sentence\"][0],\n",
    "                        (10, 30),\n",
    "                        cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                        0.5,\n",
    "                        (0, 255, 0),\n",
    "                        2,\n",
    "                    )\n",
    "                    cv2.imwrite(\n",
    "                        os.path.join(\n",
    "                            experiments_root,\n",
    "                            \"visualization\",\n",
    "                            \"sample_e%04d_s%04d.png\" % (epoch, id_sample),\n",
    "                        ),\n",
    "                        img[:, :, ::-1],\n",
    "                    )\n",
    "                    print(\"-> Image generated in:\", os.path.join(\n",
    "                            experiments_root,\n",
    "                            \"visualization\",\n",
    "                            \"sample_e%04d_s%04d.png\" % (epoch, id_sample),\n",
    "                        ))\n",
    "            break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e620629-9f39-4f06-85ba-fe8df627127d",
   "metadata": {},
   "source": [
    "### train_seg.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8eed9b3-e428-4c5a-98b9-11cd34eff289",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!torchrun T2I-Adapter/train_seg.py --local_rank 1 --bsize 1 --epochs 1 --num_workers 1 --n_samples 1 --scale 1 --gpus 1 --H 256 --W 256"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efaf8c37-6b60-4037-874e-e9a0f976c849",
   "metadata": {},
   "source": [
    "### train_depth.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4c8549-adff-4842-97da-b014c61bdff0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!torchrun T2I-Adapter/train_depth.py --local_rank 1 --bsize 1 --epochs 1 --num_workers 1 --n_samples 1 --scale 1 --gpus 1 --H 256 --W 256"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f877134-7cf7-4a8b-95d0-6bda0343881b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## 2 Conditions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6c17b7f-8d4b-49b3-9a44-fd2da26c5d86",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train.py -t --base configs/pl_train/coadapter-v1-train.yaml --train True --no-test True --gpus 1 --gpu 1 --scale_lr False \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --adapter_finetune_from models/t2iadapter_canny_sd14v1.pth \\\n",
    "    --coadapter_finetune_from models/t2iadapter_depth_sd14v1.pth \n",
    "    #--name coadapter-v1-train --autoresume"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fb13415-243d-4e8a-ad5f-7363b2c54863",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 3 Conditions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d35a73ac-ab79-462d-be94-22347ab76c61",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef82e64f-a673-4550-8649-2d2fc3024583",
   "metadata": {},
   "source": [
    "BE ADVISED: I did not use open pose, key pose or style conditionings since these do not seem apropiate for the road-image generation task (see [here](https://github.com/TencentARC/T2I-Adapter/blob/SD/docs/examples.md#style-adapter)). Canny is also discarded due to its similarity to sketch conditioning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8470648-83cd-4723-a149-e27f0a731ca1",
   "metadata": {},
   "source": [
    "WARNING: First I remove all checkpoints to train them from scratch "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "a83042e6-dccb-4fd9-9db8-ee2674c894e0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf T2I-Adapter/logs/*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2baad8c6-e439-4d16-ad64-2bc1c5248780",
   "metadata": {},
   "source": [
    "Depth, Segmentation Mask, Sketch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "980148bd-41b7-4d66-9240-4059d6356539",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train.py -t --base configs/pl_train/coadapter-v1-train_depth_seg_sketch.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --adapter_finetune_from models/t2iadapter_seg_sd14v1.pth \\\n",
    "    --coadapter_finetune_from models/t2iadapter_depth_sd14v1.pth \\\n",
    "    --coadapter2_finetune_from models/t2iadapter_sketch_sd14v1.pth \\\n",
    "    --name coadapter-depth-seg-sketch \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e37cd7c-5372-4bba-8392-6bbf64585cb3",
   "metadata": {},
   "source": [
    "Color, Segmentation Mask, Sketch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c93ca48-d6a3-49c7-bae3-9bf14141c442",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train.py -t --base configs/pl_train/coadapter-v1-train_color_seg_sketch.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --adapter_finetune_from models/t2iadapter_seg_sd14v1.pth \\\n",
    "    --coadapter_finetune_from models/t2iadapter_color_sd14v1.pth \\\n",
    "    --coadapter2_finetune_from models/t2iadapter_sketch_sd14v1.pth \\\n",
    "    --name coadapter-color-seg-sketch \\\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d2d9388-0b0e-43d5-9258-c6eb3d499ac1",
   "metadata": {},
   "source": [
    "Color, Depth, Segmentation Mask "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a23095a3-5caf-4a81-a8a9-1499a25469e2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train.py -t --base configs/pl_train/coadapter-v1-train_color_depth_seg.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --adapter_finetune_from models/t2iadapter_seg_sd14v1.pth \\\n",
    "    --coadapter_finetune_from models/t2iadapter_depth_sd14v1.pth \\\n",
    "    --coadapter2_finetune_from models/t2iadapter_color_sd14v1.pth \\\n",
    "    --name coadapter-color-depth-seg \\\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b3f9705-8978-44dc-8d82-5abdab0fd5d2",
   "metadata": {
    "tags": []
   },
   "source": [
    "Color, Depth, Sketch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e6d69f-6643-4116-bd0f-b819f51f8640",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train.py -t --base configs/pl_train/coadapter-v1-train_color_depth_sketch.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr False \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --adapter_finetune_from models/t2iadapter_color_sd14v1.pth \\\n",
    "    --coadapter_finetune_from models/t2iadapter_depth_sd14v1.pth \\\n",
    "    --coadapter2_finetune_from models/t2iadapter_sketch_sd14v1.pth \\\n",
    "    --name coadapter-color-depth-sketch \\\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e91581-c79f-4b2b-9cbc-0f8ee0df83af",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Train 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b1c4017-39f9-4bbc-a0fd-1439cd4d4c0d",
   "metadata": {},
   "source": [
    "BE ADVISED: I did not use open pose, key pose or style conditionings since these do not seem apropiate for the road-image generation task (see [here](https://github.com/TencentARC/T2I-Adapter/blob/SD/docs/examples.md#style-adapter)). Canny is also discarded due to its similarity to sketch conditioning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8d5bb7d-4c3f-4f3e-90f4-937613afec57",
   "metadata": {},
   "source": [
    "WARNING: First I remove all checkpoints to train them from scratch "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "d4790b1b-5ed1-4861-9116-3e5679ccf4ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!rm -rf T2I-Adapter/logs/*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee608a22-1792-4074-a408-9fb5211db7a4",
   "metadata": {},
   "source": [
    "Color, depth, sketch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d232de98-827a-4fb8-9aa5-b7dac0fccdd4",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train2.py -t --base configs/pl_train/coadapter-v1-train_color_depth_sketch.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --coadapter_finetune_from models/coadapter-fuser-sd15v1.pth \\\n",
    "    --name coadapter-color-depth-sketch-short-prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad76280-cef0-43e3-be56-cea5c4554d7f",
   "metadata": {},
   "source": [
    "Depth, Segmentation Mask, Sketch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2863d393-e09e-4acb-9f35-4f2b5c18d3e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train2.py -t --base configs/pl_train/coadapter-v1-train_depth_seg_sketch.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --coadapter_finetune_from models/coadapter-fuser-sd15v1.pth \\\n",
    "    --name coadapter-depth-seg-sketch-long-prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "728e8849-154b-4f71-a00c-ae9325b2e741",
   "metadata": {},
   "source": [
    "Color, Segmentation Mask, Sketch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "892e91f4-37f4-4ef9-8ade-22f9bb96e51d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train2.py -t --base configs/pl_train/coadapter-v1-train_color_seg_sketch.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --coadapter_finetune_from models/coadapter-fuser-sd15v1.pth \\\n",
    "    --name coadapter-color-seg-sketch-short-prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76cdd7b4-dda8-4d38-a144-633db248a86a",
   "metadata": {},
   "source": [
    "Color, Depth, Segmentation Mask "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b756bd7-1503-46c5-a523-459c5651a9c4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/train2.py -t --base configs/pl_train/coadapter-v1-train_color_depth_seg.yaml --train True --test False --gpus 1 --gpu 0 --scale_lr True \\\n",
    "    --num_nodes 1 --sd_finetune_from models/sd-v1-4.ckpt \\\n",
    "    --coadapter_finetune_from models/coadapter-fuser-sd15v1.pth \\\n",
    "    --name coadapter-color-depth-seg-short-prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf4592cd-79ca-4a91-b969-8d88ade038c3",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Plot loss vs epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "id": "40c02779-3c20-446b-bef0-b6ceb0915851",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def plot_losses(csv_file):\n",
    "    # Read the CSV file into a pandas DataFrame, ignoring rows with missing values\n",
    "    df = pd.read_csv(csv_file).dropna(subset=['train/loss_simple_epoch', 'train/loss_vlb_epoch', 'train/loss_epoch', 'epoch'])\n",
    "\n",
    "    # Extract relevant columns\n",
    "    epochs = df['epoch']\n",
    "    loss_simple = df['train/loss_simple_epoch']\n",
    "    loss_vlb = df['train/loss_vlb_epoch']\n",
    "    loss = df['train/loss_epoch']\n",
    "\n",
    "    # Create figure and axes for subplots\n",
    "    fig, axes = plt.subplots(2, 1, figsize=(10, 10), sharex=True)\n",
    "\n",
    "    # Plot each type of loss in separate subplot\n",
    "    axes[0].plot(epochs, loss_simple, label='Loss Simple', color='b')\n",
    "    axes[0].set_ylabel('Loss Simple')\n",
    "    axes[0].legend()\n",
    "    axes[0].grid(True)\n",
    "\n",
    "    axes[1].plot(epochs, loss_vlb, label='Loss VLB', color='g')\n",
    "    axes[1].set_ylabel('Loss VLB')\n",
    "    axes[1].legend()\n",
    "    axes[1].grid(True)\n",
    "    \"\"\"\n",
    "    axes[2].plot(epochs, loss, label='Loss', color='r')\n",
    "    axes[2].set_ylabel('Loss')\n",
    "    axes[2].set_xlabel('Epoch')\n",
    "    axes[2].legend()\n",
    "    axes[2].grid(True)\n",
    "    \"\"\"\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dea3708-0a80-4b15-a40e-3cf58bc9095e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_losses(\"T2I-Adapter/logs/coadapter-color-depth-sketch-long-prompt/csvlogger/version_0/metrics.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb54a87-9da1-4e68-a841-023f36084bd5",
   "metadata": {},
   "source": [
    "Save to targz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e41c157d-df84-4e6e-be21-1df69956185b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!tar chvfz T2I-Adapter/logs/logs.tar.gz T2I-Adapter/logs/*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf395ee-7175-429a-a619-0ad66295c349",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b74ccc62-9ebe-4aaf-b204-81686ea78d79",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Script 1: uses SD (and adapters) v. 1.4 by default"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c5909f-6f29-4723-a3a1-40d46216f752",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# From SD 1.4 vainilla, without training\n",
    "#!python T2I-Adapter/app.py --sd_ckpt T2I-Adapter/models/sd-v1-4.ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f987f2b1-a59b-406f-a57a-26f22bfe388e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# With gradio\n",
    "#!python T2I-Adapter/app.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6688280-ece4-40d7-8f17-ff3cc7a9b62c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python T2I-Adapter/app_script.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "408934f2-a4ae-4072-aa07-d36713c2322e",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Script 2: uses SD (and adapters) v. 1.5 by default but works with SD 1.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af13e3e7-d044-447d-9bac-f57e7d041627",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# From SD 1.4 vainilla, without training\n",
    "#!python T2I-Adapter/app_coadapter.py --sd_ckpt T2I-Adapter/models/sd-v1-4.ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34ec82ae-fba0-4f41-83f4-1c9b45b16292",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# From training done above. Checkpoint saved in last.ckpt\n",
    "!python T2I-Adapter/app_coadapter.py --sd_ckpt T2I-Adapter/logs/coadapter-color-depth-sketch/checkpoints/last.ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57876a76-68fe-4884-a83b-44d86e276a17",
   "metadata": {
    "tags": []
   },
   "source": [
    "### See metrics and results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0550a36b-bf3b-459e-8e08-d4fa1aca17af",
   "metadata": {},
   "source": [
    "Stored in \"test/\" and \"test/metrics.csv\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a54255a5-f1d9-4c52-ab67-1744d5f65dc8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import ipyplot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9625108-0c25-4fe1-823c-106b5a34797b",
   "metadata": {},
   "source": [
    "### Long prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53df4773-44ad-40b4-869c-44f0bb08112d",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Conditionings results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "482226e2-106f-4d1f-8ee0-45e9b978bbd0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metrics = pd.read_csv(\"T2I-Adapter/dataset/test_long_prompt/metrics.csv\")\n",
    "labels = pd.read_csv(\"T2I-Adapter/dataset/train.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bcea58b-10c5-4c92-8f1f-de44b8ae8bd4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metrics.loc[:,[\"Conditionings\", \"MSE\",\"Perceptual_Loss\"]].groupby(\"Conditionings\").mean().sort_values(by=[\"MSE\", \"Perceptual_Loss\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "830cb727-74a9-4542-8954-f6fbc0e6dbc0",
   "metadata": {},
   "source": [
    "Best images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14f8b0e2-0935-4f65-b4d8-c9a8478027e0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metrics.loc[:,[\"Image\", \"MSE\",\"Perceptual_Loss\"]].groupby(\"Image\").mean().sort_values(by=[\"MSE\", \"Perceptual_Loss\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffbc5c1c-94bf-4094-b053-dd3118a08178",
   "metadata": {},
   "source": [
    "Best conditionings+weights for each metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea9848fb-6ba5-434e-9765-da82c21fe27a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metrics.iloc[metrics[\"MSE\"].argmin(),:], metrics.iloc[metrics[\"Perceptual_Loss\"].argmin(),:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af098f8-c7f7-4b2b-999f-62f218442ce5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "metrics.iloc[metrics[\"MSE\"].argmax(),:], metrics.iloc[metrics[\"Perceptual_Loss\"].argmax(),:]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03496731-b6a3-4ec6-a36e-deeaa47b7da6",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Best example for MSE and Perceptual Loss:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "90d84f69-1d7c-42a9-b2a0-147f951a8399",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_example(img, conds, weights, plot_conds=False):\n",
    "    images_list = [f'T2I-Adapter/dataset/train/{img}',\n",
    "                   f'T2I-Adapter/dataset/test/{conds}_{weights}/{img}_generated_0.png']\n",
    "    conds_list = [f'T2I-Adapter/dataset/test/{conds}_{weights}/{img}_{conds.split(\"_\")[0]}.png',\n",
    "                 f'T2I-Adapter/dataset/test/{conds}_{weights}/{img}_{conds.split(\"_\")[1]}.png',\n",
    "                 f'T2I-Adapter/dataset/test/{conds}_{weights}/{img}_{conds.split(\"_\")[2]}.png']\n",
    "\n",
    "    ipyplot.plot_images(images_list, max_images=20, img_width=500)\n",
    "    if plot_conds:\n",
    "        ipyplot.plot_images(conds_list, max_images=20, img_width=350)\n",
    "    print(labels[labels.Image == img].Labels.values)\n",
    "    print(metrics[(metrics.Image==img) & (metrics.Conditionings == conds) & (metrics.Weights == weights)].loc[:,[\"MSE\",\"Perceptual_Loss\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a789814-5572-44c3-915f-0efab21892f1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_example(\"00000500_img_front.jpg\", \"color_seg_sketch\", \"1.0_0.5_1.0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fcffddd-c9ee-4c3c-9265-b28169e41e63",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_example(\"00000360_img_front.jpg\", \"color_seg_sketch\", \"0.5_0.5_1.0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aad6e285-6469-44dd-a158-f21ac70650dd",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Best examples for best overall conditionings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd862c20-adeb-452f-be95-bbbad2d6aa18",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mse_min = metrics[metrics.Conditionings==\"color_depth_sketch\"][\"MSE\"].argmin()\n",
    "pl_min = metrics[metrics.Conditionings==\"color_depth_sketch\"][\"Perceptual_Loss\"].argmin()\n",
    "\n",
    "metrics[metrics.Conditionings==\"color_depth_sketch\"].iloc[mse_min,:], metrics[metrics.Conditionings==\"color_depth_sketch\"].iloc[pl_min,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d7b1dc5-3133-4cbb-83fb-87e444e05229",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_example(\"00000500_img_front.jpg\", \"color_depth_sketch\", \"1.0_0.5_0.5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b815cc01-5afd-471b-8843-55ca2fbc621b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plot_example(\"00000360_img_front.jpg\", \"color_depth_sketch\", \"1.0_1.0_0.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "390c3b51-740f-4607-b2ae-194f0856e878",
   "metadata": {},
   "source": [
    "#### Worse examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "005ed4aa-9ec6-44af-a708-22a79f826703",
   "metadata": {},
   "outputs": [],
   "source": [
    "mse_max = metrics[metrics.Conditionings==\"color_depth_sketch\"][\"MSE\"].argmax()\n",
    "pl_max = metrics[metrics.Conditionings==\"color_depth_sketch\"][\"Perceptual_Loss\"].argmax()\n",
    "\n",
    "metrics[metrics.Conditionings==\"color_depth_sketch\"].iloc[mse_max,:], metrics[metrics.Conditionings==\"color_depth_sketch\"].iloc[pl_min,:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4946e8b-9c00-4d19-a7fb-0688d0f5fb91",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_example(\"00000500_img_front.jpg\", \"color_depth_sketch\", \"1.0_0.5_0.5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e79b9f6-52d1-4c50-91b2-3187ca93bb6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_example(\"00000500_img_front.jpg\", \"color_depth_sketch\", \"1.0_0.5_0.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a72d53eb-115b-43e0-8bff-a4361424fbd8",
   "metadata": {},
   "source": [
    "#### Some examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db6e6bd7-466b-4c4f-a233-2ff471101316",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from numpy import random\n",
    "conditionings = metrics.loc[:,\"Conditionings\"].unique()\n",
    "weights = metrics.loc[:,\"Weights\"].unique()\n",
    "images = metrics.loc[:,\"Image\"].unique()\n",
    "\n",
    "\n",
    "NUM_IMAGES = 10\n",
    "\n",
    "for cond, w, img in zip(random.choice(conditionings, NUM_IMAGES), random.choice(weights, NUM_IMAGES), random.choice(images, NUM_IMAGES)):\n",
    "    plot_example(img, cond, w, True)"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}