4M_Tutorial/Part_01_4M_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cd60a039-9bbe-4b7e-a649-52f126112b46",
   "metadata": {},
   "source": [
    "# **Communications Course**  \n",
    "\n",
    "## **Exploring Multi-Modal Generation with Foundation Models**  \n",
    "\n",
    "In this exercise, you will explore the multi-modal generation capabilities of a multimodal foundation model. We base this tutorial on [4M: Massively Multimodal Masked Modeling](https://4m.epfl.ch/), the model you will later build and train yourself from scratch as part of the `nano4M` project. This tutorial provides an overview of the capabilities of a pre-trained multimodal foundation model and may inspire possible extensions you would like to implement later. While we focus on 4M, many concepts introduced in this project apply to a broader class of multimodal foundation models.\n",
    "\n",
    "## **What is a Foundation Model?**  \n",
    "\n",
    "[Foundation models](https://arxiv.org/abs/2108.07258) are large-scale deep neural networks trained on massive web-scale datasets, often covering multiple tasks in their pretraining. Due to their exposure to diverse and extensive data, these models are inherently multi-task and capable of generalizing across different distributions.  \n",
    "\n",
    "Based on their capabilities and the tasks they can perform, foundation models can be broadly categorized into **three** types:\n",
    "\n",
    "### **Large Language Models (LLMs)**  \n",
    "\n",
    "Large Language Models (LLMs) are deep learning models designed to process and generate text. Trained on vast amounts of text data, these models learn complex linguistic patterns, enabling them to perform various natural language processing (NLP) tasks such as text generation, summarization, translation, and reasoning. Notable Large Language Models include:  \n",
    "\n",
    "- [GPT-4o](https://openai.com/index/hello-gpt-4o/)  \n",
    "- [GPT-3.5-turbo](https://openai.com/index/openai-o3-mini/)  \n",
    "- [DeepSeek R1](https://api-docs.deepseek.com/news/news250120)  \n",
    "- [Claude](https://www.anthropic.com/news/claude-3-7-sonnet)  \n",
    "\n",
    "### **Vision Foundation Models (VFMs)**  \n",
    "\n",
    "Vision foundation models are designed primarily for vision-centric tasks, which are often discriminative rather than generative. These models excel in open-vocabulary object recognition, multi-modal retrieval, and segmentation. Some key VFMs include:  \n",
    "\n",
    "- **CLIP** – [*Learning Transferable Visual Models from Natural Language Supervision*  ](https://arxiv.org/abs/2103.00020)\n",
    "  CLIP is a vision-language model that learns visual representations from natural language supervision. Without additional training, it can perform zero-shot tasks such as object recognition, segmentation, and image-to-text and text-to-image retrieval.  \n",
    "\n",
    "- **SAM** – [*Segment Anything*](https://arxiv.org/abs/2304.02643)\n",
    "  SAM is a foundation model for image segmentation, capable of segmenting any object in an image based on user input. You can try the online demo [here](https://sam2.metademolab.com/demo).  \n",
    "\n",
    "### **Multi-Modal Any-to-Any Foundation Models**  \n",
    "\n",
    "This category includes foundation models inherently designed to process multi-modal inputs while simultaneously generating multi-modal outputs. Imagine a model that can process various inputs such as images, text, and depth while also generating outputs like images, surface normals, and captions. Some notable examples include:  \n",
    "\n",
    "- **4M** – [*Massively Multimodal Masked Modeling*](https://4m.epfl.ch/)\n",
    "  4M is a multimodal framework that integrates multiple modalities, such as images, text, and depth maps, through a masked modeling pre-training objective, enhancing multi-modal understanding and generation. This tutorial will focus specifically on 4M.  \n",
    "\n",
    "- **Janus Pro** [Unifying multimodal understanding and generation](https://arxiv.org/abs/2501.17811)\n",
    "  Janus is a specialist model that integrates vision and language for both understanding and generation tasks. For instance, you can provide Janus with an image and ask it to describe it, or you can prompt it to generate an image based on your preferences. You can explore the online demo [here](https://huggingface.co/spaces/deepseek-ai/Janus-Pro-7B).  \n",
    "\n",
    "Beyond these, many other any-to-any multi-modal foundation models exist, such as:  \n",
    "\n",
    "- [Unified-IO 2](https://arxiv.org/abs/2312.17172)  \n",
    "- [Emu3](https://arxiv.org/pdf/2409.18869)  \n",
    "\n",
    "\n",
    "<!-- ## **References**  \n",
    "\n",
    "1. OpenAI, GPT-4 Technical Report\n",
    "2. Radford et al., *Learning Transferable Visual Models From Natural Language Supervision*\n",
    "3. Kirillov et al., *Segment Anything*\n",
    "4. Mizrahi and Bachman et al., *4M: Massively Multimodal Masked Modeling*\n",
    "5.  -->\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8fabaeb-ed2d-4800-894a-2237701ae376",
   "metadata": {},
   "source": [
    "# Playing with 4M: An Any-to-Any Vision Model for Multi-Modal Generation\n",
    "\n",
    "In the first part of this exercise, you will perform a hands-on exploration of a pre-trained 4M model. Before delving into the hands-on part, you will first familiarize yourself with the basic building blocks of 4M, which will be useful to learn and understand for subsequent weeks' labs and exercises.\n",
    "\n",
    "## About 4M:\n",
    "\n",
    "4M, which stands for Massively Multimodal Masked Modeling, is a transformer-based encoder-decoder deep neural network trained on paired multi-modal data using a masked-modeling training objective. This design enables the model to handle various data modalities, including text, images, geometric data, and semantic information. Below, we discuss the primary components of the 4M model:\n",
    "\n",
    "### 1. Model Architecture\n",
    "\n",
    "The 4M model utilizes a unified Transformer encoder-decoder architecture as shown in the figure below (right). This architecture is designed to process and generate multiple data modalities, allowing it to learn generative relationships between different types of data. By converting all modalities into discrete tokens (below Fig. left), also called tokenization, 4M ensures compatibility across various data types, facilitating efficient learning and generation.\n",
    "\n",
    "<img src=\"assets/4M_architecture.png\" width=\"1000\"/>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15703a26-b7a9-4b64-891a-441322ade45a",
   "metadata": {},
   "source": [
    "### 2. Training Data and Tokenization\n",
    "\n",
    "The latest variant of 4M, called 4M-21 [1], is trained on 21 diverse sets of paired multi-modal data. This includes RGB, depth, surface normals, semantic segmentation, SAM instances, Canny & SAM edges, 3D human poses, CLIP-B/16 features, DINOv2-B/14 features, ImageBind-H/14 features, captions, metadata, color palette, and bounding boxes. Training 4M on such data allows it to learn the relationships between different modalities to generate coherent and contextually relevant outputs across various data types.\n",
    "\n",
    "**Tokenization:** As seen in the left subplot of the previous figure, different data types have different formats, and it is not trivial to unify these different data modalities for training. For example, the caption is made up of 1D sequences of text, while CLIP features are 2D high-dimensional feature grids. 4M achieves unification of these different modalities through a procedure called tokenization. Tokenization converts the different data modalities into discrete sets of integer values (usually called tokens), which allows for seamless training.\n",
    "\n",
    "We use dedicated models to tokenize the data, which are called tokenizers. Different tokenizers used by 4M are shown in the diagram below:\n",
    "\n",
    "<img src=\"assets/4M_tokenization.png\" width=\"1000\"/>\n",
    "\n",
    "1. Bachman and Kar et al., *4M-21: An Any-to-Any Vision Model for Tens of Tasks and Modalities*\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2bff633-458c-4454-8671-9d66477db32d",
   "metadata": {},
   "source": [
    "### 3. Training Objective\n",
    "\n",
    "The training objective of 4M is based on masked modeling (e.g., like [BERT](https://arxiv.org/abs/1810.04805) or [MaskGIT](https://arxiv.org/abs/2202.04200)). During training, a portion of all tokens are masked-out, and the model is tasked with predicting these masked tokens from the non-masked ones. In our case, we extend the masking objective to include various modalities, i.e. perform [multimodal masked modeling](https://multimae.epfl.ch/), where the model is tasked to predict both across and within modalities. This approach encourages the model to develop a deep understanding of the underlying structures and relationships within and between modalities, leading to improved understanding and generative capabilities."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaa77721-68cf-464a-9efc-6cf22f142779",
   "metadata": {},
   "source": [
    "### 4. Generative AI Capabilities of 4M\n",
    "\n",
    "After pretraining, the 4M models can perform generation from any subset of modalities to any other set of modalities. For example, if you have an RGB image, you can generate captions and depth maps seamlessly, and vice versa.\n",
    "\n",
    "If our input consists of just one modality, we can generate an ((N x N) - N) matrix by performing generation from all N modalities to the other N-1 modalities, as shown in the figure below. Note that the order of modality generation matters, and we have a large number of possible chaining orders (factorial of N, N!). Also note that this matrix only shows generations from a single, full input, but generation can be performed even given multiple, partial modalities.\n",
    "\n",
    "\n",
    "<img src=\"assets/4m_matrix.png\" width=\"1000\"/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc9ff7c2-3004-4184-a374-21790708fc2f",
   "metadata": {},
   "source": [
    "# **Getting Started with 4M-21: An Any-to-Any Vision Model for Tens of Tasks and Modalities**\n",
    "\n",
    "Now that you are familiar with the 4M framework, it's time to have some fun and practically perform some cool generations! In this hands-on session, you will explore interactive demos and generate different modalities yourself.\n",
    "\n",
    "### **What We Will Explore**\n",
    "\n",
    "We will experiment with 4M models in the following ways:\n",
    "\n",
    "1. **Understanding and Quick-Start with Tokenizers**\n",
    "2. **RGB to Any-Other-Modality Generation**\n",
    "3. **Improved Generations Using Chaining**\n",
    "4. **Any to RGB Modality Generation**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d6590b-4b0c-46f1-accd-ec002b7b53a9",
   "metadata": {},
   "source": [
    "#### Setting Up the Coding Environment\n",
    "\n",
    "In order to successfully run this code, we need to first install the required packages. Make sure you have followed the installation instructions in `4M_Tutorial/README.md`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e2404be-b57a-433f-9640-6484b9f4238e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2ac4733-1f25-421d-9dec-88163c7f3cd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "import torch\n",
    "from torchvision.transforms.functional import center_crop\n",
    "from tokenizers import Tokenizer\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Video related: \n",
    "import cv2\n",
    "from IPython.display import Video\n",
    "\n",
    "from fourm.data.multimodal_dataset_folder import MultiModalDatasetFolder\n",
    "from fourm.models.fm import FM\n",
    "from fourm.vq.vqvae import VQVAE, DiVAE\n",
    "from fourm.models.generate import GenerationSampler, build_chained_generation_schedules, init_empty_target_modality, init_full_input_modality, custom_text\n",
    "# from utils.generation_abstract_functions import create_generation_schedule_rgb_to_others\n",
    "from fourm.data.modality_transforms import RGBTransform, DepthTransform, MetadataTransform\n",
    "from fourm.data.modality_info import MODALITY_INFO, MODALITY_TRANSFORMS\n",
    "from fourm.utils.plotting_utils import decode_dict, visualize_bboxes, plot_text_in_square\n",
    "from fourm.utils import denormalize, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD, IMAGENET_DEFAULT_STD, IMAGENET_DEFAULT_MEAN\n",
    "from fourm.data.modality_transforms import SemsegTransform\n",
    "from fourm.data.image_augmenter import CenterCropImageAugmenter\n",
    "from torchvision import transforms\n",
    "from fourm.data.modality_transforms import UnifiedDataTransform\n",
    "from fourm.data.dataset_utils import SubsampleDatasetWrapper\n",
    "from fourm.data.masking import UnifiedMasking\n",
    "from einops import rearrange\n",
    "from utils.semseg_helper_utils import semseg_to_rgb, plot_rgb2semseg, get_dataset, get_semseg_metrics, total_intersect_and_union, intersect_and_union, mean_iou, mean_dice, eval_metrics, tokens_per_target_dict, autoregression_schemes_dict, cfg_schedules_dict\n",
    "from tqdm import tqdm\n",
    "import matplotlib.colors as mcolors\n",
    "\n",
    "# The flag below controls whether to allow TF32 on matmul. This flag defaults to False in PyTorch 1.12 and later.\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.\n",
    "torch.backends.cudnn.allow_tf32 = True\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "torch.set_grad_enabled(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9b90650-a8d8-4b40-ac14-43ed92d4b709",
   "metadata": {},
   "source": [
    "# 1. **Unifying Data Format Using Tokenization**\n",
    "\n",
    "First, we will understand the fundamentals of using tokenizers. Recall that 4M uses different tokenizer models that convert each data type (like text, image, depth map, etc.) into discrete tokens, which are then used as input to 4M models.\n",
    "\n",
    "To encode modalities into discrete tokens and, in turn, decode the tokens that 4M predicts, we use modality-specific tokenizers. We can easily load them from the [Hugging Face hub](https://huggingface.co/EPFL-VILAB) with the following lines.\n",
    "\n",
    "The tokenizer checkpoint names are formatted as: `f'4M_tokenizers_{modality}_{vocab_size}_{min_res}-{max_res}'`.\n",
    "\n",
    "1. **Modality**: Can be any data type such as edges, semsegs, RGB image, etc.\n",
    "2. **vocab_size**: Vocabulary size denotes the capacity or number of unique tokens reserved by the tokenizer to represent the data type. If the vocabulary size is 2, it means the tokenizer can represent data using only two different tokens (i.e., token 0 and token 1). Increasing the vocabulary size usually improves the quality of the tokenizer but at the cost of more computations.\n",
    "3. **min_res**: Minimum resolution for the input data that the tokenizer was trained on.\n",
    "4. **max_res**: Maximum resolution for the input data that the tokenizer was trained on.\n",
    "\n",
    "All tokenizers here are trained to work on resolutions between 224 and 448, in steps of 32. Below, we first load all the pretrained tokenizers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ed9b32-726b-422a-98be-3d948fb22252",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_tok = Tokenizer.from_file('toks/text_tokenizer_4m_wordpiece_30k.json')\n",
    "\n",
    "toks = {\n",
    "    'tok_rgb': DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_rgb_16k_224-448').eval().to(device),\n",
    "    'tok_depth': DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_depth_8k_224-448').eval().to(device),\n",
    "    'tok_normal': DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_normal_8k_224-448').eval().to(device),\n",
    "    'tok_canny_edge': DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_edge_8k_224-512').eval().to(device),\n",
    "    'tok_sam_edge': DiVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_edge_8k_224-512').eval().to(device),\n",
    "    'tok_semseg': VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_semseg_4k_224-448').eval().to(device),\n",
    "    'tok_clip': VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_CLIP-B16_8k_224-448').eval().to(device),\n",
    "    'tok_dinov2': VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_DINOv2-B14_8k_224-448').eval().to(device),\n",
    "    'tok_imagebind': VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_ImageBind-H14_8k_224-448').eval().to(device),\n",
    "    'sam_instance': VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_sam-instance_1k_64').eval().to(device),\n",
    "    'human_poses': VQVAE.from_pretrained('EPFL-VILAB/4M_tokenizers_human-poses_1k_8').eval().to(device),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ca27a0d-a3d9-4668-8192-69645a1096d3",
   "metadata": {},
   "source": [
    "We will now explore how to use a tokenizer to encode input data and then decode it back. We will RGB image as the data modality.\n",
    "\n",
    "### **Using RGB Tokenizer to Encode and Decode an Image**:\n",
    "\n",
    "The main modality used in 4M is the RGB image. RGB images are the primary input used in major computer vision tasks such as object recognition, image segmentation, and object detection.\n",
    "\n",
    "Let's first load an input image, which we will use in our experimentation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea888a7d-03e3-4bae-9def-926842194399",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download image from the specified URL and preprocess it Hint: you can also load any custom image that you like.\n",
    "image_url = 'https://storage.googleapis.com/four_m_site/images/demo_rgb.png'\n",
    "!curl $image_url --output input.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69722209-f0c6-4078-b0a5-70d19ee96590",
   "metadata": {},
   "source": [
    "Let's perform basic preprocessing to transform the input image into a suitable format for the tokenizer encoder.\n",
    "\n",
    "For each input data type, we usually perform three steps:\n",
    "\n",
    "1. **Loading the Data**\n",
    "2. **Preprocessing**\n",
    "3. **Postprocessing**\n",
    "\n",
    "**Recommendation:** Look into the function definitions of these three steps that can be found in `RGBTransform` class "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16933c25-7cbb-4cf6-9c06-6ab5db4e5e3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initalizing the RGB transform class\n",
    "rgb_transform = RGBTransform(imagenet_default_mean_and_std=False)\n",
    "img_pil = rgb_transform.load('./input.jpg')\n",
    "img_pil = rgb_transform.preprocess(img_pil)\n",
    "img_pil = center_crop(img_pil, (min(img_pil.size), min(img_pil.size))).resize((224,224))\n",
    "img = rgb_transform.postprocess(img_pil).unsqueeze(0).to(device)\n",
    "# verfiy the normalization\n",
    "plt.imshow(denormalize(img, mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD)[0].cpu().permute(1, 2, 0))\n",
    "print(f\"Shape of the input image is: {img.shape}, \\n that stands for [Batch Size, Channels, Height, Width]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa72adf-eb4a-48ef-9a7a-725ee7e6ec8b",
   "metadata": {},
   "source": [
    "Now lets encode the input image using the RGB tokenizer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e62bdaa-b54c-4217-9ab2-85144bd0f39a",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_rgb = toks['tok_rgb'].tokenize(img)\n",
    "print(f\"We have tokenized the Image input which results in discrete tokens. The shape of the tokenized tensor is f{tokenized_rgb.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf91e3ac-c0aa-4869-beaa-a5ca058ea45b",
   "metadata": {},
   "source": [
    "Note that we can also say that the tokenized representation of the data type compresses the data present in its original format. 4M uses such tokenized representation of the data in its training."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56b0b600-ce19-4aa2-82f6-9e245e76c5f7",
   "metadata": {},
   "source": [
    "### We should carefully examine what does the tokenized tensor for RGB image contains. Here each entry of the tensor represent one token index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5d73a54-3a2f-4398-8bfb-3c39428c3a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_rgb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d003c320-0f3b-44d8-9c17-b32b67c8f261",
   "metadata": {},
   "source": [
    "### Detokenization\n",
    "\n",
    "One of the core properties of a tokenizer is that we can always retrieve a (lossy) version of the original input through a process called detokenization. This process takes the tokenized data and passes it through the tokenizer decoder, which reconstructs the tokens back into the original data.\n",
    "\n",
    "Note that this process is lossy, and the quality of the reconstructed image may be degraded.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "233ae7bd-0e04-4d90-a661-187aa9763ea2",
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_rgb = toks['tok_rgb'].decode_tokens(tokenized_rgb, image_size=224, timesteps=19)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5c19bbb-6c94-4026-bbde-0d586cf8fd4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstructed_rgb.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dde0e8d5-de47-48a4-94c0-9d9ac40a370d",
   "metadata": {},
   "source": [
    "Nice, we see that we have obtained a reconstructed RGB image data type of same shape as of the original RGB. Lets try to visualize it, and observe how the reconstruction looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "290b0e4f-24c6-43dd-bf90-0bc5daf81c75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reconstructed_rgb\n",
    "# Create a figure with two subplots (1 row, 2 columns)\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "\n",
    "# Display the original image on the left\n",
    "axes[0].imshow(denormalize(img, mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD)[0].permute(1, 2, 0).cpu())\n",
    "axes[0].set_title(\"Original Image\")\n",
    "axes[0].axis(\"off\")  # Hide axes\n",
    "\n",
    "# Display the reconstructed image on the right\n",
    "axes[1].imshow(denormalize(reconstructed_rgb, mean=IMAGENET_INCEPTION_STD, std=IMAGENET_INCEPTION_STD)[0].permute(1, 2, 0).cpu())\n",
    "axes[1].set_title(\"Reconstructed Image\")\n",
    "axes[1].axis(\"off\")  # Hide axes\n",
    "\n",
    "# Show the figure\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c34d2a5-799f-41f8-ba38-98726236c16b",
   "metadata": {},
   "source": [
    "Not bad! Although the resulting image is not of same quality, we can still see the high-level semantic similarity between the two images."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9c18164-b290-43bb-aeb3-a8c8ce46d633",
   "metadata": {},
   "source": [
    "### Text tokenization\n",
    "\n",
    "We show another example of how we can tokenize text data type. This is infact more simpler than image/2D based tokenizers, both interms of training and quality.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4af8748-0bd8-404f-8abd-06e120bb49df",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_text = \"Welcome to the Communications Course at EPFL\"\n",
    "\n",
    "# Encode the text\n",
    "encoded = text_tok.encode(input_text)\n",
    "\n",
    "# Get token IDs\n",
    "token_ids = encoded.ids\n",
    "print(\"Token IDs:\", token_ids)\n",
    "\n",
    "# Get tokens\n",
    "tokens = encoded.tokens\n",
    "print(\"Tokens:\", tokens)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b101dd1-43da-46e0-9e8a-20ccd6c4bbb9",
   "metadata": {},
   "source": [
    "Notice how each token represents a word in the English Vocabulary. For example, token 8580 represents the word 'welcome'. \n",
    "\n",
    "Now lets see the detokenization as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "016ab041-d072-4e6f-bd0f-b9ded2673c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decode the token IDs back to text\n",
    "decoded_text = text_tok.decode(token_ids)\n",
    "print(\"**Decoded Text:**\", decoded_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159c6417-c60b-4b2d-8610-cf47dc7ca22b",
   "metadata": {},
   "source": [
    "# **Exercise 1: Tokenization [25 points]**\n",
    "\n",
    "By now, you should have a basic understanding of tokenizers, their purpose, and how to use them.\n",
    "\n",
    "For this graded exercise, provide solutions using markdown blocks (like this current cell is a markdown) and code cell blocks (where code implementation is necessary). Please provide solutions to the following questions:\n",
    "\n",
    "### **1. Exploring the resolution / quality trade-off**\n",
    "- Based on the quality of the reconstructed RGB input, can we improve the reconstruction using the same tokenizer?\n",
    "- **Ablation on Resolution:**\n",
    "    - i) Perform reconstructions using image resolutions between 224x224 and 448x448 (in resolution steps of 32, i.e. using side-lengths [224, 256, 288, 320, 352, 384, 416, 448]). What do you notice? What are the implications of using more tokens to represent an image?\n",
    "    - ii) Imagine if you have an image of size 640x640 and you could use the same tokenizer to perform a reconstruction. Without performing any reconstruction experiment, can you tell if the reconstruction will be better if we use a very high-resolution image (like 640x640) for this same trained tokenizer? Provide evidence to your claim by performing experimentation with images having resolution more than 448.\n",
    "      \n",
    "      **Caution**: We do not expect it to work outside the resolution bounds but that you should try and see for yourselves, and provide intuitive explanation why this could be the case.\n",
    "    - iii) We have seen the values represented by the different entries in the tokenized tensor `tokenized_rgb`. What is the minimum possible value of the tensor, and what is the maximum possible value of the tensor?\n",
    "    \n",
    "      **Hint:** You must revisit and properly understand the tokenization section discussed earlier, in order to provide correct answer to this question.\n",
    "- **Ablation on the diffusion decoder timesteps:** The RGB detokenizer is a [diffusion model](https://scale.com/guides/diffusion-models-guide) that denoises a random image in `t` timesteps, using the tokens as conditioning. Show reconstructions (using both 224, and 448 res. images) for 1, 10, 20, 50, 100 timesteps. What do you notice?\n",
    "- Visualize your solutions and provide an intuitive explanation of how your modifications improve the reconstruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beefff37",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a10131ea",
   "metadata": {},
   "source": [
    "### **2. Exploring Different Tokenizers**\n",
    "- We have provided 10 additional tokenizers in this notebook, all following the same usage style (see the tokenizer loading section).\n",
    "- Choose the semantic segmentation (semseg) tokenizer from the list and implement a working solution similar to the previous question. Use the provided segmentation images at location `data/semseg_images_exercise_1`.\n",
    "- Do the findings from Question 1 also apply to the new tokenizer? Explain why or why not.\n",
    "- **Hint 1:** The most critical aspects to consider are the correct data loading and the necessary transformations.\n",
    "- **Hint 2:** You will find useful guidance on modality-specific transformations here: [Tokenizer-Specific Transforms](https://github.com/apple/ml-4m/blob/9e2a6931f625429b3280e8b355708dc4c30e881d/fourm/data/modality_info.py#L412), For tokenization: [Tokenization Script for 4M Training Data](https://github.com/apple/ml-4m/blob/main/save_vq_tokens.py) and plotting [semseg results.](https://github.com/apple/ml-4m/blob/9e2a6931f625429b3280e8b355708dc4c30e881d/fourm/utils/plotting_utils.py#L249)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fa556ec-2660-4fb6-990a-62cd7bbce76c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ade48e5",
   "metadata": {},
   "source": [
    "### **3. Tokenizing Video Data**  \n",
    "- In the previous exercises and examples, we showed how to use different tokenizers. Now, Let us consider that we want to add another modality, **Videos**, into the 4M model.\n",
    "- A usual step would be to train a dedicated video tokenizer for this task, however, we will explore if we can already get some results using our existing single-image tokenizers.\n",
    "- Your task is to implement and demonstrate the tokenization process on a sample video and compare the reconstruction results side by side. What do you notice?\n",
    "    - Load the video\n",
    "    - process/apply required transformations to the video\n",
    "    - tokenize and detokenize the processed video\n",
    "    - visualize the results\n",
    "- We provide sample videos for the exercise at location `data/videos_exercise_1`\n",
    "- **Hint:** For loading and visualizing the videos, we below provide helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd81fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from decord import VideoReader, cpu\n",
    "\n",
    "# Use this function to load the video and obtain a tensor of video-frames\n",
    "def load_sampled_frames(video_path, step=10):\n",
    "    \"\"\"\n",
    "    Load a video and return sampled frames.\n",
    "\n",
    "    Args:\n",
    "        video_path (str): Path to the video file.\n",
    "        step (int): Step size for sampling frames (default is every 10 frames).\n",
    "\n",
    "    Returns:\n",
    "        np.ndarray: Array of sampled frames (shape: N x H x W x C).\n",
    "    \"\"\"\n",
    "    vr = VideoReader(video_path, ctx=cpu(0))  # Use CPU\n",
    "    num_frames = len(vr)\n",
    "\n",
    "    # Sample frames at the given step interval\n",
    "    sampled_indices = np.arange(0, num_frames, step)\n",
    "    sampled_frames = vr.get_batch(sampled_indices).asnumpy()  # Convert to NumPy array\n",
    "\n",
    "    return sampled_frames\n",
    "\n",
    "# After you have processed/transformed video, use the below function to save it\n",
    "def save_video_tensor(processed_video, video_file_path, save_name='original'):\n",
    "    # processed_video: Processed Video Tensor of shape: T, H, W, C\n",
    "    # video_file_path: path to the original video\n",
    "    # save_name: either original or reconstructed\n",
    "\n",
    "    # returns: the path to the video saved on disk\n",
    "    video_filename = os.path.basename(video_file_path) + f\"_{save_name}.webm\"\n",
    "    fourcc = cv2.VideoWriter_fourcc(*'VP90')  # VP9 codec for WebM\n",
    "    fps = 5  # Frames per second\n",
    "    T, C, H, W = processed_video.shape\n",
    "    print(processed_video.shape)\n",
    "    processed_video_save = denormalize(processed_video, mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD)\n",
    "    processed_video_np = (np.array(processed_video_save.permute(0, 2, 3, 1).cpu()) * 255).astype(np.uint8)\n",
    "    print(processed_video_np.shape)\n",
    "    # Create video writer\n",
    "    video_writer = cv2.VideoWriter(video_filename, fourcc, fps, (W, H))\n",
    "    \n",
    "    # Write frame\n",
    "    for i in range(T):\n",
    "        # print(processed_video_np[i].shape)\n",
    "        frame = cv2.cvtColor(processed_video_np[i], cv2.COLOR_RGB2BGR)  # Convert RGB to BGR for OpenCV\n",
    "        video_writer.write(frame)\n",
    "    \n",
    "    video_writer.release()  # Save the file\n",
    "    print(f\"{save_name} video has been successfully saved: {video_filename}\")\n",
    "    return video_filename\n",
    "\n",
    "# After saving the videos (both original and reconstructed), use the below code snippet to display both videos\n",
    "from IPython.display import display, HTML, Video\n",
    "\n",
    "# Uncomment the below to display the videos after they have been saved. \n",
    "# video_html = f\"\"\"\n",
    "# <table>\n",
    "#     <tr>\n",
    "# <td style=\"text-align: center;\"><b>Original</b></td>\n",
    "#         <td style=\"text-align: center;\"><b>Reconstructed</b></td>\n",
    "#     </tr>\n",
    "#     <tr>\n",
    "#         <td>{Video(saved_path_original, embed=True)._repr_html_()}</td>\n",
    "#         <td>{Video(saved_path_reconstructed, embed=True)._repr_html_()}</td>\n",
    "#     </tr>\n",
    "# </table>\n",
    "# \"\"\"\n",
    "\n",
    "# display(HTML(video_html))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1384b88f-0f52-4e78-8080-a7c0cf0df945",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e1c97a0-6274-4373-b20e-5704f6385fa0",
   "metadata": {},
   "source": [
    "# 2. **RGB to Any-Other-Modality Generation**\n",
    "\n",
    "Now that you have learned the basics of tokenizers and how to use them, we will proceed with using 4M, which takes tokenized data inputs and generates any other tokenized data as output.\n",
    "\n",
    "First, we will understand how to generate new predictions from RGB images. RGB is the most common visual input, as it can be easily captured using common cameras. It is interesting to see that using the 4M model, we can directly generate other types of information from RGB images, such as captions, depth maps, or bounding boxes.\n",
    "\n",
    "Let's get started!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1693dab3-2876-475a-a802-372ab1196dc6",
   "metadata": {},
   "source": [
    "### Load a 4M-21 model\n",
    "\n",
    "Let's load a 4M-21 Base model that was trained on 21 modalities, including RGB, depth, surface normals, semantic segmentation, SAM instances, Canny & SAM edges, 3D human poses, CLIP-B/16 features, DINOv2-B/14 features, ImageBind-H/14 features, captions, metadata, color palette and bounding boxes. It can take any combination of those modalities as input, and can predict all of them. We wrap the model in a `GenerationSampler` which provides inference utilities.\n",
    "\n",
    "Please see [official repo's](https://github.com/apple/ml-4m) `README.md` for all available 4M and tokenizer models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fcd83ac-e1d3-45d7-9a16-03fdc7485612",
   "metadata": {},
   "outputs": [],
   "source": [
    "fm = FM.from_pretrained('EPFL-VILAB/4M-21_B').eval().to(device)\n",
    "sampler = GenerationSampler(fm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbe9510c-da04-43ea-93cc-6b91fcdfcc34",
   "metadata": {},
   "source": [
    "Lets play with the same image, we utilized for the tokenization part "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a8bcd2d-94c0-4575-9ae3-acc8717c5afb",
   "metadata": {},
   "outputs": [],
   "source": [
    "rgb_transform = RGBTransform(imagenet_default_mean_and_std=True) # for input to 4M, we use normalization with default imagenet normalization parameters\n",
    "img_pil = rgb_transform.load('./input.jpg')\n",
    "img_pil = rgb_transform.preprocess(img_pil)\n",
    "img_pil = center_crop(img_pil, (min(img_pil.size), min(img_pil.size))).resize((224,224))\n",
    "img = rgb_transform.postprocess(img_pil).unsqueeze(0).to(device)\n",
    "img_pil"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0319ea98-2471-420d-b3f3-086cbcfbb0e6",
   "metadata": {},
   "source": [
    "Given an RGB input, predict the other user-chosen modality. Inference or generation is performed by providing a *generation schedule* to the `GenerationSampler`.\n",
    "\n",
    "A *generation schedule* specifies the order of modalities to generate, and for each, it includes generation parameters such as the number of steps or temperature.\n",
    "\n",
    "For example:\n",
    "\n",
    "We can generate RGB-to-Depth directly, or by using chaining, which allows us to first generate the intermediate modalities and then generate the final modality (e.g., RGB to Surface Normal to Depth).\n",
    "\n",
    "In this section, we will consider generations *without* using the intermediate modality.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfb0c79f-6838-4a54-a095-ad3d2771e9ec",
   "metadata": {},
   "source": [
    "4M operates on a dictionary of the input and target modalities that for each modality hold the actual token values (or pixels for RGB), input and target masks that specify which of the tokens are used as input and which need to be predicted, as well as attention masks for the decoder. Wherever the `input_mask` is `False`, the corresponding entries are used as input, wherever it is `True` the tokens are ignored. Wherever the `target_mask` is `False`, the corresponding entries are predicted as targets, wherever it is `True` nothing is predicted. \n",
    "\n",
    "The dictionary is formatted in the following way, with B = batch size, and N = sequence length of each respective modality:\n",
    "\n",
    "```python\n",
    "sample_dict = {\n",
    "    modality_id: {\n",
    "        'tensor': ..., # Contains the discrete tokens ()\n",
    "        'input_mask': ..., # Boolean Tensor of shape B x N, where False = used as input, and True = ignored.\n",
    "        'target_mask': ..., # Boolean Tensor of shape B x N, where False = predicted as target, and True = ignored.\n",
    "        'decoder_attention_mask': ..., # Tensor containing the decoder attention pattern. Not used during inference.\n",
    "    },\n",
    "    ...\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0decaf71-81c7-4c20-a632-43699d785548",
   "metadata": {},
   "source": [
    "Below we define a parent function `create_generation_schedule_rgb_to_others`, that uses the `build_chained_generation_schedules` which instructs the model about user-provided input and output modality generation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cf3f8a3-ff07-4b1e-9240-ccaea0347cdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_generation_schedule_rgb_to_others(\n",
    "    target_domains, decoding_steps, temps,\n",
    "    cfg_scales, img, cfg_grow_conditioning=True\n",
    "):\n",
    "    \"\"\"\n",
    "    Create RGB to any other conditional domain\n",
    "    \"\"\"\n",
    "        \n",
    "    cond_domain = ['rgb@224']\n",
    "    cfg_schedules = [cfg_schedules_dict[target_domain] for target_domain in target_domains]\n",
    "    temp_schedules = cfg_schedules\n",
    "    autoregression_schemes = [autoregression_schemes_dict[target_domain] for target_domain in target_domains]\n",
    "    tokens_per_targets = [tokens_per_target_dict[target_domain] for target_domain in target_domains]\n",
    "    token_decoding_schedules = ['linear' if 'tok_' in target_domain else None for target_domain in target_domains]\n",
    "    schedule = build_chained_generation_schedules(\n",
    "        cond_domains=cond_domain, \n",
    "        target_domains=target_domains, \n",
    "        tokens_per_target=tokens_per_targets, \n",
    "        autoregression_schemes=autoregression_schemes, \n",
    "        decoding_steps=decoding_steps, \n",
    "        token_decoding_schedules=token_decoding_schedules, \n",
    "        temps=temps, \n",
    "        temp_schedules=temp_schedules,\n",
    "        cfg_scales=cfg_scales, \n",
    "        cfg_schedules=cfg_schedules, \n",
    "        cfg_grow_conditioning=cfg_grow_conditioning\n",
    "    )\n",
    "    \n",
    "    batched_sample = {\n",
    "        'rgb@224': {\n",
    "            'tensor': img, # Batched tensor\n",
    "            'input_mask': torch.zeros(1, 196, dtype=torch.bool, device=device), # False = used as input, True = ignored\n",
    "            'target_mask': torch.ones(1, 196, dtype=torch.bool, device=device), # False = predicted as target, True = ignored\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # Initialize target modalities\n",
    "    for target_mod, ntoks in zip(target_domains, tokens_per_targets):\n",
    "        batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)\n",
    "    \n",
    "    # Initialize input modalities\n",
    "    for cond_mod in cond_domain:\n",
    "        batched_sample = init_full_input_modality(batched_sample, MODALITY_INFO, cond_mod, device, eos_id=text_tok.token_to_id(\"[EOS]\"))\n",
    "    \n",
    "    return schedule, batched_sample\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "359888bf-13af-4372-a2e2-55f1df69376c",
   "metadata": {},
   "source": [
    "### Notes on function `create_generation_schedule_rgb_to_others`\n",
    "\n",
    "The above defined function defines a generation pipeline, where the starting modality is RGB, and we predict other modalities from it. This function acts as a wrapper on top of `build_chained_generation_schedules` so that we abstract away most hyper-parameters and configuration used for the generation process.\n",
    "\n",
    "### Function Arguments:\n",
    "\n",
    "The arguments used in the function are as follows:\n",
    "\n",
    "- `target_domains`: Defines the output modality; in our first case, we will use a caption (text) as the output modality.\n",
    "- `decoding_steps`: Specifies the number of decoding steps (if applicable).\n",
    "\n",
    "**NOTE**: The `decoding_steps` argument is only valid for 2D modalities such as Depth, Surface Normals etc. For textual modalities, we perform auto-regressive prediction, so it is set to `None`. For image-like modalities, we use iterative decoding prediction that produces the final prediction using a number of iterations set to `decoding_steps`. Remember that, the `decoding_steps` is different from the `timestamps` variable we explored during tokenization. Please read the details on different decoding strategies explained in [the 4M paper](https://arxiv.org/pdf/2312.06647) (Sec. A.3, Page 21).\n",
    "\n",
    "- `temps`: Controls randomness in sampling (higher values lead to more diverse outputs).\n",
    "- `cfg_scales`: [Classifier-free guidance](https://arxiv.org/abs/2207.12598) scale; higher values bias the output towards deterministic behavior. Setting it to 1 means that no classifier-free guidance is performed.\n",
    "- `img`: The input image for the generation process.\n",
    "- `cfg_grow_conditioning`: (True or False) Determines whether the generated modalities are used as positive guidance signal when generating subsequent modalities.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa8d3dcb-3222-48c1-8b72-db2741a15a99",
   "metadata": {},
   "source": [
    "### Example 1: Image to Caption using 4M\n",
    "We have now defined the function for generation pipeline for RGB as input. Lets use the image as input and predict caption. Note and understand each argument used while calling the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88a48d36-e00f-46f9-98a7-7b3bf73fae20",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example 1 usage: RGB to Caption\n",
    "schedule, batched_sample = create_generation_schedule_rgb_to_others(\n",
    "    target_domains=['caption'], # using caption as output\n",
    "    decoding_steps=[None], # Put None for autoregressive decoding\n",
    "    temps=[0.01], # temperature value\n",
    "    cfg_scales=[1.0], # Classifier-free guidance scale\n",
    "    img=img,\n",
    "    cfg_grow_conditioning=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e74e9537-02a2-4578-bc08-88e67aef16ca",
   "metadata": {},
   "source": [
    "After defining our inputs and targets, we perform inference with 4M model to obtain `out_dict` that contains the token predictions from 4M. To convert the tokens into respective modality, we perform de-tokenization, which converts the discrete tokens back to original data type "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3706ded7-5a7a-4479-b26a-61ac052cf354",
   "metadata": {},
   "source": [
    "Now we are ready to perform the generation. The `GenerationSampler` has a `.generate` function that performs the chained generation on a given sample dictionary, following the previously specified generation schedule.\n",
    "It outputs a dictionary that is formatted in the same manner as the sample dictionary, but contains also the predicted tokens. You can change the seed to get different outputs, or set it to None to randomly sample.\n",
    "\n",
    "4M / the sampler outputs discrete tokens, and we still need to decode them to images, feature maps, text, etc using the modality-specific tokenizers. \n",
    "For that, we provide the `decode_dict` function that takes as input the sample dictionary and the tokenizers, and returns plottable representations of each modality.\n",
    "Some modalities like RGB, depth and normals use a diffusion model as the tokenizer decoder. You can specify the number of DDIM steps for decoding with `decoding_steps`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f25361-a526-445a-8e29-1a8edae2e68d",
   "metadata": {},
   "source": [
    "#### Additional hyper-parameters for generation:\n",
    "\n",
    "We can additionally choose the top-k & top-p hyperparameters:\n",
    "\n",
    "**top_p:** When top_p > 0.0, keep only the top tokens with cumulative probability >= top_p (a.k.a. nucleus filtering).\n",
    "\n",
    "**top_k:** When top_k > 0, keep only the top k tokens with highest probability (a.k.a. top-k filtering)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41565139-6566-40b1-b58f-48ce617f1e08",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_p = 0.8\n",
    "top_k = 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b50fc21-5760-46a5-8b8f-f9320507f2e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dict = sampler.generate(\n",
    "    batched_sample, schedule, text_tokenizer=text_tok, \n",
    "    verbose=True, seed=0,\n",
    "    top_p=top_p, top_k=top_k,\n",
    ")\n",
    "# performing detokenization\n",
    "dec_dict = decode_dict(\n",
    "    out_dict, toks, text_tok, \n",
    "    image_size=224, patch_size=16,\n",
    "    decoding_steps=50\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a63a702-1f44-4abf-874e-168df4d79310",
   "metadata": {},
   "source": [
    "Image to Caption generation has been completed. It is now time to visualize the results!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a030769-5986-4e5f-a5d6-678b2c2be44e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(8,4), facecolor=(1, 1, 1))\n",
    "\n",
    "ax[0].imshow(img_pil)\n",
    "ax[0].set_title('RGB input', fontsize=18)\n",
    "\n",
    "plot_text_in_square(ax[1], dec_dict['caption'][0], wrap_width=16, fontsize=14)\n",
    "ax[1].set_title('Caption pred.', fontsize=18)\n",
    "\n",
    "\n",
    "for axis in ax.flatten():\n",
    "    axis.set_axis_off()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "820d6a0d-6bb6-4d1a-96aa-0fb9f280c4cd",
   "metadata": {},
   "source": [
    "We see that a caption data type (or modality) is predicted by the 4M model - Not bad! We will later see how to improve the generations, using chaining."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ac990c1-d7ac-4ceb-86a7-2f248d096ffd",
   "metadata": {},
   "source": [
    "### Example 2: Image to Semantic Segmentation Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef70476d-85ca-408f-86f4-a3671aa78fff",
   "metadata": {},
   "source": [
    "Now, we will see how we can generate Semantic Segmentation using the 4M model. Here, again our input is the RGB image, and our target (output) modality is the tokenized semantic segmentation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5181748a-bf33-4721-a916-a1a568eef263",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example 2 usage: RGB to Semantic Segmentation\n",
    "schedule, batched_sample = create_generation_schedule_rgb_to_others(\n",
    "    target_domains=['tok_semseg@224'], # using caption as output\n",
    "    decoding_steps=[1], # Decode in a single step, i.e. one-shot\n",
    "    temps=[0.01], # temperature value\n",
    "    cfg_scales=[2.0], # Some slight classifier-free guidance\n",
    "    img=img,\n",
    "    cfg_grow_conditioning=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24939fee-5025-467f-bcae-da5876e7fb91",
   "metadata": {},
   "outputs": [],
   "source": [
    "# As we want to predict the segmentation tokens, our tensor at the start for segmentation is all initialized from zero.\n",
    "batched_sample['tok_semseg@224']['tensor']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c10b7a-190c-40e0-aa95-594431d5e3da",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_p=0.8\n",
    "top_k=0.0\n",
    "out_dict = sampler.generate(\n",
    "    batched_sample, schedule, text_tokenizer=text_tok, \n",
    "    verbose=True, seed=0,\n",
    "    top_p=top_p, top_k=top_k,\n",
    ")\n",
    "dec_dict = decode_dict(\n",
    "    out_dict, toks, text_tok, \n",
    "    image_size=224, patch_size=16,\n",
    "    decoding_steps=50\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5fd6e84-41e4-46aa-93b3-9b7993349b12",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(8,4), facecolor=(1, 1, 1))\n",
    "\n",
    "ax[0].imshow(img_pil)\n",
    "ax[0].set_title('RGB input', fontsize=18)\n",
    "\n",
    "ax[1].imshow(dec_dict['tok_semseg@224'])\n",
    "ax[1].set_title('Semantic segmentation pred.', fontsize=18)\n",
    "\n",
    "for axis in ax.flatten():\n",
    "    axis.set_axis_off()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8002dcf1-6495-4793-8fb9-c87fbbbdb064",
   "metadata": {},
   "outputs": [],
   "source": [
    "dec_dict['tok_semseg@224'].shape # remember that for plotting purposes, the Semseg map is converted to 3 channel RGB image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b02af67-17a6-4a3f-a102-bac1320922b9",
   "metadata": {},
   "source": [
    "# **Exercise 2: RGB to Other Modality Generation [15 points]**\n",
    "\n",
    "In this exercise, you will explore generating different modalities using **only** an RGB image as the input. Additionally, you will analyze the impact of different hyperparameters in the scheduler function and assess the consistency of the generated outputs.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "191e1cbf-4cae-45fa-85dd-324beffabe38",
   "metadata": {},
   "source": [
    "### **1. Generate and Visualize Predictions**\n",
    "\n",
    "Using the 4M framework, generate and visualize the following modalities using an **RGB image as the only conditioning input**:\n",
    "\n",
    "- **CLIP**\n",
    "- **Depth**\n",
    "- **Metadata**\n",
    "- **Bounding Boxes**\n",
    "- **SAM Edges**\n",
    "- **Canny Edges**\n",
    "- **ImageBind**\n",
    "- **DINOV2**\n",
    "- **Color Palette**\n",
    "\n",
    "The corresponding **target names** for these modalities can be found in the official resource:\n",
    "[MODALITY_INFO](https://github.com/apple/ml-4m/blob/ee81e40983616e9737f444068572f0bac9961abe/fourm/data/modality_info.py#L32).\n",
    "\n",
    "**Optional (ungraded):** Use a custom image of your choice and perform the same exercise. \n",
    "\n",
    "*Below, we have provided plotting helper codes at the end of the questions to assist in visualizing the outputs.*\n",
    "\n",
    "```python\n",
    "# displaying Visual/2D Modalties\n",
    "\n",
    "ax[1].imshow(dec_dict['tok_depth@224'])\n",
    "ax[1].set_title('Depth pred.', fontsize=18)\n",
    "\n",
    "# displaying bounding boxes\n",
    "ax[1].imshow(visualize_bboxes(np.array(img_pil), dec_dict['det'][0],))\n",
    "ax[1].set_title('Bounding boxes pred.', fontsize=18)\n",
    "\n",
    "# displaying text\n",
    "plot_text_in_square(ax[1], dec_dict['caption'][0], wrap_width=16, fontsize=14)\n",
    "\n",
    "# displaying meta data\n",
    "metadata_pred = ',\\n'.join([f'{k}: {v:.2f}' if isinstance(v, float) else f'{k}: {v}' for k, v in dec_dict['metadata'].items()])\n",
    "plot_text_in_square(ax[1], metadata_pred, wrap_width=36, fontsize=13)\n",
    "\n",
    "# SAM/Canny edges\n",
    "ax[1].imshow(dec_dict['tok_sam_edge@224'], cmap='gray')\n",
    "ax[1].set_title('SAM edges pred.', fontsize=18)\n",
    "\n",
    "\n",
    "```\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "730e729a-284a-498c-aab5-9c487fad65cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5795da-3d36-4006-a975-1fc1dbafb67d",
   "metadata": {},
   "source": [
    "### **2. Understanding and Ablating Function Arguments**  \n",
    "\n",
    "The `create_generation_schedule_rgb_to_others` function takes several arguments that control the **generation process**. In this exercise, we will analyze their impact on the results generated for 2D modalities.  \n",
    "\n",
    "- For each **hyperparameter**, including `decoding_steps`, `temps`, and `cfg_scales`, re-generate and visualize Depth, Surface Normal, Semseg as the output modalities using the following configurations:  \n",
    "    - `decoding_steps`: 1, 10, 30, 50  \n",
    "    - `temps`: 0.0, 0.01, 0.05, 0.1, 0.5, 0.9, 3.0  \n",
    "    - `cfg_scales`: 1, 2, 3, 5, 10  \n",
    "\n",
    "- What trends do you observe in generation quality as each hyperparameter is varied?  \n",
    "  **Note:** When performing ablations on one hyperparameter, keep the values of the other two fixed.  \n",
    "\n",
    "- **Ablations on sampling:** Previously, we used an additional sampling parameter, **top_p**. Now, we will examine its effect on the **caption** output modality. Before proceeding, make sure you have set the `temps` to **0.5** and set `seed=None` in the `sampler.generate` function. Analyze the impact of `top_p` using the following values:  \n",
    "    - [0.0, 0.01, 0.2, 0.5, 0.7, 0.9, 0.95, 1]  \n",
    "    - What is the optimal value for achieving a maximally deterministic/most probable result? Generate examples using the optimal value.\n",
    "    - What is the optimal value for achieving a maximally diverse caption generation? Generate examples using the optimal value.\n",
    "    - Hint: you should understand **Top P**, that would help you to provide correct answers especially for the last two questions, (without even running any experiment!). More details on [Top P can be found here.](https://www.vellum.ai/llm-parameters/top-p)\n",
    "\n",
    "**Hint:** Additional relevant information on the generation arguments can be found in [this resource](https://github.com/apple/ml-4m/blob/main/README_GENERATION.md).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc5e5281-797e-4363-b594-2297344e3667",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8def0cc7-d2d5-4933-a58b-b4ce83e17963",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c67e426a-678f-4e56-94f4-d8ed3a07fcbc",
   "metadata": {},
   "source": [
    "# 3. **RGB to Any-Other-Modality Generation with Chaining**\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "852eb487-928c-4d2c-8b47-fa0f921d855e",
   "metadata": {},
   "source": [
    "### RGB → All\n",
    "\n",
    "Given an RGB input, predict all other modalities in a sequential manner.\n",
    "We utilize the same function, `build_chained_generation_schedules`, which also allows for building arbitrary chained generation schedules.\n",
    "We call it chained generation because every newly generated output is looped back into the input and serves as conditioning for subsequently generated modalities.\n",
    "This enables the generation of multiple modalities that are all consistent with each other, which is particularly important when the conditioning is underspecified.\n",
    "Please see the paper [here](https://arxiv.org/pdf/2312.06647) for more details on chained generation.\n",
    "\n",
    "In the following, we provide an RGB image as input (pixels, not tokens) and predict all other modalities from it.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbe53b0c-718e-45f0-90da-16007dde1446",
   "metadata": {},
   "source": [
    "We will explore if we can improve the generation quality and consistency by performing chained generations. \n",
    "\n",
    "As a first example, let us include all possible modalities as the middle ones that are chained together to generate the semantic segmentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1b78d8c-869a-4862-94f1-19ccd2263645",
   "metadata": {},
   "outputs": [],
   "source": [
    "img_pil"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e15e048-cfd9-40bc-902f-79a8fc7aaf0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_domains = [\n",
    "    'tok_clip@224', 'tok_dinov2@224', 'tok_imagebind@224', 'tok_depth@224', \n",
    "    'tok_normal@224', 'tok_canny_edge@224', 'tok_sam_edge@224', \n",
    "    'caption', 'det', 'human_poses', 'sam_instance', 'color_palette', 'metadata', 'tok_semseg@224'\n",
    "]\n",
    "decoding_steps = [1, 1, 1, 1, 1, 1, 1, None, None, None, None, None, None, 1]\n",
    "temps = [0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.6, 0.7, 0.1, 0.01, 0.1, 0.1, 0.01]\n",
    "cfg_scales = [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 2.0]\n",
    "cfg_grow_conditioning = True\n",
    "top_p, top_k = 0.8, 0.0\n",
    "\n",
    "schedule, batched_sample = create_generation_schedule_rgb_to_others(\n",
    "    target_domains=target_domains, # generation order will respect the order of modalities defined in this list\n",
    "    decoding_steps=decoding_steps,\n",
    "    temps=temps, # temperature value\n",
    "    cfg_scales=cfg_scales,\n",
    "    img=img,\n",
    "    cfg_grow_conditioning=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e63cdb7a-77f3-4d3b-a197-c988a823fc2f",
   "metadata": {},
   "source": [
    "Now, lets use 4M and generate the outputs!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93bcacd4-fa76-435c-bffa-2c62cf900b10",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dict = sampler.generate(\n",
    "    batched_sample, schedule, text_tokenizer=text_tok, \n",
    "    verbose=True, seed=0,\n",
    "    top_p=top_p, top_k=top_k,\n",
    ")\n",
    "dec_dict = decode_dict(\n",
    "    out_dict, toks, text_tok, \n",
    "    image_size=224, patch_size=16,\n",
    "    decoding_steps=50\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "929d3e80-4a54-4daf-af70-462ec9f6c9f9",
   "metadata": {},
   "source": [
    "Let's plot the RGB input and all predicted outputs!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c055fb-dc81-41b3-8a40-efd24a7e9731",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=3, ncols=6, figsize=(25,12), facecolor=(1, 1, 1))\n",
    "\n",
    "ax[0,0].imshow(img_pil)\n",
    "ax[0,0].set_title('RGB input', fontsize=18)\n",
    "\n",
    "ax[0,1].imshow(dec_dict['tok_depth@224'])\n",
    "ax[0,1].set_title('Depth pred.', fontsize=18)\n",
    "\n",
    "ax[0,2].imshow(dec_dict['tok_normal@224'])\n",
    "ax[0,2].set_title('Surface normals pred.', fontsize=18)\n",
    "\n",
    "ax[0,3].imshow(dec_dict['sam_instance'])\n",
    "ax[0,3].set_title('SAM instances pred. (single pass)', fontsize=18)\n",
    "\n",
    "ax[0,4].imshow(visualize_bboxes(np.array(img_pil), dec_dict['det'][0],))\n",
    "ax[0,4].set_title('Bounding boxes pred.', fontsize=18)\n",
    "\n",
    "plot_text_in_square(ax[1,1], dec_dict['caption'][0], wrap_width=16, fontsize=14)\n",
    "ax[1,1].set_title('Caption pred.', fontsize=18)\n",
    "\n",
    "ax[1,2].imshow(dec_dict['tok_canny_edge@224'], cmap='gray')\n",
    "ax[1,2].set_title('Canny edges pred.', fontsize=18)\n",
    "\n",
    "ax[1,3].imshow(dec_dict['tok_sam_edge@224'], cmap='gray')\n",
    "ax[1,3].set_title('SAM edges pred.', fontsize=18)\n",
    "\n",
    "ax[1,4].imshow(dec_dict['human_poses'])\n",
    "ax[1,4].set_title('Human poses pred.', fontsize=18)\n",
    "\n",
    "ax[1,5].imshow(dec_dict['color_palette'])\n",
    "ax[1,5].set_title('Color palette pred.', fontsize=18)\n",
    "\n",
    "ax[2,1].imshow(dec_dict['tok_clip@224'])\n",
    "ax[2,1].set_title('CLIP pred. (PCA viz.)', fontsize=18)\n",
    "\n",
    "ax[2,2].imshow(dec_dict['tok_dinov2@224'])\n",
    "ax[2,2].set_title('DINOv2 pred. (PCA viz.)', fontsize=18)\n",
    "\n",
    "ax[2,3].imshow(dec_dict['tok_imagebind@224'])\n",
    "ax[2,3].set_title('ImageBind pred. (PCA viz.)', fontsize=18)\n",
    "\n",
    "metadata_pred = ',\\n'.join([f'{k}: {v:.2f}' if isinstance(v, float) else f'{k}: {v}' for k, v in dec_dict['metadata'].items()])\n",
    "plot_text_in_square(ax[2,4], metadata_pred, wrap_width=36, fontsize=13)\n",
    "ax[2,4].set_title('Metadata pred.', fontsize=18)\n",
    "\n",
    "ax[2,5].imshow(dec_dict['tok_semseg@224'])\n",
    "ax[2,5].set_title('Semantic segmentation pred.', fontsize=18)\n",
    "\n",
    "for axis in ax.flatten():\n",
    "    axis.set_axis_off()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38367043-fbf1-4e27-877e-5c1030e7209b",
   "metadata": {},
   "source": [
    "### Quantitative Evaluation of 4M Predications\n",
    "\n",
    "So far, we have mainly relied on the qualitative visualizations to study and compare different kinds of generations. Additionally, it is a common practice to benchmark model results on evaluation datasets for a holistic overview of model performance. \n",
    "\n",
    "We will now evaluate 4M Semantic Segmentation performance on a sub-sampled test set from COCO dataset. The COCO (Common Objects in Context) dataset is a large-scale dataset for object detection, segmentation, and captioning. It contains over 330K natural images with 80 object categories that cover multiple indoor and outdoor settings with rich contextual information.\n",
    "\n",
    "\n",
    "For evaluation, we select 50 test images from COCO Validation set, and generate segmentation results using 4M model. \n",
    "\n",
    "##### **Evaluation Metric: Mean Intersection over Union (mIoU)**  \n",
    "\n",
    "To quantitatively evaluate our 4M model's semantic segmentation performance, we employ the Mean Intersection over Union (mIoU) metric, which is widely considered the standard evaluation measure in semantic segmentation tasks.\n",
    "\n",
    "The mIoU metric measures how accurately the predicted segmentation masks align with ground truth annotations. For each semantic class, an Intersection over Union (IoU) score is calculated:\n",
    "\n",
    "$$\\text{IoU} = \\frac{\\text{Area of Overlap}}{\\text{Area of Union}} = \\frac{\\text{TP}}{\\text{TP} + \\text{FP} + \\text{FN}}$$\n",
    "\n",
    "Where the \"Area of Overlap\" represents correctly predicted pixels (true positives), and the \"Area of Union\" encompasses all pixels that are either predicted or ground truth for that class (true positives + false positives + false negatives). This ratio ranges from 0 to 1, with higher values indicating better segmentation accuracy.\n",
    "\n",
    "The Mean IoU (mIoU) is then calculated by averaging individual IoU scores across all classes present in the dataset:\n",
    "\n",
    "$$\\text{mIoU} = \\frac{1}{n} \\sum_{i=1}^{n} \\text{IoU}_i$$\n",
    "\n",
    "Where n is the number of classes (80 for the COCO dataset). This metric provides a balanced assessment of segmentation quality across all object categories, regardless of their size or frequency in the dataset.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c986488-4181-4ef3-be3f-9fcff18de8e5",
   "metadata": {},
   "source": [
    "We provide the path for the dataset through a shared directory that contains the COCO validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66b3568f-d33e-4381-834d-1a87efeff242",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_PATH = '/work/com-304/datasets/coco/val/' # path on IZAR cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a998b4ee-4767-4826-85cc-ade88c7b899c",
   "metadata": {},
   "source": [
    "### mIoU calculation using RGB --> Semseg generation pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d404020c-1dda-466e-bb32-624a715069ba",
   "metadata": {},
   "source": [
    "Similar to our previous procedure for generation, we will use 4M to generate semseg results by providing RGB as the input. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b0e0d45-a8e4-4ad4-a854-ac7f2fc91990",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_domains = ['tok_semseg@224']\n",
    "decoding_steps = [1]\n",
    "temps =  [0.01]\n",
    "cfg_scales = [2.0]\n",
    "\n",
    "cfg_grow_conditioning = True\n",
    "top_p, top_k = 0.8, 0.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecfc1332-3f6c-4465-8c2e-94f2ce45278c",
   "metadata": {},
   "source": [
    "Below we define several helper functions to properly load the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3adf3cbf-0efb-4ae6-9fdb-2bd7aab7b5fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dataset(data_path, image_size=224, patch_size=16, text_tokenizer=None):\n",
    "    \"\"\"Get dataset for semantic segmentation evaluation.\"\"\"\n",
    "    modality_transforms = MODALITY_TRANSFORMS\n",
    "    modality_transforms['semseg_coco'] = SemsegTransform(shift_idx_by_one=True)\n",
    "    \n",
    "    # Setup modality info\n",
    "    loaded_domains = ['rgb@224', 'semseg_coco']\n",
    "    modality_info = {mod: MODALITY_INFO[mod] for mod in loaded_domains}\n",
    "    \n",
    "    # Configure max tokens\n",
    "    for k in modality_info:\n",
    "        num_patches = (image_size // patch_size) ** 2\n",
    "        if modality_info[k]['type'] == 'img':\n",
    "            modality_info[k]['max_tokens'] = num_patches\n",
    "            modality_info[k]['min_tokens'] = num_patches\n",
    "        modality_info[k][\"input_alphas\"] = [0.]\n",
    "        modality_info[k][\"target_alphas\"] = [0.]\n",
    "        modality_info[k][\"keep\"] = ['all']\n",
    "\n",
    "    # Setup transforms\n",
    "    image_augmenter = CenterCropImageAugmenter(target_size=image_size, main_domain='rgb@224')\n",
    "    transform = transforms.Compose([\n",
    "        UnifiedDataTransform(\n",
    "            transforms_dict=modality_transforms, \n",
    "            image_augmenter=image_augmenter\n",
    "        ),\n",
    "        UnifiedMasking(\n",
    "            modality_info=modality_info,\n",
    "            text_tokenizer=text_tokenizer,\n",
    "            input_tokens_range=768,\n",
    "            target_tokens_range=768\n",
    "        ),\n",
    "    ])\n",
    "\n",
    "    # Setup domains and paths\n",
    "    domains_without_vq = [domain for domain in loaded_domains \n",
    "                         if not modality_info[domain].get(\"requires_tokenizer\", False)]\n",
    "    modality_paths = {mod: modality_info[mod]['path'] \n",
    "                     for mod in modality_info \n",
    "                     if modality_info[mod].get('path', None) is not None}\n",
    "    \n",
    "    # Create dataset\n",
    "    dataset = MultiModalDatasetFolder(\n",
    "        data_path,\n",
    "        domains_without_vq,\n",
    "        modality_paths=modality_paths,\n",
    "        modality_transforms=modality_transforms,\n",
    "        transform=transform\n",
    "    )\n",
    "    \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "884594d5-2f8d-4115-b486-bb0053e2ab3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup dataset and dataloader\n",
    "# Note: Do not change values in this cell!\n",
    "dataset = get_dataset(\n",
    "    data_path=DATASET_PATH,\n",
    "    image_size=224,\n",
    "    patch_size=16,\n",
    "    text_tokenizer=text_tok\n",
    ")\n",
    "dataset = SubsampleDatasetWrapper(dataset, dataset_size=50, seed=0)\n",
    "loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=10)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7abb4080-8a2c-4895-b3e8-075aa115a6bb",
   "metadata": {},
   "source": [
    "Now we will generate the predictions on the 50 images of COCO Validation Set and compute the mIoU score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07d5f1c9-f1ed-45e9-85c2-a108d540ec6c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Run evaluation\n",
    "errors = []\n",
    "logged_img_counter = 0\n",
    "plotting_logs = []\n",
    "for batch, _ in tqdm(loader):\n",
    "    batched_sample = {\n",
    "        'rgb@224': {k: v.to(device) for k, v in batch['rgb@224'].items()}\n",
    "    }\n",
    "    semseg_gt = batch['semseg_coco']['tensor'].unsqueeze(1).to(device)\n",
    "    rgb = denormalize(batch['rgb@224']['tensor'], mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD).permute(0,2,3,1)\n",
    "    schedule, batched_sample = create_generation_schedule_rgb_to_others(\n",
    "    target_domains=target_domains, # generation order will respect the order of modalities defined in this list\n",
    "    decoding_steps=decoding_steps,\n",
    "    temps=temps, # temperature value\n",
    "    cfg_scales=cfg_scales,\n",
    "    img=batched_sample['rgb@224']['tensor'],\n",
    "    cfg_grow_conditioning=True,\n",
    "    )\n",
    "    out_dict = sampler.generate(\n",
    "    batched_sample, schedule, text_tokenizer=text_tok, \n",
    "    verbose=True, seed=0,\n",
    "    top_p=top_p, top_k=top_k,\n",
    "    )\n",
    "    dec_dict = decode_dict(\n",
    "        out_dict, toks, text_tok, \n",
    "        image_size=224, patch_size=16,\n",
    "        decoding_steps=50, to_rgb=False\n",
    "    )\n",
    "    output_argmax = dec_dict['tok_semseg@224'].argmax(1)\n",
    "   \n",
    "    output_argmax = rearrange(output_argmax, 'b h w -> b 1 h w')\n",
    "\n",
    "    for i in range(len(semseg_gt)):\n",
    "        # Save visualizations for first 10 images\n",
    "        if i < 9:\n",
    "            plotting_logs.append([rgb[i], semseg_gt[i], output_argmax[i]])\n",
    "\n",
    "        # Calculate metrics\n",
    "        metrics_argmax = get_semseg_metrics(output_argmax[[i]], semseg_gt[[i]])\n",
    "\n",
    "        errors.append({\n",
    "            'argmax': metrics_argmax,\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b052ee8-6a78-4b75-976d-23ab638f0fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets print the average mIoU over the selected tiny test set\n",
    "final_performance = [single_img_stats['argmax']['mean_iou'] for single_img_stats in errors]\n",
    "final_performance = (sum(final_performance) / len(final_performance)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb867d18-e33f-4023-85f1-cd155c1e05bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"mIoU result on the Test Set is : {final_performance}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2af04d31-973d-4339-8809-477940321b83",
   "metadata": {},
   "source": [
    "**Lets also plot the results!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbd698de-1f48-412a-9d78-456e1317302b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for index, single_prediction_list in enumerate(plotting_logs):\n",
    "    if index < 5: # we plot first 5 predictions\n",
    "        rgb_image = single_prediction_list[0]\n",
    "        gt_segmentation_map = single_prediction_list[1]\n",
    "        fourm_predicted_segmentation_map = single_prediction_list[2]\n",
    "        plot_rgb2semseg(rgb_image, gt_segmentation_map, fourm_predicted_segmentation_map)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c621138c-9bcc-408c-a70a-29c219c45206",
   "metadata": {},
   "source": [
    "# **Exercise 3: Chained Generation for Improved Semantic Sementation on COCO Images [30 points]**\n",
    "\n",
    "We have studied and explored how we can perform a chained generation that allows us to generate intermediate modalities before performing the final modality generation. Additionally, we quantitatively evaluated the 4M performance on a subset of COCO images.\n",
    "\n",
    "In this part of the exercise, we will perform zero-shot COCO evaluation (on the same 50 images) and quantitatively measure the results with and without chained generation.  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91b6f2a8-4b80-49e4-8199-5628e60460d0",
   "metadata": {},
   "source": [
    "### 1: SemSeg Evaluation using Chained Generation\n",
    "\n",
    "Perform the following tasks:\n",
    "1. Similar to RGB-only COCO evaluation, evaluate the performance of the chained generation using a maximum of 3 middle modalities. Specifically, perform the following chaining configurations and compute quantitative results:\n",
    "    - RGB -> Depth -> SemSeg\n",
    "    - RGB -> Depth -> Meta Data -> SemSeg\n",
    "    - RGB -> -> CLIP -> DINOv2 -> SemSeg\n",
    "    - RGB -> Meta Data -> Depth -> SemSeg\n",
    "    - RGB -> CLIP -> SemSeg\n",
    "    - RGB -> DINOv2 -> SemSeg\n",
    "2. Compare the performance obtained in the first question with RGB-only evaluations and observe the trends.\n",
    "    - Which configuration leads to the highest SemSeg performance? Based on your understanding so far, what could be the possible reason for this?\n",
    "    - Which configuration leads to the lowest SemSeg performance? Based on your understanding so far, what could be the possible reason for this?\n",
    "3. Based on the tutorial and exercise experiments:\n",
    "   - Which order of modality generations leads to the highest mIoU score?  \n",
    "   - If we are restricted to only a single modality as the middle one (options: Depth, DINOv2, CLIP, and Meta Data), which one should we use and why?\n",
    "\n",
    "- **Recommendation:** For an intuitive answer for part 3, we recommend you to understand each modality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4781b26-2d0e-4110-a546-3796b901516c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40cde740-12a0-400a-a5d5-0ca065666e3e",
   "metadata": {},
   "source": [
    "# 4. **Multiple inputs Conditioned Generation (Any to RGB Image Generation)**  \n",
    "\n",
    "So far, we have explored leveraging 4M to generate other set of modalities from RGB images as input. Inference from RGB images to the other spatial tasks can be performed well in one shot, but for reverse generative tasks like caption-to-image we need to perform inference autoregressively. Since 4M can perform chained generation, we can break down the generation process into a schedule of first generating intermediate modalities before generating the ones we care about most. For example, we found that for text-to-image generation, generating CLIP features before RGB can improve generation fidelity.\n",
    "\n",
    "Here, we aim to study and explore this opposite case and perform image generation from other set of modalities.\n",
    "\n",
    "As the first step, lets try to generate an image just based on the caption:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b8df07-3b21-4938-ab3a-eecbb61b0562",
   "metadata": {},
   "outputs": [],
   "source": [
    "caption = 'a nice restaurant in santorini at sunset [S_1]'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451b2a33-7f2a-4dd8-800a-745ac2318b2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "cond_domains = ['caption'] \n",
    "target_domains = ['tok_clip@224', 'tok_depth@224', 'tok_normal@224', 'tok_semseg@224', 'tok_rgb@224']\n",
    "decoding_steps = [50, 8, 8, 8, 25]\n",
    "temps = [5.0, 3.0, 3.0, 3.0, 3.0]\n",
    "cfg_grow_conditioning = True\n",
    "cfg_scales = [3.0, 2.0, 2.0, 2.0, 2.0]\n",
    "top_p, top_k = 0.8, 0.0\n",
    "# FIXED PARAMETERS, DO NOT CHANGE. (No need to understand them now).\n",
    "cfg_schedules = [cfg_schedules_dict[target_domain] for target_domain in target_domains]\n",
    "temp_schedules = ['onex:0.5:0.5' for single_entry in target_domains]\n",
    "autoregression_schemes = [autoregression_schemes_dict[target_domain] for target_domain in target_domains]\n",
    "tokens_per_targets = [tokens_per_target_dict[target_domain] for target_domain in target_domains]\n",
    "token_decoding_schedules = ['linear' if 'tok_' in target_domain else None for target_domain in target_domains]\n",
    "\n",
    "schedule = build_chained_generation_schedules(\n",
    "    cond_domains=cond_domains, target_domains=target_domains, tokens_per_target=tokens_per_targets, autoregression_schemes=autoregression_schemes, \n",
    "    decoding_steps=decoding_steps, token_decoding_schedules=token_decoding_schedules, temps=temps, temp_schedules=temp_schedules,\n",
    "    cfg_scales=cfg_scales, cfg_schedules=cfg_schedules, cfg_grow_conditioning=cfg_grow_conditioning, \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02c94f92-8d31-42f6-a88a-6166adaddf77",
   "metadata": {},
   "outputs": [],
   "source": [
    "batched_sample = {}\n",
    "\n",
    "# Initialize target modalities\n",
    "for target_mod, ntoks in zip(target_domains, tokens_per_targets):\n",
    "    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)\n",
    "    \n",
    "batched_sample = custom_text(\n",
    "    batched_sample, input_text=caption, eos_token='[EOS]', \n",
    "    key='caption', device=device, text_tokenizer=text_tok\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db0cdf4a-22aa-4b30-b225-df80d7fc5a28",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dict = sampler.generate(\n",
    "    batched_sample, schedule, text_tokenizer=text_tok, \n",
    "    verbose=True, seed=0,\n",
    "    top_p=top_p, top_k=top_k,\n",
    ")\n",
    "dec_dict = decode_dict(\n",
    "    out_dict, toks, text_tok, \n",
    "    image_size=224, patch_size=16,\n",
    "    decoding_steps=50\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47e73a42-da16-4bf0-8c97-4a7d28c3b692",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=2, ncols=4, figsize=(15,8), facecolor=(1, 1, 1))\n",
    "\n",
    "plot_text_in_square(ax[0,0], caption, wrap_width=32)\n",
    "ax[0,0].set_title('Caption input.')\n",
    "\n",
    "ax[0,1].imshow(dec_dict['tok_clip@224'])\n",
    "ax[0,1].set_title('CLIP pred. (PCA viz.)')\n",
    "\n",
    "ax[0,2].imshow(dec_dict['tok_depth@224'])\n",
    "ax[0,2].set_title('Depth pred.')\n",
    "\n",
    "ax[0,3].imshow(dec_dict['tok_normal@224'])\n",
    "ax[0,3].set_title('Surface normals pred.')\n",
    "\n",
    "ax[1,1].imshow(dec_dict['tok_semseg@224'])\n",
    "ax[1,1].set_title('Semantic segmentation pred.')\n",
    "\n",
    "ax[1,2].imshow(dec_dict['tok_rgb@224'])\n",
    "ax[1,2].set_title('RGB pred.')\n",
    "\n",
    "for axis in ax.flatten():\n",
    "    axis.set_axis_off()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf66b6c5-4efb-4350-959d-ccd654efae15",
   "metadata": {},
   "source": [
    "### Multiple inputs Conditioned Generation\n",
    "\n",
    "We note that, the generation performance of 4M by only conditoning on caption has some visual artifacts. We can improve the image generation by conditioning 4M on multiple input modalities. \n",
    "\n",
    "Here we will explore using both caption and bounding boxes to generate other set of modalities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2e60bd6-1cbc-478a-aff5-2736f2d1c9bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "caption = 'a nice restaurant in santorini at sunset [S_1]'\n",
    "bboxes = '[S_1] v0=0 v1=250 v2=420 v3=999 potted plant ' \\\n",
    "         'v0=400 v1=750 v2=950 v3=999 dining table ' \\\n",
    "         'v0=350 v1=250 v2=700 v3=350 boat ' \\\n",
    "         'v0=700 v1=720 v2=740 v3=850 bottle [S_2]'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6420fb94-f23e-4611-90f5-001b66221c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "cond_domains = ['caption', 'det'] \n",
    "target_domains = ['tok_clip@224', 'tok_depth@224', 'tok_normal@224', 'tok_semseg@224', 'tok_rgb@224']\n",
    "decoding_steps = [50, 8, 8, 8, 25]\n",
    "temps = [5.0, 3.0, 3.0, 3.0, 3.0]\n",
    "cfg_grow_conditioning = True\n",
    "cfg_scales = [3.0, 2.0, 2.0, 2.0, 2.0]\n",
    "top_p, top_k = 0.8, 0.0\n",
    "# FIXED PARAMETERS, DO NOT CHANGE\n",
    "cfg_schedules = [cfg_schedules_dict[target_domain] for target_domain in target_domains]\n",
    "temp_schedules = ['onex:0.5:0.5' for single_entry in target_domains]\n",
    "autoregression_schemes = [autoregression_schemes_dict[target_domain] for target_domain in target_domains]\n",
    "tokens_per_targets = [tokens_per_target_dict[target_domain] for target_domain in target_domains]\n",
    "token_decoding_schedules = ['linear' if 'tok_' in target_domain else None for target_domain in target_domains]\n",
    "\n",
    "schedule = build_chained_generation_schedules(\n",
    "    cond_domains=cond_domains, target_domains=target_domains, tokens_per_target=tokens_per_targets, autoregression_schemes=autoregression_schemes, \n",
    "    decoding_steps=decoding_steps, token_decoding_schedules=token_decoding_schedules, temps=temps, temp_schedules=temp_schedules,\n",
    "    cfg_scales=cfg_scales, cfg_schedules=cfg_schedules, cfg_grow_conditioning=cfg_grow_conditioning, \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00251fae-a34d-4395-821c-c5413d460f1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "batched_sample = {}\n",
    "\n",
    "# Initialize target modalities\n",
    "for target_mod, ntoks in zip(target_domains, tokens_per_targets):\n",
    "    batched_sample = init_empty_target_modality(batched_sample, MODALITY_INFO, target_mod, 1, ntoks, device)\n",
    "    \n",
    "batched_sample = custom_text(\n",
    "    batched_sample, input_text=caption, eos_token='[EOS]', \n",
    "    key='caption', device=device, text_tokenizer=text_tok\n",
    ")\n",
    "batched_sample = custom_text(\n",
    "    batched_sample, input_text=bboxes, eos_token='[EOS]', \n",
    "    key='det', device=device, text_tokenizer=text_tok\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d14eb5a2-c349-4dfc-aa35-5f1b3e6bcf5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dict = sampler.generate(\n",
    "    batched_sample, schedule, text_tokenizer=text_tok, \n",
    "    verbose=True, seed=0,\n",
    "    top_p=top_p, top_k=top_k,\n",
    ")\n",
    "dec_dict = decode_dict(\n",
    "    out_dict, toks, text_tok, \n",
    "    image_size=224, patch_size=16,\n",
    "    decoding_steps=50\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b84e91e-3b52-43c5-a467-33203b380bff",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(nrows=2, ncols=4, figsize=(15,8), facecolor=(1, 1, 1))\n",
    "\n",
    "plot_text_in_square(ax[0,0], caption, wrap_width=32)\n",
    "ax[0,0].set_title('Caption input.')\n",
    "\n",
    "ax[1,0].imshow(visualize_bboxes(None, dec_dict['det'][0].replace('[PAD] ', '')))\n",
    "ax[1,0].set_title('Bounding boxes input')\n",
    "\n",
    "ax[0,1].imshow(dec_dict['tok_clip@224'])\n",
    "ax[0,1].set_title('CLIP pred. (PCA viz.)')\n",
    "\n",
    "ax[0,2].imshow(dec_dict['tok_depth@224'])\n",
    "ax[0,2].set_title('Depth pred.')\n",
    "\n",
    "ax[0,3].imshow(dec_dict['tok_normal@224'])\n",
    "ax[0,3].set_title('Surface normals pred.')\n",
    "\n",
    "ax[1,1].imshow(dec_dict['tok_semseg@224'])\n",
    "ax[1,1].set_title('Semantic segmentation pred.')\n",
    "\n",
    "ax[1,2].imshow(dec_dict['tok_rgb@224'])\n",
    "ax[1,2].set_title('RGB pred.')\n",
    "\n",
    "for axis in ax.flatten():\n",
    "    axis.set_axis_off()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d24876b8-379a-493b-9433-2a42ada15826",
   "metadata": {},
   "source": [
    "#### Observations: \n",
    "\n",
    "As we add more conditioning modalities, the image generation task becomes relatively less easy for the model, as now it has multiple input sources that serves as extra context."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b32b5488-cc73-4896-ae41-92152d4c28e6",
   "metadata": {},
   "source": [
    "# Exercise 4: Any to RGB [30 points]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b9087ee-1e77-4673-9509-a1dd29ab188d",
   "metadata": {},
   "source": [
    "By now, you should have acquired basic understanding on generating RGB images using single-modality and multiple-modality inputs. In the below exercise, we would explore the role of chained generation and conditioning with additional inputs. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d497418-4509-46aa-9ad5-a03b7be11348",
   "metadata": {},
   "source": [
    "### 1. Exploring the role of Chained Generation\n",
    "\n",
    "1) While using the caption and detection as conditioning modalities, reduce the middle modalities to **zero** (i-e no chaining), and compare the results visually. Do we need middle modalities in order to generate much better RGB images? \n",
    "2) If you are restricted to a single modality as the middle one, which one would you choose from [CLIP, Color Palette, Meta Data] and why?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ba7daff-9151-41e2-be26-973da01bf8b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dcdbd90-58d4-4f97-ae47-41d4990151aa",
   "metadata": {},
   "source": [
    "### 2. Using extra modalities as conditioning signals\n",
    "\n",
    "1) In addition to bounding boxes and caption, perform and visualize the generation using **zero middle modality** while utilizing meta-data modality as the initial 3rd conditioning modality. What do you notice?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "377b3b4f-fda4-4bf0-99ce-3ce13f3c7f43",
   "metadata": {},
   "source": [
    "**Hint** For using meta-data modality: you can define them in a similar way as we did for bounding boxes and text.\n",
    "\n",
    "```python\n",
    "metadata_transform = MetadataTransform(shuffle=False, random_trunc=False, return_chunks=False)\n",
    "metadata_dict = {\n",
    "    'semantic_diversity': 20,\n",
    "}\n",
    "metadata_str = metadata_transform.metadata_to_string(metadata_dict) + ' [S_1]'\n",
    "\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "206e3ac5-b453-43d0-b301-e82f9ef0e81d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show your solution here. You can use multiple cells."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}