{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "116eb01e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac6c81923ae642dda82e7be3d8b8172c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/9 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ad5c1247f56746ea842c366a5a81eaf7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b31d22599263437fa2032cb7eef468b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You set `add_prefix_space`. The tokenizer needs to be converted from the slow tokenizers\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import os\n",
    "model_id = \"stabilityai/stable-diffusion-3.5-large-turbo\"\n",
    "from src.sd3_pipeline import VSFStableDiffusion3Pipeline\n",
    "pipe = VSFStableDiffusion3Pipeline.from_pretrained(\n",
    "    \"stabilityai/stable-diffusion-3.5-large-turbo\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    ") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e61d2537",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "negative_prompt_image = Image.open(\"wheel.webp\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0dd57774",
   "metadata": {},
   "outputs": [],
   "source": [
    "import open_clip\n",
    "\n",
    "model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms('hf-hub:laion/CLIP-ViT-g-14-laion2B-s34B-b88K')\n",
    "tokenizer = open_clip.get_tokenizer('hf-hub:laion/CLIP-ViT-g-14-laion2B-s34B-b88K')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0b8a3ccc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import clip\n",
    "from PIL import Image\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model_l, preprocess = clip.load(\"ViT-L/14\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c169092b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = model.to(device)\n",
    "image = preprocess(negative_prompt_image).unsqueeze(0).to(device)\n",
    "with torch.no_grad():\n",
    "    image_features_l = model_l.encode_image(image)\n",
    "    \n",
    "image = preprocess_val(negative_prompt_image).unsqueeze(0).to(device)\n",
    "with torch.no_grad():\n",
    "    image_features_g = model.encode_image(image) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a4d7d1bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_feature = pipe.encode_prompt('a wheel', 'a wheel', \"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "71778ceb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 333, 4096])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text_feature[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f1af1d43",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_embed, pooled_prompt_embed = pipe._get_clip_prompt_embeds(\n",
    "    prompt=\"a wheel\",\n",
    "    clip_model_index=0,\n",
    ")\n",
    "prompt_embed_2, pooled_prompt_embed_2 = pipe._get_clip_prompt_embeds(\n",
    "    prompt=\"a wheel\",\n",
    "    clip_model_index=1,\n",
    ")\n",
    "clip_prompt_embeds = torch.cat([prompt_embed, prompt_embed_2], dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a2dd5734",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 77, 768]), torch.Size([1, 768]))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt_embed.shape, image_features_l.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "deba7b67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 77, 1280]), torch.Size([1, 1024]))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt_embed_2.shape, image_features_g.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "19a4af36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 4096])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "feature = torch.cat([image_features_l], dim=1)\n",
    "feature = torch.cat([feature, torch.zeros(1, 4096 - 768).to(device)], dim=1)\n",
    "feature.shape "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5a69baa6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.0525], device='cuda:0', grad_fn=<SumBackward1>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.nn.functional.cosine_similarity(feature, text_feature[0][:,:77].mean(1).to(device)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "faf28cbe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([1, 77, 2048]), torch.Size([1, 1792]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clip_prompt_embeds.shape, torch.cat([image_features_l, image_features_g], dim=1).shape "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "44e2cd9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.0223], device='cuda:0', grad_fn=<SumBackward1>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.nn.functional.cosine_similarity(feature.mean(1), text_feature[0][:,:77].mean(1).to(device)) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "neg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}