inferencing_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cNHvQBhzyXCI",
    "outputId": "0a79e979-8484-4c62-96d9-7c79b1835162",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# necessary packages and repos - may have to restart runtime/kernel after running\n",
    "!apt update\n",
    "!apt install python3-opencv -y\n",
    "\n",
    "!git clone https://github.com/alembics/disco-diffusion.git\n",
    "!mv disco-diffusion/disco_xform_utils.py disco_xform_utils.py\n",
    "!git clone https://github.com/isl-org/MiDaS.git\n",
    "!git clone https://github.com/MSFTserver/pytorch3d-lite.git\n",
    "!mv MiDaS/utils.py MiDaS/midas_utils.py\n",
    "!git clone https://github.com/shariqfarooq123/AdaBins.git\n",
    "\n",
    "!git clone https://github.com/CompVis/latent-diffusion.git\n",
    "!git clone https://github.com/CompVis/taming-transformers\n",
    "!git clone https://github.com/crowsonkb/k-diffusion\n",
    "!git clone https://github.com/mlfoundations/open_clip.git\n",
    "!git clone https://github.com/assafshocher/ResizeRight.git\n",
    "!pip install -e ./taming-transformers\n",
    "!pip install omegaconf>=2.0.0 pytorch-lightning>=1.0.8 torch-fidelity einops\n",
    "!pip install git+https://github.com/openai/CLIP.git huggingface_hub lpips timm opencv-python matplotlib pandas\n",
    "!pip install transformers kornia imageio imageio_ffmpeg pillow scikit-image jsonmerge clean-fid resize-right torchdiffeq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SD weights download/setup\n",
    "import os\n",
    "from huggingface_hub import hf_hub_download\n",
    "auth_token = ''\n",
    "model_location = hf_hub_download(repo_id=\"CompVis/stable-diffusion-v-1-4-original\", filename=\"sd-v1-4.ckpt\", use_auth_token=auth_token)\n",
    "model_location = model_location.split('snapshots')[0]+'blobs/'+os.listdir(model_location.split('snapshots')[0]+'blobs/')[0]\n",
    "!mkdir models/ldm/stable-diffusion-v1\n",
    "!mv $model_location models/ldm/stable-diffusion-v1/model.ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MiDaS depth model - only needed if using disco style animations\n",
    "!mkdir models/depth/\n",
    "!mkdir models/depth/midas/\n",
    "!wget -O models/depth/midas/model.ckpt https://github.com/intel-isl/DPT/releases/download/1_0/dpt_large-midas-2f21e586.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# AdaBins depth model - only needed if using disco style animations\n",
    "!mkdir pretrained/\n",
    "!wget -O pretrained/AdaBins_nyu.pt https://cloudflare-ipfs.com/ipfs/Qmd2mMnDLWePKmgfS8m6ntAg4nhV5VkUyAydYBp8cWWeB7/AdaBins_nyu.pt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# imports\n",
    "import re\n",
    "import os\n",
    "import cv2\n",
    "import gc\n",
    "import time\n",
    "import math\n",
    "import torch\n",
    "import imageio\n",
    "import requests\n",
    "import torchvision.transforms as T\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from skimage.exposure import match_histograms\n",
    "from ipywidgets import Output\n",
    "from torch.nn import functional as F\n",
    "import torchvision.transforms.functional as TF\n",
    "from omegaconf import OmegaConf\n",
    "from types import SimpleNamespace\n",
    "from PIL import Image, ImageOps, ImageFilter\n",
    "from tqdm.auto import tqdm, trange\n",
    "from itertools import islice\n",
    "from einops import rearrange, repeat\n",
    "from torchvision.utils import make_grid\n",
    "\n",
    "from pytorch_lightning import seed_everything\n",
    "from torch import autocast\n",
    "from contextlib import contextmanager, nullcontext\n",
    "\n",
    "from ldm.util import instantiate_from_config\n",
    "from ldm.models.diffusion.ddim import DDIMSampler\n",
    "from ldm.models.diffusion.plms import PLMSSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split up imports to avoid weird error when run in one cell\n",
    "import sys\n",
    "\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}/k-diffusion')\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}/ResizeRight')\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}/open_clip/src')\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}')\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}/MiDaS')\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}/pytorch3d-lite')\n",
    "sys.path.append(f'{os.path.abspath(os.getcwd())}/AdaBins/')\n",
    "\n",
    "from resize_right import resize\n",
    "import open_clip\n",
    "import py3d_tools as p3dT\n",
    "from midas.dpt_depth import DPTDepthModel\n",
    "from midas.midas_net import MidasNet\n",
    "from midas.midas_net_custom import MidasNet_small\n",
    "from midas.transforms import Resize, NormalizeImage, PrepareForNet\n",
    "import k_diffusion as K\n",
    "import disco_xform_utils as dxf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# variables\n",
    "use_laion400m = False\n",
    "use_fp16 = False\n",
    "\n",
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "if use_laion400m:\n",
    "    print(\"Falling back to LAION 400M model...\")\n",
    "    config_location = \"configs/latent-diffusion/txt2img-1p4B-eval.yaml\"\n",
    "    model_location = \"models/ldm/text2img-large/model.ckpt\"\n",
    "    outdir = \"outputs/txt2img-samples-laion400m\"\n",
    "else:\n",
    "    print(\"Using Stable Diffusion model...\")\n",
    "    config_location = 'configs/stable-diffusion/v1-inference.yaml'\n",
    "    model_location = 'models/ldm/stable-diffusion-v1/model.ckpt'\n",
    "    outdir = 'outputs/txt2img-samples'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "fnGwQRhtyBhb",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# utility functions\n",
    "def load_model_from_config(config, ckpt):\n",
    "    print(f\"Loading model from {ckpt}\")\n",
    "    pl_sd = torch.load(ckpt)#, map_location=\"cpu\")\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 use_fp16:\n",
    "        model = model.half()\n",
    "    model = model.to(device)\n",
    "    model = model.eval()\n",
    "    return model\n",
    "\n",
    "\n",
    "def get_model():\n",
    "    config = OmegaConf.load(config_location)\n",
    "    model = load_model_from_config(config, model_location)\n",
    "    return model\n",
    "\n",
    "\n",
    "def chunk(it, size):\n",
    "    it = iter(it)\n",
    "    return iter(lambda: tuple(islice(it, size)), ())\n",
    "\n",
    "\n",
    "def load_img(path, size, mask=False, mask_offset_size=1.0, outcrop_animation=False, factor=0.5, show_generated_mask_and_image=False, blur_strength=0):\n",
    "    if path.startswith('http://') or path.startswith('https://'):\n",
    "        path = requests.get(path, stream=True).raw\n",
    "    if not outcrop_animation:\n",
    "        factor = 1.0\n",
    "        if mask:\n",
    "            image = Image.open(path).convert(\"L\")\n",
    "            if blur_strength > 0:\n",
    "                image = image.filter(ImageFilter.GaussianBlur(blur_strength))\n",
    "        else:\n",
    "            image = Image.open(path).convert(\"RGB\")\n",
    "        w_, h_ = image.size\n",
    "    elif outcrop_animation and not mask:\n",
    "        image = Image.open(path).convert(\"RGB\")\n",
    "        w_, h_ = image.size\n",
    "    w, h = map(lambda x: x - x % 32, (size[0], size[1]))  # resize to integer multiple of 32\n",
    "    if outcrop_animation:\n",
    "        w_scaled, h_scaled = map(lambda x: int(x*factor), (w, h))  # scale image\n",
    "    if outcrop_animation and mask:\n",
    "        assert 0.0 <= factor <= 1.0\n",
    "        w_scaled, h_scaled = map(lambda x: int(x*mask_offset_size), (w_scaled, h_scaled))\n",
    "        image = Image.new('L', (w_scaled, h_scaled), 'white')\n",
    "        padding = Image.new(image.mode, (w, h), 'black')\n",
    "        padding.paste(image, (int((w-w_scaled)/2), int((h-h_scaled)/2)))\n",
    "        if blur_strength > 0:\n",
    "            padding = padding.filter(ImageFilter.GaussianBlur(blur_strength))\n",
    "        image = padding\n",
    "    else:\n",
    "        print(f\"loaded input image of size ({w_}, {h_}) and resized to ({w}, {h}) from {path}\")\n",
    "        if outcrop_animation:\n",
    "            assert 0.0 <= factor <= 1.0\n",
    "            image = image.resize((w_scaled, h_scaled), resample=Image.LANCZOS)\n",
    "            padding = Image.new(image.mode, (w, h), 'black')\n",
    "            padding.paste(image, (int((w-w_scaled)/2), int((h-h_scaled)/2)))\n",
    "            image = padding\n",
    "        else:\n",
    "            image = image.resize((w, h), resample=Image.LANCZOS)\n",
    "    if show_generated_mask_and_image:\n",
    "        display(image)\n",
    "    image = np.array(image).astype(np.float32) / 255.0\n",
    "    if mask:\n",
    "        image = image[None,None]\n",
    "        # image[image < 0.5] = 0\n",
    "        # image[image >= 0.5] = 1\n",
    "        return torch.from_numpy(image)\n",
    "    else:\n",
    "        image = image[None].transpose(0, 3, 1, 2)\n",
    "    image = torch.from_numpy(image)\n",
    "    return 2.*image - 1.\n",
    "\n",
    "\n",
    "def slerp(val, low, high):\n",
    "    low_norm = low/torch.norm(low)\n",
    "    high_norm = high/torch.norm(high)\n",
    "    omega = torch.acos((low_norm*high_norm).sum())\n",
    "    so = torch.sin(omega)\n",
    "    res = (torch.sin((1.0-val)*omega)/so)*low + (torch.sin(val*omega)/so) * high\n",
    "    return res\n",
    "\n",
    "\n",
    "def get_slerp_vectors(start, end, frames=20):\n",
    "    factor = 1.0 / (frames - 1)\n",
    "    if isinstance(start, int):\n",
    "        use_lerp = True\n",
    "        out = []\n",
    "    else:\n",
    "        use_lerp = False\n",
    "        out = torch.Tensor(frames, start.shape[0]).to(device)\n",
    "    for i in range(frames):\n",
    "        if use_lerp:\n",
    "            out.append(int(start + (end - start) * factor*i))\n",
    "        else:\n",
    "            out[i] = slerp(factor*i, start, end)\n",
    "    return out\n",
    "\n",
    "\n",
    "def get_conditioning_vector(prompt):\n",
    "    split_prompt = re.split(r':([-\\d.]+)', prompt)\n",
    "    if len(split_prompt) == 0:\n",
    "        raise(AttributeError('not a valid prompt'))\n",
    "    elif len(split_prompt) == 1:\n",
    "        return model.get_learned_conditioning(prompt)\n",
    "\n",
    "    split_prompt = iter(split_prompt)\n",
    "    c_tensors = []\n",
    "    weights = []\n",
    "    for text in split_prompt:\n",
    "        if text == '':\n",
    "            continue\n",
    "        text = text.strip()\n",
    "        try:\n",
    "            weight = float(next(split_prompt))\n",
    "            weights.append(weight)\n",
    "            c_tensors.append(model.get_learned_conditioning(text))\n",
    "        except:\n",
    "            weights.append(0.5)\n",
    "            c_tensors.append(model.get_learned_conditioning(text)) \n",
    "            continue\n",
    "\n",
    "    original_c_shape = c_tensors[0].shape\n",
    "    c = c_tensors[0].flatten()\n",
    "    for i in range(1,len(c_tensors)):\n",
    "        weight = weights[i-1]\n",
    "        # slerp between subprompts by weight between 0-1\n",
    "        c = slerp(weight, c, c_tensors[i].flatten())\n",
    "    c = c.reshape(*original_c_shape)\n",
    "    return c\n",
    "\n",
    "\n",
    "def get_starting_code_and_conditioning_vector(seed, prompt, mode):\n",
    "    if seed is None:\n",
    "        seed = np.random.randint(np.iinfo(np.int32).max)\n",
    "    seed_everything(seed)\n",
    "    if mode == 'text_interpolation':\n",
    "        torch.manual_seed(seed)\n",
    "        model_wrap = K.external.CompVisDenoiser(model)\n",
    "        sigmas = model_wrap.get_sigmas(ddim_steps)\n",
    "        start_code = torch.randn([1, C, H // f, W // f], device=device) * sigmas[0]\n",
    "    else:\n",
    "        start_code = torch.randn([1, C, H // f, W // f], device=device)\n",
    "    c = get_conditioning_vector(prompt)\n",
    "    return (c, start_code, seed)\n",
    "\n",
    "\n",
    "def unflatten(l, n):\n",
    "    res = []\n",
    "    t = l[:]\n",
    "    while len(t) > 0:\n",
    "        res.append(t[:n])\n",
    "        t = t[n:]\n",
    "    return res\n",
    "\n",
    "\n",
    "def zoom_at(org_img, zoom, x, y):\n",
    "    w, h = org_img.size\n",
    "\n",
    "    xz = x / zoom\n",
    "    yz = y / zoom\n",
    "    box = (int(x - xz), int(y - yz), int(x - xz + w / zoom), int(y - yz + h / zoom))\n",
    "\n",
    "    img = org_img.crop(box)\n",
    "    img = img.resize((w, h), Image.LANCZOS)\n",
    "    return img\n",
    "\n",
    "\n",
    "def outcropping_frames_to_mp4(out_crop_dir, project_name, out_mp4, frame_count, frames, fps, factor, zoom_cropping):\n",
    "    log_scale = math.log(1/factor)\n",
    "    anim_size = None\n",
    "    with imageio.get_writer(out_mp4, mode='I', fps=fps) as writer:\n",
    "        for i in range(frame_count, 0, -1):\n",
    "            if i == 1 and original_image_location != '':\n",
    "                img = Image.open(original_image_location)\n",
    "                img = img.resize(anim_size, Image.LANCZOS)\n",
    "            else:\n",
    "                img = Image.open(f\"{out_crop_dir}/{project_name}_{str(i)}.png\")\n",
    "                anim_size = img.size\n",
    "            w, h = img.size\n",
    "            for frame in tqdm(range(frames)):\n",
    "                if i == 1 and frame == frames // 2:\n",
    "                    break\n",
    "                zoom = math.exp(log_scale * frame / (frames))\n",
    "                zoomed = zoom_at(img, zoom, w//2, h//2)\n",
    "                w, h = zoomed.size\n",
    "                if zoom_cropping != (0, 0):\n",
    "                    zoomed = zoomed.crop((0 + zoom_cropping[0], 0 + zoom_cropping[1], w - zoom_cropping[0], h - zoom_cropping[1]))\n",
    "                writer.append_data(np.array(zoomed))\n",
    "    print(f\"Outcropping animation is available at {out_mp4}\")\n",
    "\n",
    "\n",
    "def make_callback(sampler, dynamic_threshold=0, static_threshold=0, inpainting=False, mix_with_x0=False, mix_factor=[0.15, 0.30, 0.60, 1.0], x0=None, noise=None, mask=None):  \n",
    "    # Creates the callback function to be passed into the samplers\n",
    "    # The callback function is applied to the image after each step\n",
    "    def dynamic_thresholding_(img, threshold): # check over implementation to ensure correctness\n",
    "        # Dynamic thresholding from Imagen paper (May 2022)\n",
    "        s = np.percentile(np.abs(img.cpu()), threshold, axis=tuple(range(1,img.ndim)))\n",
    "        s = np.max(np.append(s,1.0))\n",
    "        # torch.clamp_(img, -1*s, s) # this causes images to become grey/brown - investigate\n",
    "        torch.FloatTensor.div_(img, s)\n",
    "\n",
    "    # Callback for samplers in the k-diffusion repo, called thus:\n",
    "    #   callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})\n",
    "    def k_callback(args_dict):\n",
    "        if static_threshold != 0 and static_threshold is not None:\n",
    "            torch.clamp_(args_dict['x'], -1*static_threshold, static_threshold)\n",
    "        if dynamic_threshold != 0 and dynamic_threshold is not None:\n",
    "            dynamic_thresholding_(args_dict['x'], dynamic_threshold)\n",
    "        if inpainting and x0 is not None and mask is not None and noise is not None:\n",
    "            x = x0 + noise * args_dict['sigma']\n",
    "            x = x * mask\n",
    "            torch.FloatTensor.add_(torch.FloatTensor.mul_(args_dict['x'], (1. - mask)), x)\n",
    "        if mix_with_x0 and x0 is not None and noise is not None:\n",
    "            x = x0 + noise * args_dict['sigma']\n",
    "            try:\n",
    "                factor = min(mix_factor[min(args_dict['i'], len(mix_factor)-1)], 1.0)\n",
    "            except KeyError:\n",
    "                factor = min(mix_factor.values[-1], 1.0)\n",
    "            torch.FloatTensor.add_(torch.FloatTensor.mul_(args_dict['x'], factor), x * (1.0 - factor))\n",
    "\n",
    "    # Function that is called on the image (img) and step (i) at each step\n",
    "    def img_callback(img, i):\n",
    "        # Thresholding functions\n",
    "        if dynamic_threshold != 0:\n",
    "            dynamic_thresholding_(img, dynamic_threshold)\n",
    "        if static_threshold != 0:\n",
    "            torch.clamp_(img, -1*static_threshold, static_threshold)\n",
    "\n",
    "    if sampler in [\"PLMS\",\"DDIM\"]: \n",
    "        # Callback function formated for compvis latent diffusion samplers\n",
    "        callback = img_callback\n",
    "    else: \n",
    "        # Default callback function uses k-diffusion sampler variables\n",
    "        callback = k_callback\n",
    "\n",
    "    return callback\n",
    "\n",
    "\n",
    "class CFGDenoiser(nn.Module):\n",
    "        def __init__(self, model):\n",
    "            super().__init__()\n",
    "            self.inner_model = model\n",
    "\n",
    "        def forward(self, x, sigma, uncond, cond, cond_scale):\n",
    "            x_in = torch.cat([x] * 2)\n",
    "            sigma_in = torch.cat([sigma] * 2)\n",
    "            cond_in = torch.cat([uncond, cond])\n",
    "            uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)\n",
    "            return uncond + (cond - uncond) * cond_scale\n",
    "\n",
    "\n",
    "def maintain_colors(prev_img, color_match_sample, mode):\n",
    "    if mode == 'HSV':\n",
    "        prev_img_hsv = cv2.cvtColor(prev_img, cv2.COLOR_RGB2HSV)\n",
    "        color_match_hsv = cv2.cvtColor(color_match_sample, cv2.COLOR_RGB2HSV)\n",
    "        matched_hsv = match_histograms(prev_img_hsv, color_match_hsv, multichannel=True)\n",
    "        return cv2.cvtColor(matched_hsv, cv2.COLOR_HSV2RGB)\n",
    "    elif mode == 'RGB':\n",
    "        return match_histograms(prev_img, color_match_sample, multichannel=True)\n",
    "    else:\n",
    "        prev_img_lab = cv2.cvtColor(prev_img, cv2.COLOR_RGB2LAB)\n",
    "        color_match_lab = cv2.cvtColor(color_match_sample, cv2.COLOR_RGB2LAB)\n",
    "        matched_lab = match_histograms(prev_img_lab, color_match_lab, multichannel=True)\n",
    "        return cv2.cvtColor(matched_lab, cv2.COLOR_LAB2RGB)\n",
    "\n",
    "\n",
    "def add_noise(sample, noise_amt):\n",
    "    return sample + torch.randn(sample.shape, device=sample.device) * noise_amt\n",
    "\n",
    "\n",
    "def sample_from_cv2(sample):\n",
    "    sample = ((sample.astype(float) / 255.0) * 2) - 1\n",
    "    sample = sample[None].transpose(0, 3, 1, 2).astype(np.float16)\n",
    "    sample = torch.from_numpy(sample)\n",
    "    return sample\n",
    "\n",
    "\n",
    "def sample_to_cv2(sample):\n",
    "    sample_f32 = rearrange(sample.squeeze().cpu().numpy(), \"c h w -> h w c\").astype(np.float32)\n",
    "    sample_f32 = ((sample_f32 * 0.5) + 0.5).clip(0, 1)\n",
    "    sample_int8 = (sample_f32 * 255).astype(np.uint8)\n",
    "    return sample_int8\n",
    "\n",
    "\n",
    "def sample_from_pil(sample):\n",
    "    sample = ((np.asarray(sample).astype(float) / 255.0) * 2) - 1\n",
    "    sample = torch.from_numpy(sample)\n",
    "    return sample\n",
    "\n",
    "\n",
    "def sample_to_pil(sample):\n",
    "    sample_f32 = sample.squeeze().cpu().numpy().astype(np.float32)\n",
    "    sample_f32 = ((sample_f32 * 0.5) + 0.5).clip(0, 1)\n",
    "    sample_int8 = (sample_f32 * 255).astype(np.uint8)\n",
    "    return Image.fromarray(sample_int8)\n",
    "\n",
    "\n",
    "####################################\n",
    "# disco style animations utility functions\n",
    "def init_midas_depth_model():\n",
    "    midas_model = None\n",
    "    net_w = None\n",
    "    net_h = None\n",
    "    resize_mode = None\n",
    "    normalization = None\n",
    "\n",
    "    print(f\"Initializing MiDaS depth model...\")\n",
    "    # load network\n",
    "    midas_model_path = 'models/depth/midas/model.ckpt'\n",
    "    midas_model = DPTDepthModel(\n",
    "        path=midas_model_path,\n",
    "        backbone=\"vitl16_384\",\n",
    "        non_negative=True,\n",
    "    )\n",
    "    net_w, net_h = 384, 384\n",
    "    resize_mode = \"minimal\"\n",
    "    normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n",
    "\n",
    "    midas_transform = T.Compose(\n",
    "        [\n",
    "            Resize(\n",
    "                net_w,\n",
    "                net_h,\n",
    "                resize_target=None,\n",
    "                keep_aspect_ratio=True,\n",
    "                ensure_multiple_of=32,\n",
    "                resize_method=resize_mode,\n",
    "                image_interpolation_method=cv2.INTER_CUBIC,\n",
    "            ),\n",
    "            normalization,\n",
    "            PrepareForNet(),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    midas_model.eval()\n",
    "\n",
    "    midas_model = midas_model.to(memory_format=torch.channels_last)  \n",
    "    midas_model = midas_model.half()\n",
    "\n",
    "    midas_model.to(device)\n",
    "\n",
    "    print(f\"MiDaS depth model initialized.\")\n",
    "    return midas_model, midas_transform, net_w, net_h, resize_mode, normalization\n",
    "\n",
    "\n",
    "def do_3d_step(img_filepath, frame_num, midas_model, midas_transform):\n",
    "    if args.key_frames:\n",
    "        translation_x = args.translation_x_series[frame_num]\n",
    "        translation_y = args.translation_y_series[frame_num]\n",
    "        translation_z = args.translation_z_series[frame_num]\n",
    "        rotation_3d_x = args.rotation_3d_x_series[frame_num]\n",
    "        rotation_3d_y = args.rotation_3d_y_series[frame_num]\n",
    "        rotation_3d_z = args.rotation_3d_z_series[frame_num]\n",
    "        print(\n",
    "            f'translation_x: {translation_x}',\n",
    "            f'translation_y: {translation_y}',\n",
    "            f'translation_z: {translation_z}',\n",
    "            f'rotation_3d_x: {rotation_3d_x}',\n",
    "            f'rotation_3d_y: {rotation_3d_y}',\n",
    "            f'rotation_3d_z: {rotation_3d_z}',\n",
    "        )\n",
    "\n",
    "    translate_xyz = [-translation_x*TRANSLATION_SCALE, translation_y*TRANSLATION_SCALE, -translation_z*TRANSLATION_SCALE]\n",
    "    rotate_xyz_degrees = [rotation_3d_x, rotation_3d_y, rotation_3d_z]\n",
    "    print('translation:',translate_xyz)\n",
    "    print('rotation:',rotate_xyz_degrees)\n",
    "    rotate_xyz = [math.radians(rotate_xyz_degrees[0]), math.radians(rotate_xyz_degrees[1]), math.radians(rotate_xyz_degrees[2])]\n",
    "    rot_mat = p3dT.euler_angles_to_matrix(torch.tensor(rotate_xyz, device=device), \"XYZ\").unsqueeze(0)\n",
    "    print(\"rot_mat: \" + str(rot_mat))\n",
    "    next_step_pil = dxf.transform_image_3d(img_filepath, midas_model, midas_transform, torch.device(\"cuda\"),\n",
    "                                          rot_mat, translate_xyz, args.near_plane, args.far_plane,\n",
    "                                          args.fov, padding_mode=args.padding_mode,\n",
    "                                          sampling_mode=args.sampling_mode, midas_weight=args.midas_weight)\n",
    "    return next_step_pil\n",
    "\n",
    "\n",
    "def interp(t):\n",
    "    return 3 * t**2 - 2 * t ** 3\n",
    "\n",
    "\n",
    "def perlin(width, height, scale=10):\n",
    "    gx, gy = torch.randn(2, width + 1, height + 1, 1, 1, device=device)\n",
    "    xs = torch.linspace(0, 1, scale + 1)[:-1, None].to(device)\n",
    "    ys = torch.linspace(0, 1, scale + 1)[None, :-1].to(device)\n",
    "    wx = 1 - interp(xs)\n",
    "    wy = 1 - interp(ys)\n",
    "    dots = 0\n",
    "    dots += wx * wy * (gx[:-1, :-1] * xs + gy[:-1, :-1] * ys)\n",
    "    dots += (1 - wx) * wy * (-gx[1:, :-1] * (1 - xs) + gy[1:, :-1] * ys)\n",
    "    dots += wx * (1 - wy) * (gx[:-1, 1:] * xs - gy[:-1, 1:] * (1 - ys))\n",
    "    dots += (1 - wx) * (1 - wy) * (-gx[1:, 1:] * (1 - xs) - gy[1:, 1:] * (1 - ys))\n",
    "    return dots.permute(0, 2, 1, 3).contiguous().view(width * scale, height * scale)\n",
    "\n",
    "\n",
    "def perlin_ms(octaves, width, height, grayscale):\n",
    "    out_array = [0.5] if grayscale else [0.5, 0.5, 0.5, 0.5]\n",
    "    # out_array = [0.0] if grayscale else [0.0, 0.0, 0.0]\n",
    "    for i in range(1 if grayscale else 4):\n",
    "        scale = 2 ** len(octaves)\n",
    "        oct_width = width\n",
    "        oct_height = height\n",
    "        for oct in octaves:\n",
    "            p = perlin(oct_width, oct_height, scale)\n",
    "            out_array[i] += p * oct\n",
    "            scale //= 2\n",
    "            oct_width *= 2\n",
    "            oct_height *= 2\n",
    "    return torch.cat(out_array)\n",
    "\n",
    "\n",
    "def create_perlin_noise(octaves=[1, 1, 1, 1], width=2, height=2, grayscale=True):\n",
    "    out = perlin_ms(octaves, width, height, grayscale)\n",
    "    if grayscale:\n",
    "        out = TF.resize(size=(args.W//args.f, args.H//args.f), img=out.unsqueeze(0))\n",
    "        out = TF.to_pil_image(out.clamp(0, 1)).convert('RGBA')\n",
    "    else:\n",
    "        out = out.reshape(-1, 4, out.shape[0]//4, out.shape[1])\n",
    "        out = TF.resize(size=(args.W//args.f, args.H//args.f), img=out)\n",
    "        out = TF.to_pil_image(out.clamp(0, 1).squeeze())\n",
    "    \n",
    "    return out\n",
    "\n",
    "\n",
    "def regen_perlin():\n",
    "    if args.perlin_mode == 'color':\n",
    "        init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
    "        init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, False)\n",
    "    elif args.perlin_mode == 'gray':\n",
    "        init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, True)\n",
    "        init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
    "    else:\n",
    "        init = create_perlin_noise([1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
    "        init2 = create_perlin_noise([1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
    "    display(init)\n",
    "    init = TF.to_tensor(init).add(TF.to_tensor(init2)).div(2).to(device).unsqueeze(0).mul(2).sub(1)\n",
    "    del init2\n",
    "    return init.expand(args.batch_size, -1, -1, -1)\n",
    "\n",
    "\n",
    "def parse_key_frames(string, prompt_parser=None):\n",
    "    \"\"\"Given a string representing frame numbers paired with parameter values at that frame,\n",
    "    return a dictionary with the frame numbers as keys and the parameter values as the values.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    string: string\n",
    "        Frame numbers paired with parameter values at that frame number, in the format\n",
    "        'framenumber1: (parametervalues1), framenumber2: (parametervalues2), ...'\n",
    "    prompt_parser: function or None, optional\n",
    "        If provided, prompt_parser will be applied to each string of parameter values.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    dict\n",
    "        Frame numbers as keys, parameter values at that frame number as values\n",
    "\n",
    "    Raises\n",
    "    ------\n",
    "    RuntimeError\n",
    "        If the input string does not match the expected format.\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    >>> parse_key_frames(\"10:(Apple: 1| Orange: 0), 20: (Apple: 0| Orange: 1| Peach: 1)\")\n",
    "    {10: 'Apple: 1| Orange: 0', 20: 'Apple: 0| Orange: 1| Peach: 1'}\n",
    "\n",
    "    >>> parse_key_frames(\"10:(Apple: 1| Orange: 0), 20: (Apple: 0| Orange: 1| Peach: 1)\", prompt_parser=lambda x: x.lower()))\n",
    "    {10: 'apple: 1| orange: 0', 20: 'apple: 0| orange: 1| peach: 1'}\n",
    "    \"\"\"\n",
    "    import re\n",
    "    pattern = r'((?P<frame>[0-9]+):[\\s]*[\\(](?P<param>[\\S\\s]*?)[\\)])'\n",
    "    frames = dict()\n",
    "    for match_object in re.finditer(pattern, string):\n",
    "        frame = int(match_object.groupdict()['frame'])\n",
    "        param = match_object.groupdict()['param']\n",
    "        if prompt_parser:\n",
    "            frames[frame] = prompt_parser(param)\n",
    "        else:\n",
    "            frames[frame] = param\n",
    "\n",
    "    if frames == {} and len(string) != 0:\n",
    "        raise RuntimeError('Key Frame string not correctly formatted')\n",
    "    return frames\n",
    "\n",
    "\n",
    "def get_inbetweens(key_frames, max_frames, integer=False):\n",
    "    \"\"\"Given a dict with frame numbers as keys and a parameter value as values,\n",
    "    return a pandas Series containing the value of the parameter at every frame from 0 to max_frames.\n",
    "    Any values not provided in the input dict are calculated by linear interpolation between\n",
    "    the values of the previous and next provided frames. If there is no previous provided frame, then\n",
    "    the value is equal to the value of the next provided frame, or if there is no next provided frame,\n",
    "    then the value is equal to the value of the previous provided frame. If no frames are provided,\n",
    "    all frame values are NaN.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    key_frames: dict\n",
    "        A dict with integer frame numbers as keys and numerical values of a particular parameter as values.\n",
    "    integer: Bool, optional\n",
    "        If True, the values of the output series are converted to integers.\n",
    "        Otherwise, the values are floats.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    pd.Series\n",
    "        A Series with length max_frames representing the parameter values for each frame.\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    >>> max_frames = 5\n",
    "    >>> get_inbetweens({1: 5, 3: 6})\n",
    "    0    5.0\n",
    "    1    5.0\n",
    "    2    5.5\n",
    "    3    6.0\n",
    "    4    6.0\n",
    "    dtype: float64\n",
    "\n",
    "    >>> get_inbetweens({1: 5, 3: 6}, integer=True)\n",
    "    0    5\n",
    "    1    5\n",
    "    2    5\n",
    "    3    6\n",
    "    4    6\n",
    "    dtype: int64\n",
    "    \"\"\"\n",
    "    key_frame_series = pd.Series([np.nan for a in range(max_frames)])\n",
    "\n",
    "    for i, value in key_frames.items():\n",
    "        key_frame_series[i] = value\n",
    "    key_frame_series = key_frame_series.astype(float)\n",
    "\n",
    "    interp_method = interp_spline\n",
    "\n",
    "    if interp_method == 'Cubic' and len(key_frames.items()) <= 3:\n",
    "        interp_method = 'Quadratic'\n",
    "\n",
    "    if interp_method == 'Quadratic' and len(key_frames.items()) <= 2:\n",
    "        interp_method = 'Linear'\n",
    "\n",
    "    key_frame_series[0] = key_frame_series[key_frame_series.first_valid_index()]\n",
    "    key_frame_series[max_frames -\n",
    "                     1] = key_frame_series[key_frame_series.last_valid_index()]\n",
    "    # key_frame_series = key_frame_series.interpolate(method=intrp_method,order=1, limit_direction='both')\n",
    "    key_frame_series = key_frame_series.interpolate(\n",
    "        method=interp_method.lower(), limit_direction='both')\n",
    "    if integer:\n",
    "        return key_frame_series.astype(int)\n",
    "    return key_frame_series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EXPERIMENTAL - optional - run if using score corrector (make sure to set any clip models you wish to use to True) (score corrector not valid for k_ samplers)\n",
    "import clip\n",
    "import lpips\n",
    "# use_secondary_model = False\n",
    "\n",
    "use_checkpoint = True #@param {type: 'boolean'}\n",
    "ViTB32 = False #@param{type:\"boolean\"}\n",
    "ViTB16 = False #@param{type:\"boolean\"}\n",
    "ViTL14 = False #@param{type:\"boolean\"}\n",
    "ViTL14_336px = False #@param{type:\"boolean\"}\n",
    "RN101 = False #@param{type:\"boolean\"}\n",
    "RN50 = False #@param{type:\"boolean\"}\n",
    "RN50x4 = False #@param{type:\"boolean\"}\n",
    "RN50x16 = False #@param{type:\"boolean\"}\n",
    "RN50x64 = False #@param{type:\"boolean\"}\n",
    "\n",
    "ViTB32_laion2b_e16 = False #@param{type:\"boolean\"}\n",
    "ViTB32_laion400m_e31 = False #@param{type:\"boolean\"}\n",
    "ViTB32_laion400m_32 = False #@param{type:\"boolean\"}\n",
    "ViTB32quickgelu_laion400m_e31 = False #@param{type:\"boolean\"}\n",
    "ViTB32quickgelu_laion400m_e32 = False #@param{type:\"boolean\"}\n",
    "ViTB16_laion400m_e31 = False #@param{type:\"boolean\"}\n",
    "ViTB16_laion400m_e32 = False #@param{type:\"boolean\"}\n",
    "RN50_yffcc15m = False #@param{type:\"boolean\"}\n",
    "RN50_cc12m = False #@param{type:\"boolean\"}\n",
    "RN50_quickgelu_yfcc15m = False #@param{type:\"boolean\"}\n",
    "RN50_quickgelu_cc12m = False #@param{type:\"boolean\"}\n",
    "RN101_yfcc15m = False #@param{type:\"boolean\"}\n",
    "RN101_quickgelu_yfcc15m = False #@param{type:\"boolean\"}\n",
    "\n",
    "# if use_secondary_model:\n",
    "#     secondary_model = SecondaryDiffusionImageNet2()\n",
    "#     secondary_model.load_state_dict(torch.load(f'{model_path}/secondary_model_imagenet_2.pth', map_location='cpu'))\n",
    "#     secondary_model.eval().requires_grad_(False).to(device)\n",
    "\n",
    "clip_models = []\n",
    "if ViTB32: clip_models.append(clip.load('ViT-B/32', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if ViTB16: clip_models.append(clip.load('ViT-B/16', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if ViTL14: clip_models.append(clip.load('ViT-L/14', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if ViTL14_336px: clip_models.append(clip.load('ViT-L/14@336px', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if RN50: clip_models.append(clip.load('RN50', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if RN50x4: clip_models.append(clip.load('RN50x4', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if RN50x16: clip_models.append(clip.load('RN50x16', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if RN50x64: clip_models.append(clip.load('RN50x64', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if RN101: clip_models.append(clip.load('RN101', jit=False)[0].eval().requires_grad_(False).to(device))\n",
    "if ViTB32_laion2b_e16: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion2b_e16').eval().requires_grad_(False).to(device))\n",
    "if ViTB32_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion400m_e31').eval().requires_grad_(False).to(device))\n",
    "if ViTB32_laion400m_32: clip_models.append(open_clip.create_model('ViT-B-32', pretrained='laion400m_e32').eval().requires_grad_(False).to(device))\n",
    "if ViTB32quickgelu_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-32-quickgelu', pretrained='laion400m_e31').eval().requires_grad_(False).to(device))\n",
    "if ViTB32quickgelu_laion400m_e32: clip_models.append(open_clip.create_model('ViT-B-32-quickgelu', pretrained='laion400m_e32').eval().requires_grad_(False).to(device))\n",
    "if ViTB16_laion400m_e31: clip_models.append(open_clip.create_model('ViT-B-16', pretrained='laion400m_e31').eval().requires_grad_(False).to(device))\n",
    "if ViTB16_laion400m_e32: clip_models.append(open_clip.create_model('ViT-B-16', pretrained='laion400m_e32').eval().requires_grad_(False).to(device))\n",
    "if RN50_yffcc15m: clip_models.append(open_clip.create_model('RN50', pretrained='yfcc15m').eval().requires_grad_(False).to(device))\n",
    "if RN50_cc12m: clip_models.append(open_clip.create_model('RN50', pretrained='cc12m').eval().requires_grad_(False).to(device))\n",
    "if RN50_quickgelu_yfcc15m: clip_models.append(open_clip.create_model('RN50-quickgelu', pretrained='yfcc15m').eval().requires_grad_(False).to(device))\n",
    "if RN50_quickgelu_cc12m: clip_models.append(open_clip.create_model('RN50-quickgelu', pretrained='cc12m').eval().requires_grad_(False).to(device))\n",
    "if RN101_yfcc15m: clip_models.append(open_clip.create_model('RN101', pretrained='yfcc15m').eval().requires_grad_(False).to(device))\n",
    "if RN101_quickgelu_yfcc15m: clip_models.append(open_clip.create_model('RN101-quickgelu', pretrained='yfcc15m').eval().requires_grad_(False).to(device))\n",
    "\n",
    "normalize = T.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])\n",
    "lpips_model = lpips.LPIPS(net='vgg').to(device)\n",
    "\n",
    "loss_values = []\n",
    "\n",
    "def parse_prompt(prompt):\n",
    "    if prompt.startswith('http://') or prompt.startswith('https://'):\n",
    "        vals = prompt.rsplit(':', 2)\n",
    "        vals = [vals[0] + ':' + vals[1], *vals[2:]]\n",
    "    else:\n",
    "        vals = prompt.rsplit(':', 1)\n",
    "    vals = vals + ['', '1'][len(vals):]\n",
    "    return vals[0], float(vals[1])\n",
    "\n",
    "def get_clip_embeds(prompts, seed):\n",
    "    seed_everything(seed)\n",
    "    target_embeds, weights = [], []\n",
    "    model_stats = []\n",
    "    for clip_model in clip_models:\n",
    "        cutn = 16\n",
    "        model_stat = {\"clip_model\": None, \"target_embeds\": [],\n",
    "                    \"make_cutouts\": None, \"weights\": []}\n",
    "        model_stat[\"clip_model\"] = clip_model\n",
    "\n",
    "        for prompt in prompts:\n",
    "            txt, weight = parse_prompt(prompt)\n",
    "            txt = clip_model.encode_text(clip.tokenize(prompt).to(device)).float()\n",
    "\n",
    "            if fuzzy_prompt:\n",
    "                for i in range(25):\n",
    "                    model_stat[\"target_embeds\"].append(\n",
    "                        (txt + torch.randn(txt.shape).cuda() * args.rand_mag).clamp(0, 1))\n",
    "                    model_stat[\"weights\"].append(weight)\n",
    "            else:\n",
    "                model_stat[\"target_embeds\"].append(txt)\n",
    "                model_stat[\"weights\"].append(weight)\n",
    "\n",
    "        if image_prompt:\n",
    "            model_stat[\"make_cutouts\"] = MakeCutouts(\n",
    "                clip_model.visual.input_resolution, cutn, skip_augs=skip_augs)\n",
    "            for prompt in image_prompt:\n",
    "                path, weight = parse_prompt(prompt)\n",
    "                img = Image.open(fetch(path)).convert('RGB')\n",
    "                img = TF.resize(img, min(side_x, side_y, *img.size),\n",
    "                                T.InterpolationMode.LANCZOS)\n",
    "                batch = model_stat[\"make_cutouts\"](TF.to_tensor(\n",
    "                    img).to(device).unsqueeze(0).mul(2).sub(1))\n",
    "                embed = clip_model.encode_image(normalize(batch)).float()\n",
    "                if fuzzy_prompt:\n",
    "                    for i in range(25):\n",
    "                        model_stat[\"target_embeds\"].append(\n",
    "                            (embed + torch.randn(embed.shape).cuda() * rand_mag).clamp(0, 1))\n",
    "                        weights.extend([weight / cutn] * cutn)\n",
    "                else:\n",
    "                    model_stat[\"target_embeds\"].append(embed)\n",
    "                    model_stat[\"weights\"].extend([weight / cutn] * cutn)\n",
    "\n",
    "        model_stat[\"target_embeds\"] = torch.cat(model_stat[\"target_embeds\"])\n",
    "        model_stat[\"weights\"] = torch.tensor(model_stat[\"weights\"], device=device)\n",
    "        if model_stat[\"weights\"].sum().abs() < 1e-3:\n",
    "            raise RuntimeError('The weights must not sum to 0.')\n",
    "        model_stat[\"weights\"] /= model_stat[\"weights\"].sum().abs()\n",
    "        model_stats.append(model_stat)\n",
    "    return model_stats\n",
    "\n",
    "#############################################################\n",
    "\n",
    "cutout_debug = False\n",
    "padargs = {}\n",
    "class MakeCutoutsDango(nn.Module):\n",
    "    def __init__(self, cut_size,\n",
    "                 Overview=4,\n",
    "                 InnerCrop=0, IC_Size_Pow=0.5, IC_Grey_P=0.2\n",
    "                 ):\n",
    "        super().__init__()\n",
    "        self.cut_size = cut_size\n",
    "        self.Overview = Overview\n",
    "        self.InnerCrop = InnerCrop\n",
    "        self.IC_Size_Pow = IC_Size_Pow\n",
    "        self.IC_Grey_P = IC_Grey_P\n",
    "        if animation_mode == 'None':\n",
    "            self.augs = T.Compose([\n",
    "                T.RandomHorizontalFlip(p=0.5),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomAffine(degrees=10, translate=(0.05, 0.05),\n",
    "                               interpolation=T.InterpolationMode.BILINEAR),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomGrayscale(p=0.1),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.ColorJitter(brightness=0.1, contrast=0.1,\n",
    "                              saturation=0.1, hue=0.1),\n",
    "            ])\n",
    "        elif animation_mode == 'Video Input':\n",
    "            self.augs = T.Compose([\n",
    "                T.RandomHorizontalFlip(p=0.5),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomAffine(degrees=15, translate=(0.1, 0.1)),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomPerspective(distortion_scale=0.4, p=0.7),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomGrayscale(p=0.15),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                # T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1),\n",
    "            ])\n",
    "        elif animation_mode == '2D' or animation_mode == '3D':\n",
    "            self.augs = T.Compose([\n",
    "                T.RandomHorizontalFlip(p=0.4),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomAffine(degrees=10, translate=(0.05, 0.05),\n",
    "                               interpolation=T.InterpolationMode.BILINEAR),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.RandomGrayscale(p=0.1),\n",
    "                T.Lambda(lambda x: x + torch.randn_like(x) * 0.01),\n",
    "                T.ColorJitter(brightness=0.1, contrast=0.1,\n",
    "                              saturation=0.1, hue=0.3),\n",
    "            ])\n",
    "            \n",
    "    def forward(self, input):\n",
    "        cutouts = []\n",
    "        gray = T.Grayscale(3)\n",
    "        sideY, sideX = input.shape[2:4]\n",
    "        max_size = min(sideX, sideY)\n",
    "        min_size = min(sideX, sideY, self.cut_size)\n",
    "        l_size = max(sideX, sideY)\n",
    "        output_shape = [1,3,self.cut_size,self.cut_size] \n",
    "        output_shape_2 = [1,3,self.cut_size+2,self.cut_size+2]\n",
    "        pad_input = F.pad(input,((sideY-max_size)//2,(sideY-max_size)//2,(sideX-max_size)//2,(sideX-max_size)//2), **padargs)\n",
    "        cutout = resize(pad_input, out_shape=output_shape)\n",
    "\n",
    "        if self.Overview>0:\n",
    "            if self.Overview<=4:\n",
    "                if self.Overview>=1:\n",
    "                    cutouts.append(cutout)\n",
    "                if self.Overview>=2:\n",
    "                    cutouts.append(gray(cutout))\n",
    "                if self.Overview>=3:\n",
    "                    cutouts.append(TF.hflip(cutout))\n",
    "                if self.Overview==4:\n",
    "                    cutouts.append(gray(TF.hflip(cutout)))\n",
    "            else:\n",
    "                cutout = resize(pad_input, out_shape=output_shape)\n",
    "                for _ in range(self.Overview):\n",
    "                    cutouts.append(cutout)\n",
    "\n",
    "            if cutout_debug:\n",
    "                if is_colab:\n",
    "                    TF.to_pil_image(cutouts[0].clamp(0, 1).squeeze(0)).save(\"/content/cutout_overview0.jpg\",quality=99)\n",
    "                else:\n",
    "                    TF.to_pil_image(cutouts[0].clamp(0, 1).squeeze(0)).save(\"cutout_overview0.jpg\",quality=99)\n",
    "\n",
    "                              \n",
    "        if self.InnerCrop >0:\n",
    "            for i in range(self.InnerCrop):\n",
    "                size = int(torch.rand([])**self.IC_Size_Pow * (max_size - min_size) + min_size)\n",
    "                offsetx = torch.randint(0, sideX - size + 1, ())\n",
    "                offsety = torch.randint(0, sideY - size + 1, ())\n",
    "                cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
    "                if i <= int(self.IC_Grey_P * self.InnerCrop):\n",
    "                    cutout = gray(cutout)\n",
    "                cutout = resize(cutout, out_shape=output_shape)\n",
    "                cutouts.append(cutout)\n",
    "            if cutout_debug:\n",
    "                TF.to_pil_image(cutouts[-1].clamp(0, 1).squeeze(0)).save(\"cutout_InnerCrop.jpg\",quality=99)\n",
    "        cutouts = torch.cat(cutouts)\n",
    "        if skip_augs is not True: cutouts=self.augs(cutouts)\n",
    "        return cutouts\n",
    "\n",
    "def spherical_dist_loss(x, y):\n",
    "    x = F.normalize(x, dim=-1)\n",
    "    y = F.normalize(y, dim=-1)\n",
    "    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)     \n",
    "\n",
    "def tv_loss(input):\n",
    "    \"\"\"L2 total variation loss, as in Mahendran et al.\"\"\"\n",
    "    input = F.pad(input, (0, 1, 0, 1), 'replicate')\n",
    "    x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]\n",
    "    y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]\n",
    "    return (x_diff**2 + y_diff**2).mean([1, 2, 3])\n",
    "\n",
    "\n",
    "def range_loss(input):\n",
    "    return (input - input.clamp(-1, 1)).pow(2).mean([1, 2, 3])\n",
    "\n",
    "def _extract_into_tensor(arr, timesteps, broadcast_shape):\n",
    "    \"\"\"\n",
    "    Extract values from a 1-D numpy array for a batch of indices.\n",
    "    :param arr: the 1-D numpy array.\n",
    "    :param timesteps: a tensor of indices into the array to extract.\n",
    "    :param broadcast_shape: a larger shape of K dimensions with the batch\n",
    "                            dimension equal to the length of timesteps.\n",
    "    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        res = torch.from_numpy(arr).to(device=device)[timesteps].float()\n",
    "    except:\n",
    "         res = arr[timesteps].float()\n",
    "    while len(res.shape) < len(broadcast_shape):\n",
    "        res = res[..., None]\n",
    "    return res.expand(broadcast_shape)\n",
    "\n",
    "def cond_fn(diffusion, x, t, c, y=None):\n",
    "    cur_t = t\n",
    "    with torch.enable_grad():\n",
    "        x_is_NaN = False\n",
    "        x = x.detach().requires_grad_()\n",
    "        n = x.shape[0]\n",
    "        # if use_secondary_model is True:\n",
    "        #     alpha = torch.tensor(\n",
    "        #         diffusion.sqrt_alphas_cumprod[cur_t], device=device, dtype=torch.float32)\n",
    "        #     sigma = torch.tensor(\n",
    "        #         diffusion.sqrt_one_minus_alphas_cumprod[cur_t], device=device, dtype=torch.float32)\n",
    "        #     cosine_t = alpha_sigma_to_t(alpha, sigma)\n",
    "        #     out = secondary_model(x, cosine_t[None].repeat([n])).pred\n",
    "        #     fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t]\n",
    "        #     x_in = out * fac + x * (1 - fac)\n",
    "        #     x_in_grad = torch.zeros_like(x_in)\n",
    "        # else:\n",
    "        my_t = torch.ones([n], device=device, dtype=torch.long) * cur_t\n",
    "        out = diffusion.p_mean_variance(x, c, my_t, clip_denoised=False, return_x0=False)\n",
    "        fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t]\n",
    "        x_in = out[-1] * fac + x * (1 - fac)\n",
    "        x_in.requires_grad_()\n",
    "        x_in_grad = torch.zeros_like(x_in)\n",
    "        if clip_guidance_scale != 0:\n",
    "            for model_stat in model_stats:\n",
    "                for i in range(cutn_batches):\n",
    "                    t_int = int(t.item())+1 # errors on last step without +1, need to find source\n",
    "                    # when using SLIP Base model the dimensions need to be hard coded to avoid AttributeError: 'VisionTransformer' object has no attribute 'input_resolution'\n",
    "                    try:\n",
    "                        input_resolution = model_stat[\"clip_model\"].visual.input_resolution\n",
    "                    except:\n",
    "                        input_resolution = 224\n",
    "                    cuts = MakeCutoutsDango(input_resolution,\n",
    "                                            Overview=cut_overview[1000-t_int],\n",
    "                                            InnerCrop=cut_innercut[1000-t_int],\n",
    "                                            IC_Size_Pow=cut_ic_pow[1000-t_int],\n",
    "                                            IC_Grey_P=cut_icgray_p[1000-t_int]\n",
    "                                            )\n",
    "                    clip_in = normalize(cuts(model.differentiable_decode_first_stage(x_in).add(1).div(2)))\n",
    "                    image_embeds = model_stat[\"clip_model\"].encode_image(\n",
    "                        clip_in).float()\n",
    "                    dists = spherical_dist_loss(image_embeds.unsqueeze(\n",
    "                        1), model_stat[\"target_embeds\"].unsqueeze(0))\n",
    "                    dists = dists.view(\n",
    "                        [cut_overview[1000-t_int]+cut_innercut[1000-t_int], n, -1])\n",
    "                    losses = dists.mul(model_stat[\"weights\"]).sum(2).mean(0)\n",
    "                    # log loss, probably shouldn't do per cutn_batch\n",
    "                    loss_values.append(losses.sum().item())\n",
    "                    x_in_grad += torch.autograd.grad(losses.sum()\n",
    "                                                    * clip_guidance_scale, x_in)[0] / cutn_batches\n",
    "        if use_tv_scale != 0:\n",
    "            tv_losses = tv_loss(model.differentiable_decode_first_stage(x_in)) # do I really need to decode these first?\n",
    "        else:\n",
    "            tv_losses = toch.tensor(0)\n",
    "        # if use_secondary_model is True:\n",
    "        #     range_losses = range_loss(out)\n",
    "        # else:\n",
    "        if range_scale != 0:\n",
    "            range_losses = range_loss(model.differentiable_decode_first_stage(x_in))\n",
    "        else:\n",
    "            range_losses = torch.tensor(0)\n",
    "        if sat_scale != 0:\n",
    "            sat_losses = torch.abs(model.differentiable_decode_first_stage(x_in) - model.differentiable_decode_first_stage(x_in).clamp(min=-1, max=1)).mean()\n",
    "        else:\n",
    "            sat_losses = torch.tensor(0)\n",
    "        loss = tv_losses.sum() * tv_scale + range_losses.sum() * \\\n",
    "            range_scale + sat_losses.sum() * sat_scale\n",
    "        if init is not None and init_scale != 0:\n",
    "            init_losses = lpips_model(model.differentiable_decode_first_stage(x_in), model.differentiable_decode_first_stage(model.q_sample(init_latent, t)))\n",
    "            loss = loss + init_losses.sum() * init_scale\n",
    "        x_in_grad += torch.autograd.grad(loss, x_in)[0]\n",
    "        if torch.isnan(x_in_grad).any() == False:\n",
    "            grad = -torch.autograd.grad(x_in, x, x_in_grad)[0]\n",
    "        else:\n",
    "            # print(\"NaN'd\")\n",
    "            x_is_NaN = True\n",
    "            grad = torch.zeros_like(x)\n",
    "    if clamp_grad and x_is_NaN == False:\n",
    "        # min=-0.02, min=-clamp_max,\n",
    "        magnitude = grad.square().mean().sqrt()\n",
    "        return grad * magnitude.clamp(max=clamp_max) / magnitude\n",
    "    return grad\n",
    "\n",
    "class ScoreCorrector():\n",
    "    def __init__():\n",
    "        return\n",
    "    def modify_score(model, e_t, x, t, c, corrector_kwargs={}):\n",
    "        \"\"\"\n",
    "        Compute what the p_mean_variance output would have been, should the\n",
    "        model's score function be conditioned by cond_fn.\n",
    "        See condition_mean() for details on cond_fn.\n",
    "        Unlike condition_mean(), this instead uses the conditioning strategy\n",
    "        from Song et al (2020).\n",
    "        \"\"\"\n",
    "        alpha_bar = _extract_into_tensor(model.alphas_cumprod, t, x.shape)\n",
    "\n",
    "        eps = model._predict_eps_from_xstart(x, t, e_t)\n",
    "        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(\n",
    "            model, x, t, c, **corrector_kwargs\n",
    "        )\n",
    "\n",
    "        # out = p_mean_var.copy()\n",
    "        out = model.predict_start_from_noise(x, t, eps)\n",
    "        # out[\"pred_xstart\"] = model._predict_xstart_from_eps(x, t, eps)\n",
    "        # out[\"mean\"], _, _ = model.q_posterior_mean_variance(\n",
    "        #     x_start=out[\"pred_xstart\"], x_t=x, t=t\n",
    "        # )\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "BPnyd-XUKbfE",
    "outputId": "0fcd10e4-0df2-4ab9-cbf5-f08f4902c954",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set up model\n",
    "model = get_model()\n",
    "\n",
    "data_type = torch.half if use_fp16 else torch.float32\n",
    "\n",
    "os.makedirs(outdir, exist_ok=True)\n",
    "sample_path = os.path.join(outdir, \"samples\")\n",
    "os.makedirs(sample_path, exist_ok=True)\n",
    "base_count = len(os.listdir(sample_path))\n",
    "grid_count = len(os.listdir(outdir)) - 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Text -> Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# parameters\n",
    "\n",
    "# \"Joe Biden:0.2 happy\" use this format for every mode but text_interpolation\n",
    "# (242662456, \"Joe Biden:0.2 happy:0.2\") use this format for text_interpolation mode (seed, prompt)\n",
    "# for outcropping animations it will cycle through the list of prompts (repeating if frames > prompts)\n",
    "# prompt weighting works by slerping between conditioning vectors using the weight between prompts\n",
    "prompts = [\n",
    "    \"Joe Biden:0.2 happy\",\n",
    "]\n",
    "\n",
    "mode = 'txt2img' # one of [\"txt2img\",\"img2img\",\"text_interpolation\",\"inpainting\"]\n",
    "\n",
    "seed = 741 # seed for reproducible generations - use None for random seed\n",
    "average_weights = True # If using prompt weights, whether or not to average them\n",
    "\n",
    "sampler_type = 'k_euler_ancestral' # one of [\"ddim\",\"plms\",\"k_lms\",\"k_dpm2\",\"k_dpm2_ancestral\",\"k_heun\",\"k_euler\",\"k_euler_ancestral\"] (PLMS not compatible with init_image)\n",
    "\n",
    "ddim_steps = 100 # ddim sampling steps\n",
    "ddim_eta = 0.0 # ddim eta (eta=0.0 corresponds to deterministic sampling) (must be 0.0 if using PLMS/k_lms sampling)\n",
    "unconditional_guidance_scale = 7.5 # unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))\n",
    "\n",
    "dynamic_threshold = 10 # set to None if not used (0-100)\n",
    "static_threshold = None # set to None if not used\n",
    "\n",
    "precision = 'autocast' # precision to evaluate at (full or autocast)\n",
    "\n",
    "n_iter = 1 # how many sample iterations\n",
    "batch_size = 1 # how many samples to generate per prompt (set to 1 for non text_interpolation animations)\n",
    "n_rows = 0 # rows in the grid (will use batch_size if set to 0)\n",
    "\n",
    "fixed_code = True # if enabled, uses the same starting code across samples\n",
    "skip_grid = True # do not save a grid, only individual samples\n",
    "skip_save = False # do not save individual samples\n",
    "show_images = True # whether or not to show images after generation\n",
    "\n",
    "H = 512 # height\n",
    "W = 512 # width\n",
    "C = 4 # channels\n",
    "f = 8 # downsampling factor\n",
    "\n",
    "###########################\n",
    "# text interpolation settings (k_ samplers don't work well with batches greater than 1, needs investigating)\n",
    "save_mp4 = 'test.mp4'\n",
    "skip_save_video = False # do not save mp4\n",
    "loop = True # whether or not to make a looping animation\n",
    "degrees_per_second = 20 # degrees to travel per second\n",
    "fps = 40 # frames per second of output mp4\n",
    "fixed_seed_per_prompt = False # whether or not to use the same seed for every prompt\n",
    "fixed_seed = None # fixed seed to use if fixed_seed_per_prompt == True\n",
    "\n",
    "###########################\n",
    "# init image settings\n",
    "init_image_location = \"\" # file location of init image (use a blank string if none desired)\n",
    "init_noise_strength = 0.75 # how much to noise init image (0-1.0 where 1.0 is full destruction of init image information)\n",
    "\n",
    "###########################\n",
    "# inpainting and outcropping settings\n",
    "project_name = 'outcrop'\n",
    "out_crop_dir = \"out_crop\" # directory to save images in\n",
    "outcropping_animation = True # whether or not to run an outcropping animation\n",
    "outcropping_factor = 0.5 # how much to outcrop each step (0.0 - 1.0)\n",
    "outcropping_frames = 10 # how many frames to repeat outcropping procedure for\n",
    "original_image_location = \"\" # file location of the original image to inpaint (not needed for outcropping_animation if wanting to start from a SD render)\n",
    "mask_location = \"\" # file location of the image mask (not used if using outcropping_animation)\n",
    "increment_seed_every_frame = False # whether or not to increment the seed by the frame (only with fixed code)\n",
    "randomize_seed_every_frame = True # whether or not to randomize the seed by the frame (only with False increment_seed_every_frame)\n",
    "blur_strength = 1\n",
    "mask_offset_size = .90\n",
    "\n",
    "###########################\n",
    "# score corrector parameters (not valid for k_ samplers)\n",
    "use_score_corrector = False # whether or not to use score correction\n",
    "fuzzy_prompt = False\n",
    "image_prompt = False\n",
    "\n",
    "cutn_batches = 2\n",
    "\n",
    "clip_guidance_scale = 0\n",
    "tv_scale = 0\n",
    "range_scale = 0\n",
    "sat_scale = 0\n",
    "\n",
    "init = None\n",
    "init_scale = 0\n",
    "\n",
    "animation_mode = 'None'\n",
    "skip_augs = True\n",
    "\n",
    "clamp_max = 0.02\n",
    "clamp_grad = True\n",
    "\n",
    "cut_overview = [12]*400+[4]*600 \n",
    "cut_innercut = [4]*400+[12]*600 \n",
    "cut_ic_pow = [1]*1000\n",
    "cut_icgray_p = [0.2]*400+[0]*600\n",
    "\n",
    "####################################################################################\n",
    "# setup (ignore)\n",
    "if n_rows == 0:\n",
    "    n_rows = batch_size\n",
    "\n",
    "if seed is None:\n",
    "    seed = np.random.randint(np.iinfo(np.int16).max) if use_fp16 else np.random.randint(np.iinfo(np.int32).max)\n",
    "seed_everything(seed)\n",
    "\n",
    "start_code = None\n",
    "if fixed_code:\n",
    "    start_code = torch.randn([batch_size, C, H // f, W // f], device=device).to(data_type)\n",
    "\n",
    "precision_scope = autocast if precision==\"autocast\" else nullcontext\n",
    "data = list(chunk(prompts, batch_size))\n",
    "\n",
    "init_image_flag = False\n",
    "\n",
    "if mode == \"img2img\":\n",
    "    assert os.path.isfile(init_image_location)\n",
    "    init_image = load_img(init_image_location, (W, H)).to(device)\n",
    "    init_image = repeat(init_image, '1 ... -> b ...', b=batch_size).to(data_type)\n",
    "    init_image_latent = model.get_first_stage_encoding(model.encode_first_stage(init_image))  # move to latent space\n",
    "\n",
    "    assert 0. <= init_noise_strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
    "    t_enc = int(init_noise_strength * ddim_steps)\n",
    "    init_image_flag = True\n",
    "    print(f\"target t_enc is {t_enc} steps\")\n",
    "else:\n",
    "    init_image_flag = False\n",
    "                          \n",
    "if isinstance(init, str):\n",
    "    assert os.path.isfile(init)\n",
    "    init_image_ = load_img(init, (W, H)).to(device)\n",
    "    init_image_ = repeat(init_image, '1 ... -> b ...', b=batch_size).to(data_type)\n",
    "    init_latent = model.get_first_stage_encoding(model.encode_first_stage(init_image_))  # move to latent space\n",
    "else:\n",
    "    init = None\n",
    "\n",
    "if mode == 'inpainting':\n",
    "    show_generated_mask_and_image = False # debug\n",
    "    if outcropping_animation:\n",
    "        batch_size = 1\n",
    "        data = prompts * (outcropping_frames//len(prompts))\n",
    "        data.extend(prompts[:outcropping_frames%len(prompts)])\n",
    "        data = [[d] for d in data]\n",
    "    else:\n",
    "        data = list(chunk(prompts, batch_size))\n",
    "    if original_image_location != '':\n",
    "        assert os.path.isfile(original_image_location)\n",
    "        init = True\n",
    "        original_image = load_img(original_image_location, (W, H), outcrop_animation=outcropping_animation, factor=outcropping_factor).to(device)\n",
    "        original_image = repeat(original_image, '1 ... -> b ...', b=batch_size).to(data_type)\n",
    "        x0 = model.get_first_stage_encoding(model.encode_first_stage(original_image))  # move to latent space\n",
    "        init_latent = x0\n",
    "    if not outcropping_animation:\n",
    "        assert os.path.isfile(mask_location)\n",
    "    mask = load_img(mask_location, (W//f, H//f), mask=True, mask_offset_size=mask_offset_size, outcrop_animation=outcropping_animation, factor=outcropping_factor, show_generated_mask_and_image=show_generated_mask_and_image, blur_strength=blur_strength).to(device)\n",
    "    mask = repeat(mask, '1 ... -> b ...', b=batch_size).to(data_type)\n",
    "    inpainting_flag = True\n",
    "else:\n",
    "    x0 = None\n",
    "    mask = None\n",
    "    inpainting_flag = False\n",
    "    outcropping_animation = False\n",
    "    \n",
    "if mode == 'text_interpolation':\n",
    "    assert len(prompts[0]) == 2\n",
    "    frames_per_degree = fps / degrees_per_second\n",
    "\n",
    "    if fixed_seed_per_prompt:\n",
    "        if fixed_seed is None:\n",
    "            fixed_seed = np.random.randint(np.iinfo(np.int32).max)\n",
    "        for i in range(len(prompts)):\n",
    "            if isinstance(prompts[i], str):\n",
    "                prompts[i] = (fixed_seed, prompts[i])\n",
    "            else:\n",
    "                prompts[i][0] = fixed_seed\n",
    "\n",
    "    # interpolation setup\n",
    "    previous_c = None\n",
    "    previous_start_code = None\n",
    "    previous_seed = None\n",
    "    slerp_c_vectors = []\n",
    "    slerp_start_codes = []\n",
    "    seeds = []\n",
    "    for i, data in enumerate(map(lambda x: get_starting_code_and_conditioning_vector(*x, mode), prompts)):\n",
    "        c, start_code, seed = data\n",
    "        if i == 0:\n",
    "            slerp_c_vectors.append(c)\n",
    "            slerp_start_codes.append(start_code)\n",
    "            seeds.append(seed)\n",
    "        else:\n",
    "            start_norm = previous_c.flatten()/torch.norm(previous_c.flatten())\n",
    "            end_norm = c.flatten()/torch.norm(c.flatten())\n",
    "            omega = torch.acos((start_norm*end_norm).sum())\n",
    "            frames_c = round(omega.item() * frames_per_degree * 57.2957795)\n",
    "            start_norm = previous_start_code.flatten()/torch.norm(previous_start_code.flatten())\n",
    "            end_norm = start_code.flatten()/torch.norm(start_code.flatten())\n",
    "            omega = torch.acos((start_norm*end_norm).sum())\n",
    "            frames_start_code = round(omega.item() * frames_per_degree * 57.2957795)\n",
    "\n",
    "            frames = frames_c if frames_c >= frames_start_code else frames_start_code\n",
    "\n",
    "            original_c_shape = c.shape\n",
    "            original_start_code_shape = start_code.shape\n",
    "            c_vectors = get_slerp_vectors(previous_c.flatten(), c.flatten(), frames=frames)\n",
    "            c_vectors = c_vectors.reshape(-1, *original_c_shape)\n",
    "            slerp_c_vectors.extend(list(c_vectors[1:])) # drop first frame to prevent repeating frames\n",
    "            start_codes = get_slerp_vectors(previous_start_code.flatten(), start_code.flatten(), frames=frames)\n",
    "            start_codes = start_codes.reshape(-1, *original_start_code_shape)\n",
    "            slerp_start_codes.extend(list(start_codes[1:])) # drop first frame to prevent repeating frames\n",
    "            seed_interp = get_slerp_vectors(previous_seed, seed, frames=frames)\n",
    "            seeds.extend(list(seed_interp[1:]))\n",
    "            if loop and i == len(prompts) - 1:\n",
    "                c_vectors = get_slerp_vectors(c.flatten(), slerp_c_vectors[0].flatten(), frames=frames)\n",
    "                c_vectors = c_vectors.reshape(-1, *original_c_shape)\n",
    "                slerp_c_vectors.extend(list(c_vectors[1:-1])) # drop first and last frame to prevent repeating frames\n",
    "                start_codes = get_slerp_vectors(start_code.flatten(), slerp_start_codes[0].flatten(), frames=frames)\n",
    "                start_codes = start_codes.reshape(-1, *original_start_code_shape)\n",
    "                slerp_start_codes.extend(list(start_codes[1:-1])) # drop first and last frame to prevent repeating frames\n",
    "                seed_interp = get_slerp_vectors(seed, seeds[0], frames=frames)\n",
    "                seeds.extend(list(seed_interp[1:-1]))\n",
    "        previous_c = c\n",
    "        previous_start_code = start_code\n",
    "        previous_seed = seed\n",
    "\n",
    "    slerp_c_vectors = unflatten(slerp_c_vectors, batch_size)\n",
    "    slerp_start_codes = unflatten(slerp_start_codes, batch_size)\n",
    "    seeds = unflatten(seeds, batch_size)\n",
    "    seeds = [int(sum(seed)/len(seed)) for seed in seeds]\n",
    "    data = ['' for _ in slerp_c_vectors]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# run generation(s)\n",
    "if mode == 'text_interpolation' and not skip_save_video:\n",
    "    video_out = imageio.get_writer(save_mp4, mode='I', fps=fps, codec='libx264')\n",
    "with torch.no_grad():\n",
    "    with precision_scope(\"cuda\"):\n",
    "        with model.ema_scope():\n",
    "            tic = time.time()\n",
    "            all_samples = list()\n",
    "            if mode != 'text_interpolation':\n",
    "                seed_everything(seed)\n",
    "            for n in trange(n_iter, desc=\"Sampling\"):\n",
    "                frame = 1\n",
    "                if inpainting_flag and outcropping_animation:\n",
    "                    if init_image_flag == False:\n",
    "                        x0 = None\n",
    "                        mask_ = mask\n",
    "                        mask = None\n",
    "                        init = None\n",
    "                    else:\n",
    "                        frame = 2\n",
    "                for prompts in tqdm(data, desc=\"data\"):\n",
    "                    if mode == 'text_interpolation':\n",
    "                        seed_everything(seeds[frame-1])\n",
    "                    if use_score_corrector:\n",
    "                        model_stats = get_clip_embeds(prompts, seed)\n",
    "                    if inpainting_flag and outcropping_animation and (frame > 2 or (original_image_location == '' and frame == 2)):\n",
    "                        init = True\n",
    "                        image = load_img('prevframe.png', (W, H), outcrop_animation=True, factor=outcropping_factor).to(device)\n",
    "                        image = repeat(image, '1 ... -> b ...', b=batch_size).to(data_type)\n",
    "                        x0 = model.get_first_stage_encoding(model.encode_first_stage(image))\n",
    "\n",
    "                        if use_score_corrector:\n",
    "                            image = load_img('prevframe.png', (W, H), outcrop_animation=False, factor=outcropping_factor).to(device) # load image without cropping\n",
    "                            image = repeat(image, '1 ... -> b ...', b=batch_size)\n",
    "                            init_latent = model.get_first_stage_encoding(model.encode_first_stage(image))\n",
    "                        if increment_seed_every_frame and fixed_code:\n",
    "                            seed_everything(seed+frame)\n",
    "                            start_code = torch.randn([batch_size, C, H // f, W // f], device=device)\n",
    "                        elif randomize_seed_every_frame:\n",
    "                            start_code = None\n",
    "                            if fixed_code:\n",
    "                                seed_everything(np.random.randint(np.iinfo(np.int16).max) if use_fp16 else np.random.randint(np.iinfo(np.int32).max))\n",
    "                                start_code = torch.randn([batch_size, C, H // f, W // f], device=device)\n",
    "                        if mask == None:\n",
    "                            mask = mask_\n",
    "                    uc = None\n",
    "                    if unconditional_guidance_scale != 1.0:\n",
    "                        uc = model.get_learned_conditioning(batch_size * [\"\"])\n",
    "                    if isinstance(prompts, tuple):\n",
    "                        prompts = list(prompts)\n",
    "                    if mode != 'text_interpolation':\n",
    "                        c = torch.cat([get_conditioning_vector(prompt) for prompt in prompts])\n",
    "                    batch_size = c.shape[0]\n",
    "\n",
    "                    callback = make_callback(sampler=sampler_type,\n",
    "                                            dynamic_threshold=dynamic_threshold, \n",
    "                                            static_threshold=static_threshold,\n",
    "                                            inpainting=inpainting_flag,\n",
    "                                            x0=x0,\n",
    "                                            mask=mask)         \n",
    "\n",
    "                    shape = [C, H // f, W // f]\n",
    "\n",
    "                    if sampler_type in [\"k_lms\",\"k_dpm2\",\"k_dpm2_ancestral\",\"k_heun\",\"k_euler\",\"k_euler_ancestral\"]:\n",
    "                        model_wrap = K.external.CompVisDenoiser(model)\n",
    "                        sigmas = model_wrap.get_sigmas(ddim_steps)\n",
    "                        if inpainting_flag and outcropping_animation and (frame > 2 or (original_image_location == '' and frame == 2)):\n",
    "                            if increment_seed_every_frame:\n",
    "                                torch.manual_seed(seed+frame) # changes manual seeding procedure\n",
    "                            elif randomize_seed_every_frame:\n",
    "                                torch.manual_seed(np.random.randint(np.iinfo(np.int16).max) if use_fp16 else np.random.randint(np.iinfo(np.int32).max))\n",
    "                            else:\n",
    "                                torch.manual_seed(seed)  \n",
    "                        elif mode == 'text_interpolation':\n",
    "                            torch.manual_seed(seeds[frame-1])\n",
    "                        else:\n",
    "                            torch.manual_seed(seed) \n",
    "                        if init_image_flag:\n",
    "                            sigmas = sigmas[len(sigmas)-t_enc-1:]\n",
    "                            x = init_image_latent + torch.randn([batch_size, *shape], device=device) * sigmas[0]\n",
    "                        elif mode == 'text_interpolation':\n",
    "                            x = torch.cat(tuple(slerp_start_codes[frame-1]))\n",
    "                            c = torch.cat(tuple(torch.stack(list(slerp_c_vectors[frame-1]), dim=0)))\n",
    "                        else:\n",
    "                            x = torch.randn([batch_size, *shape], device=device) * sigmas[0] # for GPU draw\n",
    "                        model_wrap_cfg = CFGDenoiser(model_wrap)\n",
    "                        extra_args = {'cond': c, 'uncond': uc, 'cond_scale': unconditional_guidance_scale}\n",
    "                        if sampler_type == \"k_lms\":\n",
    "                            samples_ddim = K.sampling.sample_lms(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_dpm2\":\n",
    "                            samples_ddim = K.sampling.sample_dpm_2(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_dpm2_ancestral\":\n",
    "                            samples_ddim = K.sampling.sample_dpm_2_ancestral(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_heun\":\n",
    "                            samples_ddim = K.sampling.sample_heun(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_euler\":\n",
    "                            samples_ddim = K.sampling.sample_euler(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_euler_ancestral\":\n",
    "                            samples_ddim = K.sampling.sample_euler_ancestral(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                    else:\n",
    "                        if init_image_flag:\n",
    "                            # encode (scaled latent)\n",
    "                            if sampler_type == \"plms\":\n",
    "                                print('init_image not compatible with plms sampler, falling back to ddim')\n",
    "                            sampler = DDIMSampler(model)\n",
    "                            sampler.make_schedule(ddim_num_steps=ddim_steps, ddim_eta=ddim_eta, verbose=False)\n",
    "                            z_enc = sampler.stochastic_encode(init_image_latent, torch.tensor([t_enc]*batch_size).to(device))\n",
    "                            # decode it\n",
    "                            samples_ddim = sampler.decode(z_enc, c, t_enc, unconditional_guidance_scale=unconditional_guidance_scale,\n",
    "                                                    unconditional_conditioning=uc)\n",
    "                        else:\n",
    "                            if mode == 'text_interpolation':\n",
    "                                start_code = torch.cat(tuple(slerp_start_codes[frame-1]))\n",
    "                                c = torch.cat(tuple(torch.stack(list(slerp_c_vectors[frame-1]), dim=0)))\n",
    "                            if sampler_type == \"plms\":\n",
    "                                sampler = PLMSSampler(model)\n",
    "                            else:\n",
    "                                sampler = DDIMSampler(model)\n",
    "                            samples_ddim, _ = sampler.sample(S=ddim_steps,\n",
    "                                                                conditioning=c,\n",
    "                                                                batch_size=c.shape[0],\n",
    "                                                                shape=shape,\n",
    "                                                                verbose=False,\n",
    "                                                                unconditional_guidance_scale=unconditional_guidance_scale,\n",
    "                                                                unconditional_conditioning=uc,\n",
    "                                                                eta=ddim_eta,\n",
    "                                                                x_T=start_code,\n",
    "                                                                score_corrector=ScoreCorrector if use_score_corrector else None,\n",
    "                                                                corrector_kwargs={},\n",
    "                                                                img_callback=callback,\n",
    "                                                                x0=x0,\n",
    "                                                                mask=mask\n",
    "                                                            )\n",
    "\n",
    "                    x_samples_ddim = model.decode_first_stage(samples_ddim)\n",
    "                    x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)\n",
    "                    \n",
    "                    if (not skip_save or not skip_save_video or show_images) and not outcropping_animation:\n",
    "                        for x_sample in x_samples_ddim:\n",
    "                            x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')\n",
    "                            if mode == 'text_interpolation' and not skip_save_video:\n",
    "                                video_out.append_data(x_sample)\n",
    "                            if not skip_save or show_images:\n",
    "                                img = Image.fromarray(x_sample.astype(np.uint8))\n",
    "                            if not skip_save:\n",
    "                                img.save(os.path.join(sample_path, f\"{base_count:05}.png\"))\n",
    "                            if show_images:\n",
    "                                display(img)\n",
    "                    \n",
    "                    if outcropping_animation:\n",
    "                        for x_sample in x_samples_ddim:\n",
    "                            x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')\n",
    "                            img = Image.fromarray(x_sample.astype(np.uint8))\n",
    "                            img.save('prevframe.png')\n",
    "                            img.save(out_crop_dir.strip('/') + '/' + project_name + '_' + str(frame) + '.png')\n",
    "                            if show_images:\n",
    "                                display(img)\n",
    "\n",
    "                    if not skip_grid:\n",
    "                        all_samples.append(x_samples_ddim)\n",
    "                    \n",
    "                    frame += 1\n",
    "                    \n",
    "            if not skip_grid:\n",
    "                # additionally, save as grid\n",
    "                grid = torch.stack(all_samples, 0)\n",
    "                grid = rearrange(grid, 'n b c h w -> (n b) c h w')\n",
    "                grid = make_grid(grid, nrow=n_rows)\n",
    "\n",
    "                # to image\n",
    "                grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()\n",
    "                Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outdir, f'grid-{grid_count:04}.png'))\n",
    "                grid_count += 1\n",
    "\n",
    "            toc = time.time()\n",
    "\n",
    "print(f\"Your samples are ready and waiting for you here: \\n{outdir}\\n\"\n",
    "        f\"\\nTime elapsed: {round(toc - tic, 2)} seconds\")\n",
    "if mode == 'text_interpolation' and not skip_save_video:\n",
    "    video_out.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save outcropping frames to mp4 animation\n",
    "out_mp4 = 'outcrop.mp4'\n",
    "frames_between_crops = 30\n",
    "fps = 30\n",
    "frame_count = None # leave none to use outcropping frames (number of outcropped frames)\n",
    "\n",
    "####################################################################################\n",
    "\n",
    "if frame_count == None:\n",
    "    frame_count = outcropping_frames\n",
    "\n",
    "outcropping_frames_to_mp4(out_crop_dir, project_name, out_mp4, frame_count, frames_between_crops, fps, outcropping_factor, zoom_cropping=(0, 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Disco Diffusion Style Animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = [\n",
    "    [0, \"A dragons lair, epic matte painting, concept art, trending on artstation\"],\n",
    "    [100, \"A wizards magical potion room, epic matte painting, concept art, trending on artstation\"]\n",
    "]\n",
    "\n",
    "average_weights = True # whether or not to average text prompt weights (recommended to leave at True)\n",
    "seed = None # seed (set to None for random)\n",
    "loop = False # whether or not to loop prompts\n",
    "interpolate = False # whether or not to interpolate between prompt conditioning vectors\n",
    "\n",
    "sampler_type = 'k_euler_ancestral' # one of [\"ddim\",\"plms\",\"k_lms\",\"k_dpm2\",\"k_dpm2_ancestral\",\"k_heun\",\"k_euler\",\"k_euler_ancestral\"] (PLMS not compatible with init_image)\n",
    "\n",
    "ddim_steps = 200 # ddim sampling steps\n",
    "ddim_eta = 0.0 # ddim eta (eta=0.0 corresponds to deterministic sampling) (must be 0.0 if using PLMS/k_* sampling)\n",
    "unconditional_guidance_scale = 10.0 # unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))\n",
    "\n",
    "dynamic_threshold = None # set to None if not used\n",
    "static_threshold = None # set to None if not used\n",
    "\n",
    "color_match = True # whether or not to use color\n",
    "color_match_mode = 'cycle' # whether or not to try to match colors between frames for added coherence - ['RGB','HSV','LAB','cycle']\n",
    "noise_between_frames = 0.005 # amount of noise to add between frames\n",
    "\n",
    "show_images = True\n",
    "show_every_x_frames = 20\n",
    "\n",
    "init_latent_mixing = True # whether or not to mix init_latent in between timesteps\n",
    "mix_factor = \"0: (0.15), 25: (1.0)\" # keyframes (in timesteps) of amount of previous frame's latent to inject in between timesteps (1.0 - mix factor = amount mixed in) \n",
    "\n",
    "previous_frame_noise_strength = 0.4\n",
    "animation_mode = '3D'\n",
    "\n",
    "perlin_init = False # whether or not to use perlin noise (currently doesn't look good)\n",
    "perlin_mode = 'color' # color or gray\n",
    "\n",
    "init_image = None\n",
    "init_noise_strength = 0.4\n",
    "\n",
    "fixed_code = True\n",
    "\n",
    "n_iter = 1\n",
    "batch_size = 1\n",
    "\n",
    "start_frame = 0\n",
    "\n",
    "resume_run = False\n",
    "\n",
    "batch_name = 'test'\n",
    "batchNum = 1\n",
    "batchFolder = 'outputs'\n",
    "\n",
    "H = 512 # height\n",
    "W = 512 # width\n",
    "C = 4 # channels\n",
    "f = 8 # downsampling factor\n",
    "\n",
    "degrees_per_second = 12 # degrees to travel per second\n",
    "fps = 12 # frames per second of output mp4\n",
    "\n",
    "# score corrector parameters\n",
    "use_score_corrector = False # whether or not to use score correction (not compatible with k_ samplers)\n",
    "fuzzy_prompt = False\n",
    "image_prompt = False\n",
    "\n",
    "cutn_batches = 2\n",
    "\n",
    "clip_guidance_scale = 0\n",
    "tv_scale = 0\n",
    "range_scale = 0\n",
    "sat_scale = 0\n",
    "\n",
    "init = None\n",
    "init_scale = 0\n",
    "skip_augs = True\n",
    "\n",
    "clamp_max = 0.02\n",
    "clamp_grad = True\n",
    "\n",
    "cut_overview = [12]*400+[4]*600 \n",
    "cut_innercut = [4]*400+[12]*600 \n",
    "cut_ic_pow = [1]*1000\n",
    "cut_icgray_p = [0.2]*400+[0]*600\n",
    "\n",
    "frames_per_degree = fps / degrees_per_second\n",
    "\n",
    "precision = 'autocast' # precision to evaluate at (full or autocast)\n",
    "device='cuda'\n",
    "\n",
    "# video init mode\n",
    "video_init_path = \"init.mp4\"  # @param {type: 'string'}\n",
    "extract_nth_frame = 2  # @param {type: 'number'}\n",
    "persistent_frame_output_in_batch_folder = True  # @param {type: 'boolean'}\n",
    "video_init_seed_continuity = False  # @param {type: 'boolean'}\n",
    "# video optical flow settings:**\n",
    "video_init_flow_warp = True\n",
    "video_init_flow_blend = 0.999 # 0 - take next frame, 1 - take prev warped frame\n",
    "video_init_check_consistency = False\n",
    "video_init_blend_mode = \"optical flow\" # ['None', 'linear', 'optical flow']\n",
    "# Call optical flow from video frames and warp prev frame with flow\n",
    "if animation_mode == \"Video Input\":\n",
    "    # suggested by Chris the Wizard#8082 at discord\n",
    "    if persistent_frame_output_in_batch_folder:\n",
    "        videoFramesFolder = f'{batchFolder}/videoFrames'\n",
    "    else:\n",
    "        videoFramesFolder = f'/content/videoFrames'\n",
    "    createPath(videoFramesFolder)\n",
    "    print(f\"Exporting Video Frames (1 every {extract_nth_frame})...\")\n",
    "    try:\n",
    "        for f in pathlib.Path(f'{videoFramesFolder}').glob('*.jpg'):\n",
    "            f.unlink()\n",
    "    except:\n",
    "        print('')\n",
    "    vf = f'select=not(mod(n\\,{extract_nth_frame}))'\n",
    "    if os.path.exists(video_init_path):\n",
    "        subprocess.run(['ffmpeg', '-i', f'{video_init_path}', '-vf', f'{vf}', '-vsync', 'vfr', '-q:v', '2', '-loglevel',\n",
    "                       'error', '-stats', f'{videoFramesFolder}/%04d.jpg'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
    "    else:\n",
    "        print(\n",
    "            f'\\nWARNING!\\n\\nVideo not found: {video_init_path}.\\nPlease check your video path.\\n')\n",
    "\n",
    "# 2D/3D animation settings\n",
    "# `zoom` is a multiplier of dimensions, 1 is no zoom.\n",
    "# All rotations are provided in degrees.\n",
    "\n",
    "key_frames = True\n",
    "max_frames = 10000\n",
    "\n",
    "if animation_mode == \"Video Input\":\n",
    "    max_frames = len(glob(f'{videoFramesFolder}/*.jpg'))\n",
    "\n",
    "interp_spline = 'Linear' # Do not change, currently will not look good. param ['Linear','Quadratic','Cubic']\n",
    "angle = \"0:(0)\"\n",
    "zoom = \"0: (1.0)\"\n",
    "translation_x = \"0: (0)\"\n",
    "translation_y = \"0: (0)\"\n",
    "translation_z = \"0: (4.0)\"\n",
    "rotation_3d_x = \"0: (0.1)\"\n",
    "rotation_3d_y = \"0: (0)\"\n",
    "rotation_3d_z = \"0: (0.0)\"\n",
    "midas_depth_model = \"dpt_large\"\n",
    "midas_weight = 0.3\n",
    "near_plane = 200\n",
    "far_plane = 10000\n",
    "fov = 40\n",
    "padding_mode = 'border'\n",
    "sampling_mode = 'bicubic'\n",
    "\n",
    "# turbo mode\n",
    "turbo_mode = False\n",
    "turbo_steps = \"3\"\n",
    "turbo_preroll = 10  # frames\n",
    "\n",
    "# insist turbo be used only w 3d anim.\n",
    "if turbo_mode and animation_mode != '3D':\n",
    "    print('=====')\n",
    "    print('Turbo mode only available with 3D animations. Disabling Turbo.')\n",
    "    print('=====')\n",
    "    turbo_mode = False\n",
    "\n",
    "# video init coherency settings\n",
    "\n",
    "video_init_frames_scale = 15000\n",
    "video_init_frames_skip_steps = '70%'\n",
    "\n",
    "# vr mode\n",
    "vr_mode = False\n",
    "vr_eye_angle = 0.5 # `vr_eye_angle` is the y-axis rotation of the eyes towards the center\n",
    "vr_ipd = 5.0 # interpupillary distance (between the eyes)\n",
    "\n",
    "#################################################\n",
    "# setup (ignore)\n",
    "keyframes = True # temporary\n",
    "\n",
    "assert 0. <= previous_frame_noise_strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
    "t_enc = int(previous_frame_noise_strength * ddim_steps)\n",
    "print(f\"target t_enc is {t_enc} steps\")\n",
    "if init_image is not None:\n",
    "    init_t_enc = int(init_noise_strength * ddim_steps)\n",
    "    assert 0. <= init_noise_strength <= 1., 'can only work with strength in [0.0, 1.0]'\n",
    "    print(f\"target init_t_enc is {init_t_enc} steps\")\n",
    "\n",
    "# insist VR be used only w 3d anim.\n",
    "if vr_mode and animation_mode != '3D':\n",
    "    print('=====')\n",
    "    print('VR mode only available with 3D animations. Disabling VR.')\n",
    "    print('=====')\n",
    "    vr_mode = False\n",
    "\n",
    "if key_frames:\n",
    "    try:\n",
    "        mix_factor_series = get_inbetweens(parse_key_frames(mix_factor), t_enc)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `mix_factor` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `mix_factor` as \"\n",
    "            f'\"0: ({mix_factor})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        mix_factor = f\"0: ({mix_factor})\"\n",
    "        mix_factor_series = get_inbetweens(parse_key_frames(mix_factor), t_enc)\n",
    "    try:\n",
    "        angle_series = get_inbetweens(parse_key_frames(angle), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `angle` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `angle` as \"\n",
    "            f'\"0: ({angle})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        angle = f\"0: ({angle})\"\n",
    "        angle_series = get_inbetweens(parse_key_frames(angle), max_frames)\n",
    "\n",
    "    try:\n",
    "        zoom_series = get_inbetweens(parse_key_frames(zoom), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `zoom` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `zoom` as \"\n",
    "            f'\"0: ({zoom})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        zoom = f\"0: ({zoom})\"\n",
    "        zoom_series = get_inbetweens(parse_key_frames(zoom), max_frames)\n",
    "\n",
    "    try:\n",
    "        translation_x_series = get_inbetweens(parse_key_frames(translation_x), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `translation_x` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `translation_x` as \"\n",
    "            f'\"0: ({translation_x})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        translation_x = f\"0: ({translation_x})\"\n",
    "        translation_x_series = get_inbetweens(parse_key_frames(translation_x), max_frames)\n",
    "\n",
    "    try:\n",
    "        translation_y_series = get_inbetweens(parse_key_frames(translation_y), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `translation_y` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `translation_y` as \"\n",
    "            f'\"0: ({translation_y})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        translation_y = f\"0: ({translation_y})\"\n",
    "        translation_y_series = get_inbetweens(parse_key_frames(translation_y), max_frames)\n",
    "\n",
    "    try:\n",
    "        translation_z_series = get_inbetweens(parse_key_frames(translation_z), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `translation_z` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `translation_z` as \"\n",
    "            f'\"0: ({translation_z})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        translation_z = f\"0: ({translation_z})\"\n",
    "        translation_z_series = get_inbetweens(parse_key_frames(translation_z), max_frames)\n",
    "\n",
    "    try:\n",
    "        rotation_3d_x_series = get_inbetweens(parse_key_frames(rotation_3d_x), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `rotation_3d_x` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `rotation_3d_x` as \"\n",
    "            f'\"0: ({rotation_3d_x})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        rotation_3d_x = f\"0: ({rotation_3d_x})\"\n",
    "        rotation_3d_x_series = get_inbetweens(parse_key_frames(rotation_3d_x), max_frames)\n",
    "\n",
    "    try:\n",
    "        rotation_3d_y_series = get_inbetweens(parse_key_frames(rotation_3d_y), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `rotation_3d_y` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `rotation_3d_y` as \"\n",
    "            f'\"0: ({rotation_3d_y})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        rotation_3d_y = f\"0: ({rotation_3d_y})\"\n",
    "        rotation_3d_y_series = get_inbetweens(parse_key_frames(rotation_3d_y), max_frames)\n",
    "\n",
    "    try:\n",
    "        rotation_3d_z_series = get_inbetweens(parse_key_frames(rotation_3d_z), max_frames)\n",
    "    except RuntimeError as e:\n",
    "        print(\n",
    "            \"WARNING: You have selected to use key frames, but you have not \"\n",
    "            \"formatted `rotation_3d_z` correctly for key frames.\\n\"\n",
    "            \"Attempting to interpret `rotation_3d_z` as \"\n",
    "            f'\"0: ({rotation_3d_z})\"\\n'\n",
    "            \"Please read the instructions to find out how to use key frames \"\n",
    "            \"correctly.\\n\"\n",
    "        )\n",
    "        rotation_3d_z = f\"0: ({rotation_3d_z})\"\n",
    "        rotation_3d_z_series = get_inbetweens(parse_key_frames(rotation_3d_z), max_frames)\n",
    "\n",
    "else:\n",
    "    angle = float(angle)\n",
    "    zoom = float(zoom)\n",
    "    translation_x = float(translation_x)\n",
    "    translation_y = float(translation_y)\n",
    "    translation_z = float(translation_z)\n",
    "    rotation_3d_x = float(rotation_3d_x)\n",
    "    rotation_3d_y = float(rotation_3d_y)\n",
    "    rotation_3d_z = float(rotation_3d_z)\n",
    "\n",
    "\n",
    "if seed is None:\n",
    "    seed = np.random.randint(np.iinfo(np.int32).max)\n",
    "\n",
    "args = {\n",
    "    'seed': seed,\n",
    "    'animation_mode': animation_mode,\n",
    "    'init_image': init_image,\n",
    "    'perlin_init': perlin_init,\n",
    "    'perlin_mode': perlin_mode,\n",
    "    'H': H,\n",
    "    'W': W,\n",
    "    'C': C,\n",
    "    'f': f,\n",
    "    'start_frame': start_frame,\n",
    "    'max_frames': max_frames,\n",
    "    'n_iter': n_iter,\n",
    "    'batch_size': batch_size,\n",
    "    'init_noise_strength': init_noise_strength,\n",
    "    'previous_frame_noise_strength': previous_frame_noise_strength,\n",
    "    'ddim_steps': ddim_steps,\n",
    "    'ddim_eta': ddim_eta,\n",
    "    'unconditional_guidance_scale': unconditional_guidance_scale,\n",
    "    'fixed_code': fixed_code,\n",
    "    'key_frames': key_frames,\n",
    "    'angle_series': angle_series,\n",
    "    'zoom_series': zoom_series,\n",
    "    'translation_x_series': translation_x_series,\n",
    "    'translation_y_series': translation_y_series,\n",
    "    'translation_z_series': translation_z_series,\n",
    "    'rotation_3d_x_series': rotation_3d_x_series,\n",
    "    'rotation_3d_y_series': rotation_3d_y_series,\n",
    "    'rotation_3d_z_series': rotation_3d_z_series,\n",
    "    'near_plane': near_plane,\n",
    "    'far_plane': far_plane,\n",
    "    'fov': fov,\n",
    "    'padding_mode': padding_mode,\n",
    "    'sampling_mode': sampling_mode,\n",
    "    'midas_weight': midas_weight\n",
    "}\n",
    "\n",
    "args = SimpleNamespace(**args)\n",
    "\n",
    "seed_everything(seed)\n",
    "\n",
    "start_code = None\n",
    "if fixed_code:\n",
    "    start_code = torch.randn([batch_size, C, H // f, W // f], device=device)\n",
    "\n",
    "precision_scope = autocast if precision==\"autocast\" else nullcontext\n",
    "\n",
    "if init_image is None:\n",
    "    init_image_flag = False\n",
    "else:\n",
    "    init_image_flag = True\n",
    "    \n",
    "if isinstance(prompts[0], list) and len(prompts) > 1:\n",
    "    keyframed_prompts = True\n",
    "    frame_series, prompts = map(list, zip(*prompts))\n",
    "    frame_series[0] = 0\n",
    "else:\n",
    "    keyframed_prompts = False\n",
    "    frame_series = [0]\n",
    "\n",
    "if interpolate and len(prompts) > 1:\n",
    "    previous_c = None\n",
    "    slerp_c_vectors = []\n",
    "    for i, c in enumerate(map(lambda x: get_conditioning_vector(x), prompts)):\n",
    "        if i == 0:\n",
    "            slerp_c_vectors.append(c)\n",
    "        else:\n",
    "            if not keyframed_prompts:\n",
    "                start_norm = previous_c.flatten()/torch.norm(previous_c.flatten())\n",
    "                end_norm = c.flatten()/torch.norm(c.flatten())\n",
    "                omega = torch.acos((start_norm*end_norm).sum())\n",
    "                frames = round(omega.item() * frames_per_degree * 57.2957795)\n",
    "                frame_series.append(frame_series[i-1]+frames)\n",
    "            else:\n",
    "                frames = frame_series[i] - frame_series[i-1]\n",
    "            assert frames > 0, 'Frames between interpolated vectors must be greater than 0'\n",
    "\n",
    "            original_c_shape = c.shape\n",
    "            c_vectors = get_slerp_vectors(previous_c.flatten(), c.flatten(), frames=frames)\n",
    "            c_vectors = c_vectors.reshape(-1, *original_c_shape)\n",
    "            slerp_c_vectors.extend(list(c_vectors[1:])) # drop first frame to prevent repeating frames\n",
    "            if loop and i == len(prompts) - 1:\n",
    "                c_vectors = get_slerp_vectors(c.flatten(), slerp_c_vectors[0].flatten(), frames=frames)\n",
    "                c_vectors = c_vectors.reshape(-1, *original_c_shape)\n",
    "                slerp_c_vectors.extend(list(c_vectors[1:-1])) # drop first and last frame to prevent repeating frames\n",
    "        previous_c = c\n",
    "    data = ['']\n",
    "else:\n",
    "    data = ['']\n",
    "    interpolate = False\n",
    "    \n",
    "prompts_ = prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize midas depth model\n",
    "if args.animation_mode == \"3D\":\n",
    "    midas_model, midas_transform, midas_net_w, midas_net_h, midas_resize_mode, midas_normalization = init_midas_depth_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Run animation loop - WIP - reach out with any bugs/questions/concerns - https://github.com/francislabountyjr/stable-diffusion\n",
    "TRANSLATION_SCALE = 1.0/200.0\n",
    "stop_on_next_loop = False\n",
    "\n",
    "# make sure sampler is DDIM if using PLMS\n",
    "if sampler_type in ['plms', 'ddim']:\n",
    "    if sampler_type == 'plms':\n",
    "        print('PLMS not compatible, falling back to DDIM')\n",
    "    sampler = DDIMSampler(model)\n",
    "else:\n",
    "    sampler = None\n",
    "\n",
    "color_match_sample = None\n",
    "\n",
    "switch_frame = frame_series[1] if len(frame_series) > 1 else frame_series[0]\n",
    "switch_idx = 0\n",
    "for frame_num in tqdm(range(args.start_frame, args.max_frames), desc='Frames'):\n",
    "    if stop_on_next_loop:\n",
    "        break\n",
    "\n",
    "    if frame_num == 0:\n",
    "        init_image = args.init_image if (args.init_image != '' or args.init_image is not None) else None\n",
    "\n",
    "    if args.animation_mode == \"2D\":\n",
    "        if args.key_frames:\n",
    "            angle = args.angle_series[frame_num]\n",
    "            zoom = args.zoom_series[frame_num]\n",
    "            translation_x = args.translation_x_series[frame_num]\n",
    "            translation_y = args.translation_y_series[frame_num]\n",
    "            print(\n",
    "                f'angle: {angle}',\n",
    "                f'zoom: {zoom}',\n",
    "                f'translation_x: {translation_x}',\n",
    "                f'translation_y: {translation_y}',\n",
    "            )\n",
    "\n",
    "        if frame_num > 0:\n",
    "            if resume_run and frame_num == start_frame:\n",
    "                img_0 = cv2.imread(\n",
    "                    batchFolder+f\"/{batch_name}({batchNum})_{start_frame-1:04}.png\")\n",
    "                if color_match:\n",
    "                    color_match_sample = cv2.imread(batchFolder + \\\n",
    "                        f\"/{batch_name}({batchNum})_{0}.png\").copy()\n",
    "            else:\n",
    "                img_0 = cv2.imread('prevFrame.png')\n",
    "                center = (1*img_0.shape[1]//2, 1*img_0.shape[0]//2)\n",
    "                trans_mat = np.float32(\n",
    "                    [[1, 0, translation_x],\n",
    "                     [0, 1, translation_y]]\n",
    "                )\n",
    "                rot_mat = cv2.getRotationMatrix2D(center, angle, zoom)\n",
    "                trans_mat = np.vstack([trans_mat, [0, 0, 1]])\n",
    "                rot_mat = np.vstack([rot_mat, [0, 0, 1]])\n",
    "                transformation_matrix = np.matmul(rot_mat, trans_mat)\n",
    "                img_0 = cv2.warpPerspective(\n",
    "                    img_0,\n",
    "                    transformation_matrix,\n",
    "                    (img_0.shape[1], img_0.shape[0]),\n",
    "                    borderMode=cv2.BORDER_WRAP\n",
    "                )\n",
    "                \n",
    "            # apply color matching\n",
    "            if color_match:\n",
    "                if color_match_sample is None:\n",
    "                    color_match_sample = img_0.copy()\n",
    "                else:\n",
    "                    if color_match_mode == 'cycle':\n",
    "                        img_0 = maintain_colors(img_0, color_match_sample, ['RGB','HSV','LAB'][frame_num%3])\n",
    "                    else:\n",
    "                        img_0 = maintain_colors(img_0, color_match_sample, color_match_mode)\n",
    "\n",
    "            # apply frame noising\n",
    "            if noise_between_frames > 0:\n",
    "                img_0 = add_noise(sample_from_cv2(img_0), noise_between_frames)\n",
    "\n",
    "            cv2.imwrite('prevFrameScaled.png', img_0)\n",
    "            init_image = 'prevFrameScaled.png'\n",
    "        \n",
    "    if args.animation_mode == \"3D\":\n",
    "        if frame_num > 0:\n",
    "            # seed += 1\n",
    "            if resume_run and frame_num == start_frame:\n",
    "                img_filepath = batchFolder + \\\n",
    "                    f\"/{batch_name}({batchNum})_{start_frame-1}.png\"\n",
    "                if color_match:\n",
    "                    color_match_sample = Image.open(batchFolder + \\\n",
    "                        f\"/{batch_name}({batchNum})_{0}.png\").copy()\n",
    "                if turbo_mode and frame_num > turbo_preroll:\n",
    "                    shutil.copyfile(img_filepath, 'oldFrameScaled.png')\n",
    "            else:\n",
    "                img_filepath = 'prevFrame.png'\n",
    "\n",
    "            next_step_pil = do_3d_step(\n",
    "                img_filepath, frame_num, midas_model, midas_transform)\n",
    "            \n",
    "            # apply color matching\n",
    "            if color_match:\n",
    "                if color_match_sample is None:\n",
    "                    color_match_sample = Image.open(img_filepath).copy()\n",
    "                else:\n",
    "                    if color_match_mode == 'cycle':\n",
    "                        next_step_pil = maintain_colors(np.asarray(next_step_pil).copy(), np.asarray(color_match_sample).copy(), ['RGB','HSV','LAB'][frame_num%3])\n",
    "                    else:\n",
    "                        next_step_pil = maintain_colors(np.asarray(next_step_pil).copy(), np.asarray(color_match_sample).copy(), color_match_mode)\n",
    "                    next_step_pil = Image.fromarray(next_step_pil)\n",
    "            \n",
    "            if turbo_mode and frame_num != turbo_preroll and frame_num > turbo_preroll and frame_num % int(turbo_steps) == 0:\n",
    "                next_step_pil.save('oldFrameScaled.png') # to prevent blending in noise for turbo mode\n",
    "\n",
    "            # apply frame noising\n",
    "            if noise_between_frames > 0 and ((turbo_mode and not frame_num == turbo_preroll and frame_num % int(turbo_steps) == 0 and frame_num > turbo_preroll) or not turbo_mode):\n",
    "                next_step_pil = add_noise(sample_from_pil(next_step_pil), noise_between_frames)\n",
    "                next_step_pil = sample_to_pil(next_step_pil)\n",
    "\n",
    "            next_step_pil.save('prevFrameScaled.png')\n",
    "\n",
    "            # Turbo mode - skip some diffusions, use 3d morph for clarity and to save time\n",
    "            if turbo_mode:\n",
    "                if frame_num == turbo_preroll:  # start tracking oldframe\n",
    "                    # stash for later blending\n",
    "                    next_step_pil.save('oldFrameScaled.png')\n",
    "                elif frame_num > turbo_preroll:\n",
    "                    # set up 2 warped image sequences, old & new, to blend toward new diff image\n",
    "                    old_frame = do_3d_step(\n",
    "                        'oldFrameScaled.png', frame_num, midas_model, midas_transform)\n",
    "                    \n",
    "                    old_frame.save('oldFrameScaled.png')\n",
    "                    if frame_num % int(turbo_steps) != 0:\n",
    "                        print(\n",
    "                            'turbo skip this frame: skipping clip diffusion steps')\n",
    "                        filename = f'{batch_name}({batchNum})_{frame_num}.png'\n",
    "                        blend_factor = (\n",
    "                            (frame_num % int(turbo_steps))+1)/int(turbo_steps)\n",
    "                        print(\n",
    "                            'turbo skip this frame: skipping clip diffusion steps and saving blended frame')\n",
    "                        # this is already updated..\n",
    "                        newWarpedImg = cv2.imread('prevFrameScaled.png')\n",
    "                        oldWarpedImg = cv2.imread('oldFrameScaled.png')\n",
    "                        blendedImage = cv2.addWeighted(\n",
    "                            newWarpedImg, blend_factor, oldWarpedImg, 1-blend_factor, 0.0)\n",
    "                        \n",
    "                        \n",
    "                        \n",
    "                        cv2.imwrite(\n",
    "                            f'{batchFolder}/{filename}', blendedImage)\n",
    "                        # save it also as prev_frame to feed next iteration\n",
    "                        # next_step_pil.save(f'{img_filepath}')\n",
    "                                        \n",
    "                        next_step_pil.save('prevFrame.png')\n",
    "                        if vr_mode:\n",
    "                            generate_eye_views(\n",
    "                                TRANSLATION_SCALE, batchFolder, filename, frame_num, midas_model, midas_transform)\n",
    "                        continue\n",
    "                    else:\n",
    "                        # if not a skip frame, will run diffusion and need to blend.\n",
    "                        # oldWarpedImg = cv2.imread('prevFrameScaled.png')\n",
    "                        # swap in for blending later\n",
    "                        # cv2.imwrite(f'oldFrameScaled.png', oldWarpedImg)\n",
    "                        print('clip/diff this frame - generate clip diff image')\n",
    "\n",
    "            init_image = 'prevFrameScaled.png'\n",
    "\n",
    "        if args.animation_mode == \"Video Input\":\n",
    "            init_scale = args.video_init_frames_scale\n",
    "            skip_steps = args.calc_frames_skip_steps\n",
    "            if not video_init_seed_continuity:\n",
    "                seed += 1\n",
    "            if video_init_flow_warp:\n",
    "                if frame_num == 0:\n",
    "                    skip_steps = args.video_init_skip_steps\n",
    "                    init_image = f'{videoFramesFolder}/{frame_num+1:04}.jpg'\n",
    "                if frame_num > 0:\n",
    "                    prev = PIL.Image.open(\n",
    "                        batchFolder+f\"/{batch_name}({batchNum})_{frame_num-1}.png\")\n",
    "\n",
    "                    frame1_path = f'{videoFramesFolder}/{frame_num}.jpg'\n",
    "                    frame2 = PIL.Image.open(\n",
    "                        f'{videoFramesFolder}/{frame_num+1:04}.jpg')\n",
    "                    flo_path = f\"/{flo_folder}/{frame1_path.split('/')[-1]}.npy\"\n",
    "\n",
    "                    init_image = 'warped.png'\n",
    "                    print(video_init_flow_blend)\n",
    "                    weights_path = None\n",
    "                    if video_init_check_consistency:\n",
    "                        # TBD\n",
    "                        pass\n",
    "\n",
    "                    warp(prev, frame2, flo_path, blend=video_init_flow_blend,\n",
    "                         weights_path=weights_path).save(init_image)\n",
    "\n",
    "            else:\n",
    "                init_image = f'{videoFramesFolder}/{frame_num+1:04}.jpg'\n",
    "\n",
    "    seed_everything(args.seed+frame_num)\n",
    "    if interpolate:\n",
    "        model_stats = []\n",
    "    else:\n",
    "        if use_score_corrector:\n",
    "            model_stats = get_clip_embeds(prompts, seed)\n",
    "    init = None\n",
    "    if init_image is not None:\n",
    "        init = load_img(init_image, (args.W, args.H)).to(device)\n",
    "        init = repeat(init, '1 ... -> b ...', b=args.batch_size).to(data_type)\n",
    "        init_latent = model.get_first_stage_encoding(model.encode_first_stage(init))  # move to latent space\n",
    "        init_image_flag = True\n",
    "\n",
    "    if args.perlin_init and frame_num == 0:\n",
    "        if args.perlin_mode == 'color':\n",
    "            init = create_perlin_noise(\n",
    "                [1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
    "            init2 = create_perlin_noise(\n",
    "                [1.5**-i*0.5 for i in range(8)], 4, 4, False)\n",
    "        elif args.perlin_mode == 'gray':\n",
    "            init = create_perlin_noise(\n",
    "                [1.5**-i*0.5 for i in range(12)], 1, 1, True)\n",
    "            init2 = create_perlin_noise(\n",
    "                [1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
    "        else:\n",
    "            init = create_perlin_noise(\n",
    "                [1.5**-i*0.5 for i in range(12)], 1, 1, False)\n",
    "            init2 = create_perlin_noise(\n",
    "                [1.5**-i*0.5 for i in range(8)], 4, 4, True)\n",
    "        init = TF.to_tensor(init).add(TF.to_tensor(init2)).div(\n",
    "            2).to(device).unsqueeze(0).mul(2).sub(1)\n",
    "        del init2\n",
    "\n",
    "    print(f'Frame {frame_num}')\n",
    "\n",
    "    image_display = Output()\n",
    "    with torch.no_grad():\n",
    "        with precision_scope(\"cuda\"):\n",
    "            with model.ema_scope():\n",
    "                tic = time.time()\n",
    "                all_samples = list()\n",
    "                print('')\n",
    "                display(image_display)\n",
    "                gc.collect()\n",
    "                torch.cuda.empty_cache()\n",
    "\n",
    "                if perlin_init and frame_num == 0:\n",
    "                    init = regen_perlin()\n",
    "                    start_code = init\n",
    "\n",
    "                for _ in tqdm(data, desc=\"data\"):\n",
    "                    uc = None\n",
    "                    if unconditional_guidance_scale != 1.0:\n",
    "                        uc = model.get_learned_conditioning(batch_size * [\"\"])\n",
    "                    # if isinstance(prompts, tuple):\n",
    "                    #     prompts = list(prompts)\n",
    "              \n",
    "                    callback = make_callback(sampler=sampler_type,\n",
    "                                            dynamic_threshold=dynamic_threshold, \n",
    "                                            static_threshold=static_threshold,\n",
    "                                            inpainting=False,\n",
    "                                            x0=None,\n",
    "                                            mask=None)\n",
    "                    \n",
    "                    if interpolate:\n",
    "                        if loop:\n",
    "                            c = slerp_c_vectors[frame_num%len(slerp_c_vectors)]\n",
    "                        else:\n",
    "                            c = slerp_c_vectors[min(frame_num,len(slerp_c_vectors)-1)]\n",
    "                        c = torch.cat([c])\n",
    "                        if switch_frame == frame_num:\n",
    "                            if switch_idx < len(frame_series):\n",
    "                                switch_idx = switch_idx + 1 \n",
    "                                if switch_idx + 1 < len(frame_series):\n",
    "                                    switch_frame = frame_series[switch_idx + 1]\n",
    "                        if not loop:\n",
    "                            print(f'Current Prompt: {prompts_[switch_idx]}')\n",
    "                    else:\n",
    "                        if keyframed_prompts:\n",
    "                            if switch_frame == frame_num:\n",
    "                                if switch_idx < len(frame_series):\n",
    "                                    switch_idx = switch_idx + 1 \n",
    "                                    if switch_idx + 1 < len(frame_series):\n",
    "                                        switch_frame = frame_series[switch_idx + 1]\n",
    "                                    \n",
    "                            prompts = [prompts_[switch_idx]]\n",
    "                            print(f'Current Prompt: {prompts[0]}')\n",
    "                        else:\n",
    "                            if loop:\n",
    "                                prompts = [prompts_[frame_num%len(prompts_)]]\n",
    "                            else:\n",
    "                                prompts = [prompts_[min(frame_num,len(prompts_)-1)]]\n",
    "                            print(f'Current Prompt: {prompts[0]}')\n",
    "                        c = torch.cat([get_conditioning_vector(prompt) for prompt in prompts])\n",
    "              \n",
    "                    shape = [args.C, args.H // args.f, args.W // args.f]\n",
    "                    if sampler_type in [\"k_lms\",\"k_dpm2\",\"k_dpm2_ancestral\",\"k_heun\",\"k_euler\",\"k_euler_ancestral\"]:\n",
    "                        model_wrap = K.external.CompVisDenoiser(model)\n",
    "                        sigmas = model_wrap.get_sigmas(ddim_steps)\n",
    "                        torch.manual_seed(args.seed+frame_num) # changes manual seeding procedure \n",
    "                        if init_image_flag:\n",
    "                            if frame_num == 0:\n",
    "                                sigmas = sigmas[len(sigmas)-init_t_enc-1:]\n",
    "                            else:\n",
    "                                sigmas = sigmas[len(sigmas)-t_enc-1:]\n",
    "                            noise = torch.randn([batch_size, *shape], device=device)\n",
    "                            x = init_latent + noise * sigmas[0]\n",
    "                            callback = make_callback(sampler=sampler_type,\n",
    "                                            dynamic_threshold=dynamic_threshold, \n",
    "                                            static_threshold=static_threshold,\n",
    "                                            inpainting=False,\n",
    "                                            mix_with_x0=init_latent_mixing,\n",
    "                                            x0=init_latent,\n",
    "                                            noise=noise,\n",
    "                                            mix_factor=mix_factor_series,\n",
    "                                            mask=None)\n",
    "                        else:\n",
    "                            x = torch.randn([batch_size, *shape], device=device) * sigmas[0] # for GPU draw\n",
    "                        model_wrap_cfg = CFGDenoiser(model_wrap)\n",
    "                        extra_args = {'cond': c, 'uncond': uc, 'cond_scale': args.unconditional_guidance_scale}\n",
    "                        if sampler_type == \"k_lms\":\n",
    "                            samples_ddim = K.sampling.sample_lms(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_dpm2\":\n",
    "                            samples_ddim = K.sampling.sample_dpm_2(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_dpm2_ancestral\":\n",
    "                            samples_ddim = K.sampling.sample_dpm_2_ancestral(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_heun\":\n",
    "                            samples_ddim = K.sampling.sample_heun(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_euler\":\n",
    "                            samples_ddim = K.sampling.sample_euler(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                        elif sampler_type == \"k_euler_ancestral\":\n",
    "                            samples_ddim = K.sampling.sample_euler_ancestral(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=False, callback=callback)\n",
    "                    else:\n",
    "                        if init_image_flag:\n",
    "                            # encode (scaled latent)\n",
    "                            sampler.make_schedule(ddim_num_steps=ddim_steps, ddim_eta=ddim_eta, verbose=False)\n",
    "                            z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(device))\n",
    "                            # decode it\n",
    "                            samples_ddim = sampler.decode(z_enc, c, t_enc, unconditional_guidance_scale=unconditional_guidance_scale,\n",
    "                                                    unconditional_conditioning=uc)\n",
    "                        else:\n",
    "                            if sampler_type == \"plms\":\n",
    "                                sampler = PLMSSampler(model)\n",
    "                            else:\n",
    "                                sampler = DDIMSampler(model)\n",
    "                            samples_ddim, _ = sampler.sample(S=ddim_steps,\n",
    "                                                                conditioning=c,\n",
    "                                                                batch_size=c.shape[0],\n",
    "                                                                shape=shape,\n",
    "                                                                verbose=False,\n",
    "                                                                unconditional_guidance_scale=args.unconditional_guidance_scale,\n",
    "                                                                unconditional_conditioning=uc,\n",
    "                                                                eta=args.ddim_eta,\n",
    "                                                                x_T=start_code,\n",
    "                                                                score_corrector=ScoreCorrector if use_score_corrector else None,\n",
    "                                                                corrector_kwargs={},\n",
    "                                                                img_callback=callback\n",
    "                                                            )\n",
    "\n",
    "                    x_samples_ddim = model.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",
    "                    for x_sample in x_samples_ddim:\n",
    "                        x_sample = 255. * \\\n",
    "                            rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')\n",
    "                        img = Image.fromarray(x_sample.astype(np.uint8))\n",
    "                        if show_images and frame_num % show_every_x_frames == 0:\n",
    "                            display(img)\n",
    "                        if args.animation_mode != \"None\":\n",
    "                            filename = f'{batch_name}({batchNum})_{frame_num}.png'\n",
    "                            img.save('prevFrame.png')\n",
    "                            img.save(f'{batchFolder}/{filename}')\n",
    "                            # if frame_num == 0:\n",
    "                            #     save_settings()\n",
    "                            if args.animation_mode == \"3D\":\n",
    "                                # If turbo, save a blended image\n",
    "                                if turbo_mode and frame_num > 0:\n",
    "                                    # Mix new image with prevFrameScaled\n",
    "                                    blend_factor = (1)/int(turbo_steps)\n",
    "                                    # This is already updated..\n",
    "                                    newFrame = cv2.imread('prevFrame.png')\n",
    "                                    prev_frame_warped = cv2.imread(\n",
    "                                        'prevFrameScaled.png')\n",
    "                                    blendedImage = cv2.addWeighted(\n",
    "                                        newFrame, blend_factor, prev_frame_warped, (1-blend_factor), 0.0)\n",
    "                                    cv2.imwrite(\n",
    "                                        f'{batchFolder}/{filename}', blendedImage)\n",
    "                                else:\n",
    "                                    img.save(f'{batchFolder}/{filename}')\n",
    "\n",
    "                                with precision_scope('cuda', enabled=False):\n",
    "                                    if vr_mode:\n",
    "                                        generate_eye_views(\n",
    "                                            TRANSLATION_SCALE, batchFolder, filename, frame_num, midas_model, midas_transform)\n",
    "                toc = time.time()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "latent-imagenet-diffusion.ipynb",
   "provenance": []
  },
  "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.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "96115ba3d133a87032cd70b9a322d48617443b59c68527f3d6b41f2029e5d2d0"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}