demo_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import argparse\n",
    "import os\n",
    "import sys\n",
    "import yaml\n",
    "import datetime\n",
    "sys.path.append('/'.join(os.getcwd().split('/')[:-1]))\n",
    "from pipelines.sd_controlnet_rave import RAVE\n",
    "from pipelines.sd_multicontrolnet_rave import RAVE_MultiControlNet\n",
    "from IPython.display import Video, Image\n",
    "import PIL\n",
    "import utils.constants as const\n",
    "import utils.video_grid_utils as vgu\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import pprint \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def init_device():\n",
    "    device_name = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    device = torch.device(device_name)\n",
    "    return device\n",
    "\n",
    "def init_paths(input_ns):\n",
    "    if input_ns.save_folder == None or input_ns.save_folder == '':\n",
    "        input_ns.save_folder = input_ns.video_name.replace('.mp4', '').replace('.gif', '')\n",
    "    else:\n",
    "        input_ns.save_folder += f\"/{input_ns.video_name.replace('.mp4', '').replace('.gif', '')}\"\n",
    "    save_dir = f'{const.OUTPUT_PATH}/{input_ns.save_folder}'\n",
    "    os.makedirs(save_dir, exist_ok=True)\n",
    "    save_idx = max([int(x[-5:]) for x in os.listdir(save_dir)])+1 if os.listdir(save_dir) != [] else 0\n",
    "    input_ns.save_path = f'{save_dir}/{input_ns.positive_prompts}-{str(save_idx).zfill(5)}'\n",
    "    \n",
    "\n",
    "    input_ns.video_path = f'{const.MP4_PATH}/{input_ns.video_name}.mp4'\n",
    "    \n",
    "    if '-' in input_ns.preprocess_name:\n",
    "        input_ns.hf_cn_path = [const.PREPROCESSOR_DICT[i] for i in input_ns.preprocess_name.split('-')]\n",
    "    else:\n",
    "        input_ns.hf_cn_path = const.PREPROCESSOR_DICT[input_ns.preprocess_name]\n",
    "    input_ns.hf_path = \"runwayml/stable-diffusion-v1-5\"\n",
    "    \n",
    "    input_ns.inverse_path = f'{const.GENERATED_DATA_PATH}/inverses/{input_ns.video_name}/{input_ns.preprocess_name}_{input_ns.model_id}_{input_ns.grid_size}x{input_ns.grid_size}_{input_ns.pad}'\n",
    "    input_ns.control_path = f'{const.GENERATED_DATA_PATH}/controls/{input_ns.video_name}/{input_ns.preprocess_name}_{input_ns.grid_size}x{input_ns.grid_size}_{input_ns.pad}'\n",
    "    os.makedirs(input_ns.control_path, exist_ok=True)\n",
    "    os.makedirs(input_ns.inverse_path, exist_ok=True)\n",
    "    os.makedirs(input_ns.save_path, exist_ok=True)\n",
    "    return input_ns\n",
    "    \n",
    "def run(input_ns):\n",
    "\n",
    "    if 'model_id' not in list(input_ns.__dict__.keys()):\n",
    "        input_ns.model_id = \"None\"\n",
    "    device = init_device()\n",
    "    input_ns = init_paths(input_ns)\n",
    "\n",
    "    input_ns.image_pil_list = vgu.prepare_video_to_grid(input_ns.video_path, input_ns.sample_size, input_ns.grid_size, input_ns.pad)\n",
    "    print(input_ns.video_path )\n",
    "    input_ns.sample_size = len(input_ns.image_pil_list)\n",
    "    print(f'Frame count: {len(input_ns.image_pil_list)}')\n",
    "\n",
    "    controlnet_class = RAVE_MultiControlNet if '-' in str(input_ns.controlnet_conditioning_scale) else RAVE\n",
    "    \n",
    "\n",
    "    CN = controlnet_class(device)\n",
    "\n",
    "\n",
    "    CN.init_models(input_ns.hf_cn_path, input_ns.hf_path, input_ns.preprocess_name, input_ns.model_id)\n",
    "    \n",
    "    input_dict = vars(input_ns)\n",
    "    pp = pprint.PrettyPrinter(indent=4)\n",
    "    pp.pprint(input_dict)\n",
    "    yaml_dict = {k:v for k,v in input_dict.items() if k != 'image_pil_list'}\n",
    "\n",
    "    start_time = datetime.datetime.now()\n",
    "    if '-' in str(input_ns.controlnet_conditioning_scale):\n",
    "        res_vid, control_vid_1, control_vid_2 = CN(input_dict)\n",
    "    else: \n",
    "        res_vid, control_vid = CN(input_dict)\n",
    "    end_time = datetime.datetime.now()\n",
    "    save_name = f\"{'-'.join(input_ns.positive_prompts.split())}_cstart-{input_ns.controlnet_guidance_start}_gs-{input_ns.guidance_scale}_pre-{'-'.join((input_ns.preprocess_name.replace('-','+').split('_')))}_cscale-{input_ns.controlnet_conditioning_scale}_grid-{input_ns.grid_size}_pad-{input_ns.pad}_model-{input_ns.model_id.split('/')[-1]}\"\n",
    "    res_vid[0].save(f\"{input_ns.save_path}/{save_name}.gif\", save_all=True, append_images=res_vid[1:], optimize=False, loop=10000)\n",
    "    # control_vid[0].save(f\"{input_ns.save_path}/control_{save_name}.gif\", save_all=True, append_images=control_vid[1:], optimize=False, loop=10000)\n",
    "\n",
    "    yaml_dict['total_time'] = (end_time - start_time).total_seconds()\n",
    "    yaml_dict['total_number_of_frames'] = len(res_vid)\n",
    "    yaml_dict['sec_per_frame'] = yaml_dict['total_time']/yaml_dict['total_number_of_frames']\n",
    "    with open(f'{input_ns.save_path}/config.yaml', 'w') as yaml_file:\n",
    "        yaml.dump(yaml_dict, yaml_file)\n",
    "    \n",
    "    if '-' in str(input_ns.controlnet_conditioning_scale):\n",
    "        return (res_vid, control_vid_1, control_vid_2)\n",
    "    else: \n",
    "        return (res_vid, control_vid)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_ns = argparse.Namespace(**{})\n",
    "input_ns.video_name = \"truck\"  # specify the name of the video existing in the data/mp4_videos folder (e.g. data/mp4_videos/truck.mp4)\n",
    "input_ns.preprocess_name = 'depth_zoe'  # specify the name of the condition used for the video (e.g. depth_zoe condition)\n",
    "\n",
    "input_ns.batch_size = 4  # denotes the batch size of grids (e.g. 4 grids run in parallel)\n",
    "input_ns.batch_size_vae = 1  # denotes the batch size for the VAE (e.g. 1 grid runs in parallel for the VAE)\n",
    "\n",
    "input_ns.cond_step_start = 0.0  # denotes the step to start conditioning\n",
    "\n",
    "input_ns.controlnet_conditioning_scale = 1.0  # denotes the scale of the conditioning\n",
    "input_ns.controlnet_guidance_end = 1.0  # denotes the end of the controlnet guidance\n",
    "input_ns.controlnet_guidance_start = 0.0  # denotes the start of the controlnet guidance\n",
    "\n",
    "input_ns.give_control_inversion = True  # denotes whether to give control to the inversion\n",
    "\n",
    "input_ns.grid_size = 3  # denotes the size of each grid (e.g. grid_size x grid_size)\n",
    "input_ns.sample_size = -1  # denotes the number of grids to be generated (-1 for the full video)\n",
    "input_ns.pad = 1  # denotes the padding of the video (if 1, use the same video)\n",
    "input_ns.guidance_scale = 7.5  # denotes the scale of the guidance\n",
    "input_ns.inversion_prompt = ''  # denotes the inversion prompt\n",
    "\n",
    "input_ns.is_ddim_inversion = True  # denotes whether to use ddim for inversion\n",
    "input_ns.is_shuffle = True  # denotes whether to applying shuffling between the grids\n",
    "\n",
    "input_ns.negative_prompts = \"\"  # denotes the negative prompts\n",
    "input_ns.num_inference_steps = 50  # denotes the number of inference steps during the sampling process\n",
    "input_ns.num_inversion_step = 50  # denotes the number of inversion steps during the inversion process\n",
    "input_ns.positive_prompts = \"Wooden trucks drive on a racetrack\"  # denotes the positive prompts\n",
    "input_ns.save_folder = ''  # denotes the name of the folder to save the results under results\n",
    "\n",
    "input_ns.seed = 0  # denotes the seed\n",
    "input_ns.model_id = 'None'  # None to use stable diffusion v1.5, otherwise use the model id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run RAVE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = run(input_ns)\n",
    "save_dir_name = 'truck'\n",
    "save_dir = f'assets/notebook-generated/{save_dir_name}'\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "if len(res) == 3:\n",
    "    res_vid, control_vid_1, control_vid_2 = res\n",
    "    control_vid_1[0].save(f\"{save_dir}/control.gif\", save_all=True, append_images=control_vid_1[1:], loop=0)\n",
    "    control_vid_2[0].save(f\"{save_dir}/control_2.gif\", save_all=True, append_images=control_vid_2[1:], loop=0)\n",
    "    res_vid[0].save(f\"{save_dir}/result.gif\", save_all=True, append_images=res_vid[1:], loop=0)\n",
    "else:\n",
    "    res_vid, control_vid = res\n",
    "    control_vid[0].save(f\"{save_dir}/control.gif\", save_all=True, append_images=control_vid[1:], loop=0)\n",
    "    res_vid[0].save(f\"{save_dir}/result.gif\", save_all=True, append_images=res_vid[1:], loop=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Original Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Video(filename=input_ns.video_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Edited Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=f\"{save_dir}/result.gif\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename=f\"{save_dir}/control.gif\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rave",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}