ldm.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8179e2c6-ffdb-4f91-a456-1b5bb44e0a31",
   "metadata": {},
   "source": [
    "- LDM trainer and sampler for triplanes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "faca32d1-1bdc-4094-abab-7ea8d7d28aef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/root/dev/jjuke_diffusion'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e506b2be-5f34-4057-a26d-afe7cff1a904",
   "metadata": {},
   "source": [
    "# Unet for LDM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d6ed669-f8e2-4662-b0c9-9b2323b0abb8",
   "metadata": {},
   "source": [
    "Neural backbone $ \\epsilon_\\theta( \\mathbf{z}_t, t, \\mathbf{y}) $ of the LDM\n",
    "\n",
    "- $ \\mathbf{y} $: Only handling text here!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1b1e554d-fb18-4478-8103-92cf01863669",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jjuke_diffusion.unet_cond.unet import UNetModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65f329a1-56a2-4cc7-9c91-c2bf46b60de0",
   "metadata": {},
   "source": [
    "# Conditioning Mechanism"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55bff526-d8ec-4330-b237-7acb52db9d16",
   "metadata": {},
   "source": [
    "Domain specific encoder $ \\tau_\\theta $\n",
    "\n",
    "- Input:\n",
    "    - $ \\mathcal{y} $\n",
    "- Output:\n",
    "    - Intermediate representation $ \\tau_\\theta(\\mathbf{y}) \\in \\mathbb{R}^{M \\times d_\\tau} $\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d6accea-d577-4ffc-aa5e-a86f6c073be1",
   "metadata": {},
   "source": [
    "## Domain specific encoder with BERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "854f24a1-e8b9-4677-87cd-a2e0f0f9f1ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from einops import rearrange, repeat\n",
    "from transformers import CLIPProcessor, CLIPModel\n",
    "import kornia\n",
    "\n",
    "from jjuke_diffusion.unet_ldm.transformer import Encoder, TransformerWrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7e526e20-3ef2-4ae2-89a5-eae062e83c03",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AbstractEncoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    def encode(self, *args, **kwargs):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "eee2ffa4-a667-45c6-99b6-1421da32e26d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTTokenizer(AbstractEncoder):\n",
    "    \"\"\" Uses pre-trained BERT tokenizer from huggingface \"\"\"\n",
    "    def __init__(self, max_len=77, device=\"cuda\"):\n",
    "        super().__init__()\n",
    "\n",
    "        from transformers import BertTokenizerFast\n",
    "        self.tokenizer = BertTokenizerFast.from_pretrained(\"bert-base-uncased\")\n",
    "        self.max_len = max_len\n",
    "        self.device = device\n",
    "\n",
    "\n",
    "    def forward(self, text):\n",
    "        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_len,\n",
    "                                        return_length=True, return_overflowing_tokens=False,\n",
    "                                        padding=\"max_length\", return_tensors=\"pt\")\n",
    "        tokens = batch_encoding[\"input_ids\"].to(self.device)\n",
    "        return tokens\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def encode(self, text):\n",
    "        return self(text)\n",
    "\n",
    "    def decode(self, text):\n",
    "        return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "899fca91-0365-48be-8fda-a56d4636b7fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTEmbedder(AbstractEncoder):\n",
    "    \"\"\" Uses BERT tokenizer model and add some transformer encoder layers \"\"\"\n",
    "    def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77,\n",
    "                 use_tokenizer=True, emb_dropout=0.0, device=\"cuda\"):\n",
    "        super().__init__()\n",
    "\n",
    "        self.use_tokenizer = use_tokenizer\n",
    "        if self.use_tokenizer:\n",
    "            self.tokenize = BERTTokenizer(max_len=max_seq_len)\n",
    "\n",
    "        encoder = Encoder(dim=n_embed, depth=n_layer)\n",
    "        self.transformer = TransformerWrapper(\n",
    "            num_tokens=vocab_size, max_seq_len=max_seq_len, attn_layers=encoder, emb_dropout=emb_dropout\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, text):\n",
    "        if self.use_tokenizer:\n",
    "            tokens = self.tokenize(text) # .to(self.device)\n",
    "        else:\n",
    "            tokens = text\n",
    "\n",
    "        z = self.transformer(tokens, return_embeddings=True)\n",
    "        return z\n",
    "\n",
    "\n",
    "    def encode(self, text):\n",
    "        return self(text) # output length: 77"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d611f6aa-134b-47d1-bf5f-2168d5842e8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_encoder = BERTEmbedder(n_embed=1280, n_layer=32).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a3fb2f82-3c75-4652-b8fe-6d2aab61290e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 77, 1280])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_example = [\"a photo of a cat\", \"a photo of dog\"]\n",
    "text_encoder(text_example).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a161a849-7444-445e-bff0-50e012123073",
   "metadata": {},
   "source": [
    "## Domain specific encoder with CLIP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "19b20e95-31c9-47c3-ac5e-afeb659d170d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/sgtd/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "(…)t-large-patch14/resolve/main/config.json: 100%|█| 4.52\n",
      "`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config[\"id2label\"]` will be overriden.\n",
      "`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config[\"bos_token_id\"]` will be overriden.\n",
      "`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The value `text_config[\"eos_token_id\"]` will be overriden.\n",
      "model.safetensors: 100%|█| 1.71G/1.71G [01:34<00:00, 18.1\n",
      "(…)14/resolve/main/preprocessor_config.json: 100%|█| 316/\n",
      "(…)tch14/resolve/main/tokenizer_config.json: 100%|█| 905/\n",
      "(…)it-large-patch14/resolve/main/vocab.json: 100%|█| 961k\n",
      "(…)it-large-patch14/resolve/main/merges.txt: 100%|█| 525k\n",
      "(…)arge-patch14/resolve/main/tokenizer.json: 100%|█| 2.22\n",
      "(…)h14/resolve/main/special_tokens_map.json: 100%|█| 389/\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 480, 640])\n"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "from transformers import CLIPProcessor, CLIPModel\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "model = CLIPModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "\n",
    "img = Image.open(requests.get(\"http://images.cocodataset.org/val2017/000000039769.jpg\", stream=True).raw)\n",
    "\n",
    "print(ToTensor()(img).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7047f25a-cd4e-4652-8fa2-04930607546b",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = processor(text=[\"a photo of a cat\", \"a photo of a dog\"], images=img, return_tensors=\"pt\", padding=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "57ed6e97-f7aa-4e1b-ae2a-1b799277c655",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "keys of inputs:  dict_keys(['input_ids', 'attention_mask', 'pixel_values'])\n",
      "shape of input_ids:  torch.Size([2, 7])\n",
      "shape of attention_mask:  torch.Size([2, 7])\n",
      "shape of pixel_values:  torch.Size([1, 3, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "print(\"keys of inputs: \", inputs.keys())\n",
    "print(\"shape of input_ids: \", inputs[\"input_ids\"].shape)\n",
    "print(\"shape of attention_mask: \", inputs[\"attention_mask\"].shape)\n",
    "print(\"shape of pixel_values: \", inputs[\"pixel_values\"].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5f972b54-27ab-46e4-9ead-038b75bb92f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = model(**inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f0739890-9070-45ff-ae54-942b467d0143",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'transformers.models.clip.modeling_clip.CLIPOutput'>\n",
      "keys of outputs:  dict_keys(['loss', 'logits_per_image', 'logits_per_text', 'text_embeds', 'image_embeds', 'text_model_output', 'vision_model_output'])\n",
      "shape of loss:  None\n",
      "shape of logits_per_image:  torch.Size([1, 2])\n",
      "shape of logits_per_text:  torch.Size([2, 1])\n",
      "shape of image_embeds:  torch.Size([1, 768])\n",
      "shape of text_embeds:  torch.Size([2, 768])\n",
      "class of model output:  <class 'transformers.modeling_outputs.BaseModelOutputWithPooling'>\n"
     ]
    }
   ],
   "source": [
    "print(outputs.__class__)\n",
    "print(\"keys of outputs: \", vars(outputs).keys())\n",
    "print(\"shape of loss: \", vars(outputs)[\"loss\"])\n",
    "print(\"shape of logits_per_image: \", vars(outputs)[\"logits_per_image\"].shape)\n",
    "print(\"shape of logits_per_text: \", vars(outputs)[\"logits_per_text\"].shape)\n",
    "print(\"shape of image_embeds: \", vars(outputs)[\"image_embeds\"].shape)\n",
    "print(\"shape of text_embeds: \", vars(outputs)[\"text_embeds\"].shape)\n",
    "print(\"class of model output: \", vars(outputs)[\"vision_model_output\"].__class__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "17af05d5-3696-4132-b606-2f61b26b4b42",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output of CLIP vision model: \n",
      "last_hidden_state:  torch.Size([1, 257, 1024])\n",
      "pooler_output:  torch.Size([1, 257, 1024])\n",
      "hidden_states:  None\n",
      "attentions:  None\n",
      "\n",
      "Output of CLIP text model: \n",
      "last_hidden_state:  torch.Size([2, 7, 768])\n",
      "pooler_output:  torch.Size([2, 7, 768])\n",
      "hidden_states:  None\n",
      "attentions:  None\n"
     ]
    }
   ],
   "source": [
    "print(\"Output of CLIP vision model: \")\n",
    "print(\"last_hidden_state: \", vars(vars(outputs)[\"vision_model_output\"])[\"last_hidden_state\"].shape)\n",
    "print(\"pooler_output: \", vars(vars(outputs)[\"vision_model_output\"])[\"last_hidden_state\"].shape)\n",
    "print(\"hidden_states: \", vars(vars(outputs)[\"vision_model_output\"])[\"hidden_states\"])\n",
    "print(\"attentions: \", vars(vars(outputs)[\"vision_model_output\"])[\"attentions\"])\n",
    "\n",
    "print(\"\\nOutput of CLIP text model: \")\n",
    "print(\"last_hidden_state: \", vars(vars(outputs)[\"text_model_output\"])[\"last_hidden_state\"].shape)\n",
    "print(\"pooler_output: \", vars(vars(outputs)[\"text_model_output\"])[\"last_hidden_state\"].shape)\n",
    "print(\"hidden_states: \", vars(vars(outputs)[\"text_model_output\"])[\"hidden_states\"])\n",
    "print(\"attentions: \", vars(vars(outputs)[\"text_model_output\"])[\"attentions\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9816d9cc-86bd-4b8d-9b28-189ebf4b925b",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs_wo_img = processor(text=[\"a photo of a cat\", \"a photo of a dog\"], padding=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "45cfba62-4521-4e09-b007-3e77debed0cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "keys of inputs:  dict_keys(['input_ids', 'attention_mask'])\n",
      "input_ids:  [[49406, 320, 1125, 539, 320, 2368, 49407], [49406, 320, 1125, 539, 320, 1929, 49407]]\n",
      "attention_mask:  [[1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1]]\n"
     ]
    }
   ],
   "source": [
    "print(\"keys of inputs: \", inputs_wo_img.keys())\n",
    "print(\"input_ids: \", inputs_wo_img[\"input_ids\"])\n",
    "print(\"attention_mask: \", inputs_wo_img[\"attention_mask\"])\n",
    "# print(\"shape of pixel_values: \", inputs_wo_img[\"pixel_values\"].shape) → error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a2175c55-64f1-46e3-8fa4-50738b86fa3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[49406,   320,  1125,   539,   320,  2368, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407],\n",
      "        [49406,   320,  1125,   539,   320,  1929, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
      "         49407, 49407, 49407, 49407, 49407, 49407, 49407]])\n",
      "torch.Size([2, 77])\n"
     ]
    }
   ],
   "source": [
    "from transformers import CLIPTokenizerFast, CLIPTextModel\n",
    "\n",
    "tokenizer = CLIPTokenizerFast.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "model = CLIPTextModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "\n",
    "text_inputs = [\"a photo of a cat\", \"a photo of a dog\"]\n",
    "max_len = 77\n",
    "\n",
    "batch_encoding = tokenizer(text_inputs, truncation=True, max_length=max_len,\n",
    "                           return_length=True, return_overflowing_tokens=False,\n",
    "                           padding=\"max_length\", return_tensors=\"pt\")\n",
    "tokens = batch_encoding[\"input_ids\"]\n",
    "\n",
    "print(tokens)\n",
    "print(tokens.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e2bab3c4-92dd-4d36-aa15-e71a4e1c8100",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import CLIPTokenizerFast, CLIPTextModel\n",
    "tokenizer = CLIPTokenizerFast.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "\n",
    "text = \"a photo of a cat\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6928a6e6-476f-47c9-8fc0-581db7668b39",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CLIPTokenizer(AbstractEncoder):\n",
    "    \"\"\" Uses pre-trained CLIP tokenizer from huggingface \"\"\"\n",
    "    def __init__(self, max_len=77, device=\"cuda\"):\n",
    "        super().__init__()\n",
    "\n",
    "        from transformers import CLIPTokenizerFast\n",
    "        self.tokenizer = CLIPTokenizerFast.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "        self.max_len = max_len\n",
    "        self.device = device\n",
    "\n",
    "    \n",
    "    def forward(self, text):\n",
    "        batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_len, \n",
    "                                        return_length=True, return_overflowing_tokens=False,\n",
    "                                        padding=\"max_length\", return_tensors=\"pt\")\n",
    "        tokens = batch_encoding[\"input_ids\"].to(self.device)\n",
    "        return tokens\n",
    "\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def encode(self, text):\n",
    "        return self(text)\n",
    "\n",
    "\n",
    "    def decode(self, text):\n",
    "        return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0044d67a-be82-4ad4-a60f-ccd74b83ec5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import CLIPTextConfig, CLIPTextModel\n",
    "\n",
    "clip_text_model = CLIPTextModel(CLIPTextConfig(vocab_size=49408, max_position_embeddings=77, hidden_size=1280, intermediate_size=1280, projection_dim=1280, num_hidden_layers=32, attention_dropout=0.0)).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "cc990ea5-4c14-4eba-8c6f-489d0aba55a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def count_params(model):\n",
    "    return sum(p.numel() for p in model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "dfe39688-911c-4a89-88e7-e98c61c6503f",
   "metadata": {},
   "outputs": [],
   "source": [
    "last_hidden_state, pooler_output = clip_text_model(CLIPTokenizer(max_len=77)(text_example), return_dict=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9ce9595e-11d4-46a7-b8ad-cff693fb44db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 77, 1280])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "last_hidden_state.shape # same to output of BERT Embedder (LDM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a7171c0e-ff21-454e-8e0c-0b429212012b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 1280])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pooler_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "fe3a08ff-7e91-4cfb-b2ed-67692d9a5c2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CLIPEmbedder(AbstractEncoder):\n",
    "    \"\"\" Uses CLIP tokenizer and add some transformer encoder layers \"\"\"\n",
    "    def __init__(self, n_embed, n_layer, vocab_size=49408, max_seq_len=77,\n",
    "                 use_tokenizer=True, emb_dropout=0.0, device=\"cuda\"):\n",
    "        super().__init__()\n",
    "\n",
    "        self.use_tokenizer = use_tokenizer\n",
    "        if self.use_tokenizer:\n",
    "            self.tokenize = CLIPTokenizer(max_len=max_seq_len)\n",
    "\n",
    "        encoder = Encoder(dim=n_embed, depth=n_layer)\n",
    "        self.transformer = TransformerWrapper(\n",
    "            num_tokens=vocab_size, max_seq_len=max_seq_len, attn_layers=encoder, emb_dropout=emb_dropout\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, text):\n",
    "        if self.use_tokenizer:\n",
    "            tokens = self.tokenize(text) # .to(self.device)\n",
    "        else:\n",
    "            tokens = text\n",
    "\n",
    "        z = self.transformer(tokens, return_embeddings=True)\n",
    "        return z\n",
    "\n",
    "\n",
    "    def decode(self, text):\n",
    "        return self(text) # output length: 77\n",
    "\n",
    "\n",
    "class CLIPEmbedderHF(AbstractEncoder):\n",
    "    \"\"\" Uses CLIP tokenizer and CLIP Text Model for the text encoder \"\"\"\n",
    "    def __init__(self, n_embed, n_layer, vocab_size=49408, max_seq_len=77,\n",
    "                 use_tokenizer=True, emb_dropout=0.0, device=\"cuda\"):\n",
    "        super().__init__()\n",
    "\n",
    "        from transformers import CLIPTextConfig, CLIPTextModel\n",
    "\n",
    "        self.use_tokenizer = use_tokenizer\n",
    "        if self.use_tokenizer:\n",
    "            self.tokenize = CLIPTokenizer(max_len=max_seq_len)\n",
    "\n",
    "        self.transformer = CLIPTextModel(CLIPTextConfig(\n",
    "            vocab_size=vocab_size, max_position_embeddings=max_seq_len, hidden_size=n_embed,\n",
    "            intermediate_size=n_embed, projection_dim=n_embed, num_hidden_layers=n_layer,\n",
    "            attention_dropout=emb_dropout\n",
    "        ))\n",
    "\n",
    "\n",
    "    def forward(self, text):\n",
    "        if self.use_tokenizer:\n",
    "            tokens = self.tokenize(text) # .to(self.device)\n",
    "        else:\n",
    "            tokens = text\n",
    "\n",
    "        z, _ = self.transformer(tokens, return_dict=False)\n",
    "        return z\n",
    "\n",
    "\n",
    "    def decode(self, text):\n",
    "        return self(text) # output length: 77"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f0337f02-131b-4188-ba64-6db64e85fb47",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of params of BERTEmbedder:  581994042\n",
      "Number of params of CLIPEmbedder:  630361088\n",
      "Number of params of CLIPEmbedderHF:  378325760\n"
     ]
    }
   ],
   "source": [
    "# compare number of parameters of CLIPEmbedder, CLIPEmbedderHF, and BERTEmbedder\n",
    "\n",
    "bert_embedder = BERTEmbedder(1280, 32).cuda()\n",
    "clip_embedder = CLIPEmbedder(1280, 32).cuda()\n",
    "clip_embedder_hf = CLIPEmbedderHF(1280, 32).cuda()\n",
    "print(\"Number of params of BERTEmbedder: \", count_params(bert_embedder.transformer))\n",
    "print(\"Number of params of CLIPEmbedder: \", count_params(clip_embedder.transformer))\n",
    "print(\"Number of params of CLIPEmbedderHF: \", count_params(clip_embedder_hf.transformer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "0c6fa8b7-856c-4699-bd27-dbc5b2faeb97",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 77, 1280])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bert_embedder(text_example).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2b9f8e53-d049-4d45-b9bb-db147b2f74dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 77, 1280])"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clip_embedder(text_example).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "0bb5bdf7-4375-4219-aa4f-77ce3bbff534",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 77, 1280])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clip_embedder_hf(text_example).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "331d03d7-2c2e-4708-9309-3b2e74de4406",
   "metadata": {},
   "source": [
    "# LDMTrainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a8609d0-e22a-4e90-9425-8782500b2a18",
   "metadata": {},
   "source": [
    "$$ \\mathcal{L}_\\text{LDM} := \\mathbb{E}_{\\mathbf{y}, \\epsilon \\sim \\mathcal{N}(0, 1), t} \\left[ \\left\\Vert \\epsilon - \\epsilon_\\theta(\\mathbf{z}_t, t, \\tau_\\theta(\\mathbf{y})) \\right\\Vert_2^2 \\right] $$\n",
    "\n",
    "- Input:\n",
    "    - x: Feature map\n",
    "    - c: Condition\n",
    "    - ...\n",
    "- Output:\n",
    "    - Losses dictinoary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "001364f5-f06a-4cb5-93a7-c0953e23b0d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jjuke_diffusion.diffusion.common import get_betas\n",
    "from jjuke_diffusion.diffusion.ddpm import DDPMTrainer\n",
    "from jjuke_diffusion.unet_cond.unet import UNetModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "15997561-8d70-4706-9135-b68bbab353b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# DDPM Trainer args\n",
    "args_model_mean_type = \"eps\"\n",
    "args_model_var_type = \"fixed_small\"\n",
    "args_loss_type = \"l2\"\n",
    "\n",
    "betas = get_betas(\"linear\", 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "ad50e96b-0fb1-40b3-9a05-0af9ae401316",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = DDPMTrainer(\n",
    "    betas,\n",
    "    model_mean_type=args_model_mean_type,\n",
    "    model_var_type=args_model_var_type,\n",
    "    loss_type=args_loss_type,\n",
    "    clip_denoised=False # TODO: check if it is True\n",
    ").cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "a601e05c-cbd6-484e-9955-10215b1db981",
   "metadata": {},
   "outputs": [
    {
     "ename": "OutOfMemoryError",
     "evalue": "CUDA out of memory. Tried to allocate 114.00 MiB (GPU 0; 23.62 GiB total capacity; 21.67 GiB already allocated; 111.75 MiB free; 22.13 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[46], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m model \u001b[38;5;241m=\u001b[39m \u001b[43mUNetModel\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      2\u001b[0m \u001b[43m    \u001b[49m\u001b[43munet_dim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[43min_channels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m96\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43mout_channels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m96\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel_channels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m320\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      6\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattention_resolutions\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m16\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m# downsampling factor = spatial resolution (h, w of feature maps) / CA resoliutions (32, 16, 8 → from LDM paper)\u001b[39;49;00m\n\u001b[1;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[43mchannel_mult\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m4\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43mnum_heads\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m8\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      9\u001b[0m \u001b[43m    \u001b[49m\u001b[43muse_spatial_transformer\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtransformer_depth\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     11\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcontext_dim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1280\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     12\u001b[0m \u001b[43m    \u001b[49m\u001b[43mattention_type\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mxformers\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m     13\u001b[0m \u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/envs/sgtd/lib/python3.9/site-packages/torch/nn/modules/module.py:905\u001b[0m, in \u001b[0;36mModule.cuda\u001b[0;34m(self, device)\u001b[0m\n\u001b[1;32m    888\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcuda\u001b[39m(\u001b[38;5;28mself\u001b[39m: T, device: Optional[Union[\u001b[38;5;28mint\u001b[39m, device]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    889\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Moves all model parameters and buffers to the GPU.\u001b[39;00m\n\u001b[1;32m    890\u001b[0m \n\u001b[1;32m    891\u001b[0m \u001b[38;5;124;03m    This also makes associated parameters and buffers different objects. So\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    903\u001b[0m \u001b[38;5;124;03m        Module: self\u001b[39;00m\n\u001b[1;32m    904\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 905\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mlambda\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/conda/envs/sgtd/lib/python3.9/site-packages/torch/nn/modules/module.py:797\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    795\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_apply\u001b[39m(\u001b[38;5;28mself\u001b[39m, fn):\n\u001b[1;32m    796\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 797\u001b[0m         \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    799\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    800\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    801\u001b[0m             \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    802\u001b[0m             \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    807\u001b[0m             \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    808\u001b[0m             \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/opt/conda/envs/sgtd/lib/python3.9/site-packages/torch/nn/modules/module.py:797\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    795\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_apply\u001b[39m(\u001b[38;5;28mself\u001b[39m, fn):\n\u001b[1;32m    796\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 797\u001b[0m         \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    799\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    800\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    801\u001b[0m             \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    802\u001b[0m             \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    807\u001b[0m             \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    808\u001b[0m             \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "    \u001b[0;31m[... skipping similar frames: Module._apply at line 797 (2 times)]\u001b[0m\n",
      "File \u001b[0;32m/opt/conda/envs/sgtd/lib/python3.9/site-packages/torch/nn/modules/module.py:797\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    795\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_apply\u001b[39m(\u001b[38;5;28mself\u001b[39m, fn):\n\u001b[1;32m    796\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 797\u001b[0m         \u001b[43mmodule\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_apply\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfn\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    799\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    800\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    801\u001b[0m             \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    802\u001b[0m             \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    807\u001b[0m             \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    808\u001b[0m             \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/opt/conda/envs/sgtd/lib/python3.9/site-packages/torch/nn/modules/module.py:820\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    816\u001b[0m \u001b[38;5;66;03m# Tensors stored in modules are graph leaves, and we don't want to\u001b[39;00m\n\u001b[1;32m    817\u001b[0m \u001b[38;5;66;03m# track autograd history of `param_applied`, so we have to use\u001b[39;00m\n\u001b[1;32m    818\u001b[0m \u001b[38;5;66;03m# `with torch.no_grad():`\u001b[39;00m\n\u001b[1;32m    819\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m--> 820\u001b[0m     param_applied \u001b[38;5;241m=\u001b[39m \u001b[43mfn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparam\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    821\u001b[0m should_use_set_data \u001b[38;5;241m=\u001b[39m compute_should_use_set_data(param, param_applied)\n\u001b[1;32m    822\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m should_use_set_data:\n",
      "File \u001b[0;32m/opt/conda/envs/sgtd/lib/python3.9/site-packages/torch/nn/modules/module.py:905\u001b[0m, in \u001b[0;36mModule.cuda.<locals>.<lambda>\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m    888\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcuda\u001b[39m(\u001b[38;5;28mself\u001b[39m: T, device: Optional[Union[\u001b[38;5;28mint\u001b[39m, device]] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m T:\n\u001b[1;32m    889\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124;03m\"\"\"Moves all model parameters and buffers to the GPU.\u001b[39;00m\n\u001b[1;32m    890\u001b[0m \n\u001b[1;32m    891\u001b[0m \u001b[38;5;124;03m    This also makes associated parameters and buffers different objects. So\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    903\u001b[0m \u001b[38;5;124;03m        Module: self\u001b[39;00m\n\u001b[1;32m    904\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 905\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_apply(\u001b[38;5;28;01mlambda\u001b[39;00m t: \u001b[43mt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcuda\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m)\n",
      "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 114.00 MiB (GPU 0; 23.62 GiB total capacity; 21.67 GiB already allocated; 111.75 MiB free; 22.13 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
     ]
    }
   ],
   "source": [
    "model = UNetModel(\n",
    "    unet_dim=2,\n",
    "    in_channels=96,\n",
    "    out_channels=96,\n",
    "    model_channels=320,\n",
    "    attention_resolutions=[16, 8, 4], # downsampling factor = spatial resolution (h, w of feature maps) / CA resoliutions (32, 16, 8 → from LDM paper)\n",
    "    channel_mult=[1, 2, 4, 4],\n",
    "    num_heads=8,\n",
    "    use_spatial_transformer=True,\n",
    "    transformer_depth=1,\n",
    "    context_dim=1280,\n",
    "    attention_type=\"xformers\",\n",
    ").cuda()\n",
    "# lr: 1e-4\n",
    "# iterations: 390,000\n",
    "# batch_size: 680"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19be254e-8e64-407f-b176-0a3a987be1c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from einops import rearrange, repeat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f14b425-511a-4d0f-a4fe-52f58832bdef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Feature maps (3 triplanes)\n",
    "x = torch.rand(2, 128, 128, 96).cuda() # (b, h, w, c)\n",
    "t = torch.randint(0, 1000, (2,)).cuda()\n",
    "\n",
    "# Conditions\n",
    "cond = bert_embedder(text_example).cuda()\n",
    "print(cond.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "974773d3-7a64-4199-9cff-c57367eaa149",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_input function of LDM official code -> get_input(batch, \"caption\").to(self.device)\n",
    "x = rearrange(x, \"b h w c -> b c h w\").contiguous().float()\n",
    "cond = rearrange(cond, \"b n c -> b c n\").contiguous().float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c83bb5b5-aca6-4679-bf42-ede09790ef89",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Shape of feature maps: \", x.shape)\n",
    "print(\"Shape of text embedding: \", cond.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d6c5354-6d1f-4fd3-bb17-31afa791961b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model(x, t, cond).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae846cbb-1976-4089-9193-0d9fd041e080",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_maps = torch.rand(2, 128, 128, 96).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99f1db40-13ab-4391-901b-3de5c79817b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = trainer(model, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa9a9479-a6ff-4ca2-bf1b-383dffb40043",
   "metadata": {},
   "outputs": [],
   "source": [
    "losses"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "230e83dd-e812-4130-9b5b-eb1d265a7236",
   "metadata": {},
   "source": [
    "# LDMSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "58201465-42d7-4537-9470-f41f0b19f58c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jjuke_diffusion.diffusion.ddim import DDIMSampler\n",
    "from jjuke_diffusion.diffusion.karras import KarrasSampler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93a267ef-3463-46bc-b614-c1d396a56da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# args for samplers\n",
    "args_n_sampler_steps = 50\n",
    "args_eta = 0.\n",
    "\n",
    "args_model_mean_type = \"eps\"\n",
    "args_model_var_type = \"fixed_small\"\n",
    "args_loss_type = \"l2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b67691-9810-446e-8f33-697108296c6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddim_sampler = DDIMSampler(\n",
    "    betas,\n",
    "    ddim_num_timesteps=50,\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24c2c200-5ec4-48ed-9cdf-9d1ced367951",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sgtd",
   "language": "python",
   "name": "sgtd"
  },
  "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}