{"id":25857756,"url":"https://github.com/fareedkhan-dev/text2video-from-scratch","last_synced_at":"2025-07-17T15:37:18.941Z","repository":{"id":274920402,"uuid":"924469656","full_name":"FareedKhan-dev/text2video-from-scratch","owner":"FareedKhan-dev","description":"A Straightforward, Step-by-Step Implementation of a Video Diffusion Model","archived":false,"fork":false,"pushed_at":"2025-01-30T05:38:59.000Z","size":4638,"stargazers_count":26,"open_issues_count":0,"forks_count":3,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-03-01T19:17:45.157Z","etag":null,"topics":["diffusion-models","openai","pytorch","text-to-video","video-diffusion-model"],"latest_commit_sha":null,"homepage":"https://medium.com/@fareedkhandev/building-a-tiny-text-to-video-model-from-scratch-using-python-f31bdab12fbb","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/FareedKhan-dev.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2025-01-30T04:06:55.000Z","updated_at":"2025-02-28T07:49:15.000Z","dependencies_parsed_at":"2025-01-30T06:27:45.090Z","dependency_job_id":"9effe209-9cae-4735-b9dc-aa88d399daf1","html_url":"https://github.com/FareedKhan-dev/text2video-from-scratch","commit_stats":null,"previous_names":["fareedkhan-dev/text2video-from-scratch"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/FareedKhan-dev/text2video-from-scratch","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Ftext2video-from-scratch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Ftext2video-from-scratch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Ftext2video-from-scratch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Ftext2video-from-scratch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/FareedKhan-dev","download_url":"https://codeload.github.com/FareedKhan-dev/text2video-from-scratch/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/FareedKhan-dev%2Ftext2video-from-scratch/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":265623336,"owners_count":23800142,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["diffusion-models","openai","pytorch","text-to-video","video-diffusion-model"],"created_at":"2025-03-01T19:17:48.638Z","updated_at":"2025-07-17T15:37:18.927Z","avatar_url":"https://github.com/FareedKhan-dev.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"others/title.gif\" style=\"width: 100%;\"\u003e\n\n\u003cdiv align=\"center\"\u003e\n\n\u003c!-- omit in toc --\u003e\n# Text2Video from Scratch\n  \n![Python](https://img.shields.io/badge/Python-3.8%2B-blue) ![License](https://img.shields.io/badge/License-MIT-green) ![Contributions](https://img.shields.io/badge/Contributions-Welcome-red) ![Lightning AI](https://img.shields.io/badge/Lightning%20AI-For%20GPU%20Resources-pink)\n\n\u003c/div\u003e\n\nThis repository shows the step-by-step implementation of a **Text2Video** model from scratch using PyTorch. The model generates videos from text prompts using a diffusion-based architecture with a 3D U-Net and attention mechanisms.\n\nFollowing are some results of the model trained on a (non commercial) shutterstock and (open source) msrvtt dataset for 10k Training steps, trained on a single GPU for 4 days:\n\nPrompt: **A person holding camera - 10K Training Steps**\n\n\u003cimg src=\"others/hold_cam.gif\" style=\"width: 100%; height: auto;\"\u003e\n\nPrompt: **Spaceship crossing the bridge - 10K Training Steps**\n\n\u003cimg src=\"others/spaceship.gif\" style=\"width: 100%; height: auto;\"\u003e\n\nOutput of model trained on **moving mnist dataset** for **5K Training Steps**:\n\n\u003cimg src=\"others/move_mnist.gif\" style=\"width: 100%; height: auto;\"\u003e\n\nPrompt: **News reporter talking - 10K Training Steps**\n\n\u003cimg src=\"others/news_report.gif\" style=\"width: 100%; height: auto;\"\u003e\n\n\n\u003c!-- omit in toc --\u003e\n## Table of Contents\n- [Code Structure](#code-structure)\n- [Step by Step Implementation](#step-by-step-implementation)\n  - [Prerequisites](#prerequisites)\n  - [Installing Module](#installing-module)\n  - [Importing Libraries](#importing-libraries)\n  - [Preparing the Training Data](#preparing-the-training-data)\n  - [Defining Helper Functions](#defining-helper-functions)\n  - [Attention Head](#attention-head)\n  - [Building Blocks](#building-blocks)\n  - [Common Components](#common-components)\n  - [Relative Position Bias](#relative-position-bias)\n  - [UNet3D Block](#unet3d-block)\n  - [Coding Utils](#coding-utils)\n  - [Dataset Transformation](#dataset-transformation)\n  - [Gaussian Diffusion](#gaussian-diffusion)\n  - [Text Handler (BERT Model)](#text-handler-bert-model)\n  - [Trainer](#trainer)\n  - [Configs](#configs)\n  - [Model instantiation and Training](#model-instantiation-and-training)\n  - [Generating Videos](#generating-videos)\n- [Usage](#usage)\n- [References](#references)\n- [Contributing](#contributing)\n\n## Code Structure\n\nThe codebase is organized as follows:\n```bash\ntext2video-from-scratch/\n├── configs/\n│   └── default.yaml          # Configuration file for training parameters and hyperparameters\n├── src/\n│   ├── architecture/\n│   │   ├── __init__.py       # Makes the architecture directory a Python package\n│   │   ├── attention.py      # Contains the Attention and EinopsToAndFrom classes for attention mechanisms\n│   │   ├── blocks.py         # Contains Block, ResnetBlock, and SpatialLinearAttention classes (building blocks for the UNet)\n│   │   ├── common.py         # Contains common layers and utilities used in the architecture\n│   │   ├── unet.py           # Contains the main Unet3D model definition\n│   │   └── relative_position_bias.py   # Contains the RelativePositionBias class for positional encoding.\n│   ├── data_generation/\n│   │   ├── msrvtt.py        # Contains the MSRVTT class for loading the MSRVTT dataset\n│   │   └── synthetic_object.py   # Contains the SyntheticObject class for generating synthetic object data\n│   ├── data/\n│   │   ├── dataset.py        # Defines the Dataset class for loading and preprocessing video data\n│   │   └── utils.py          # Utility functions for handling video and image data.\n│   ├── diffusion/\n│   │   ├── __init__.py       # Makes the diffusion directory a Python package\n│   │   └── gaussian_diffusion.py  # Contains the GaussianDiffusion class, which implements the diffusion process\n│   ├── text/\n│   │   └── text_handler.py   # Functions for handling text input using a pre-trained BERT model (tokenization, embedding)\n│   ├── trainer/\n│   │   └── trainer.py        # Contains the Trainer class, which handles the training loop, optimization, EMA, saving, and sampling\n│   └── utils/\n│       └── helper_functions.py   # General-purpose helper functions (exists, noop, is_odd, default, cycle, etc.)\n├── train.py                  # Main training script: loads config, creates model, diffusion, trainer, and starts training\n├── generate.py                  # Main generation script: loads config, creates model, diffusion, trainer, and starts generation \n└── requirements.txt          # Lists the required Python packages for the project\n```\n\n`train.py` is the main script for training the Tiny SORA model. It loads the configuration from `configs/default.yaml`, creates the model, diffusion process, and trainer, and starts the training loop. The `generate.py` script is used for generating videos from text prompts using the trained model. The `src` directory contains the source code for the model architecture, data generation, diffusion process, text handling, and training utilities. The `requirements.txt` file lists the required Python packages for the project.\n\n\n## Step by Step Implementation\n\nTake a look at our diffusion architecture main components.\n\n 1. The main part is a **3D U-Net**, which is good at working with videos because they have frames that change over time. This U-Net isn’t just simple layers; it also uses attention. Attention helps the model focus on important parts of the video. Temporal attention looks at how frames relate to each other in time, and spatial attention focuses on the different parts of each frame. These attention layers, along with special blocks, help the model learn from the video data. Also, to make the model generate video that match the text prompt we provide, we use text embeddings.\n\n 2. The model works using something called **diffusion**. Think of it like this, first, we add noise to the training videos until they’re just random. Then, the model learns to undo this noise. To generate a video, it starts with pure noise and then gradually removes the noise using the U-Net, using the text embeddings that we provided as a guide. Important steps here include adding noise and then removing it. The text prompt is converted to embeddings using **BERT** and passed to **UNet**, enabling to generate videos from text. By doing this again and again, we get a video that matches the text we gave, which lets us make videos from words.\n\nInstead of looking at the original complex diagram, let’s visualize a simpler and easier architecture diagram that we will be coding.\n\n![Video Diffusion Architecture Created By [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/10304/1*GgBhUh8GO1WLzb97IgPJOA.png)\n\nLet’s read through the flow of our architecture that we will be coding:\n\n 1. **Input Video:** The process begins with a video that we want to learn from or use as a starting point.\n\n 2. **UNet3D Encoding:** The video goes through the UNet3D’s encoder, which is like a special part of the system that shrinks the video and extracts important features from it.\n\n 3. **UNet3D Bottleneck Processing:** The shrunken video features are then processed in the UNet3D’s bottleneck, the part that has the smallest spatial dimensions of the video in feature map.\n\n 4. **UNet3D Decoding:** Next, the processed features are sent through the UNet3D’s decoder, which enlarges the features back to a video, adding the learned structure.\n\n 5. **Text Conditioning:** The provided text prompt, used to guide the generation, is converted into a text embedding, which provides input to the UNet3D at various points, allowing the video to be generated accordingly.\n\n 6. **Diffusion Process:** During training, noise is added to the video, and the model learns to remove it. During video generation, we starts with noise, and the model uses the UNet3D to gradually remove the noise, generating the video.\n\n 7. **Output Video:** Finally, we get the output video that is generated by the model and is based on the input video or noise and the text provided as a prompt.\n\n### Prerequisites\n\nMake sure you have a good understanding of object-oriented programming (OOP) and neural networks (NN). Familiarity with PyTorch will also be helpful for coding.\n\n| Topic               | Video Link                                                |\n|---------------------|-----------------------------------------------------------|\n| OOP                 | [OOP Video](https://www.youtube.com/watch?v=Ej_02ICOIgs\u0026pp=ygUKb29wIHB5dGhvbg%3D%3D) |\n| Neural Network      | [Neural Network Video](https://www.youtube.com/watch?v=Jy4wM2X21u0\u0026pp=ygUbbmV1cmFsIG5ldHdvcmsgcHl0aG9uIHRvcmNo) |\n| Pytorch             | [Pytorch Video](https://www.youtube.com/watch?v=V_xro1bcAuA\u0026pp=ygUbbmV1cmFsIG5ldHdvcmsgcHl0aG9uIHRvcmNo) |\n\n### Installing Module\n\nMake sure Git is installed in your environment. You first need to clone the repository:\n```bash\n git clone https://github.com/FareedKhan-dev/text2video-from-scratch\n cd text2video-from-scratch\n```\n\nThen you can install the required dependencies:\n```bash\npip install -r requirements.txt\n```\n\n### Importing Libraries\n\nLet’s import the required libraries that will be used throughout this blog:\n```python\n# General system operations\nimport os  # for interacting with the operating system, like file handling\nimport yaml  # for parsing and writing YAML files\nfrom pathlib import Path  # for handling paths in a platform-independent way\nimport subprocess  # to execute system commands\nimport zipfile  # to handle ZIP files\n\n# Data manipulation and processing\nimport pandas as pd  # for working with structured data (like DataFrames)\nfrom tqdm import tqdm  # for adding progress bars to loops\n\n# Image, Video, and Dataset handling\nfrom PIL import Image  # for working with image files\nfrom moviepy.editor import VideoFileClip  # for editing and processing video files\nfrom datasets import load_dataset  # for loading datasets from the Hugging Face library\n\n# Core PyTorch modules and helpers\nimport torch  # for tensor operations and building neural networks\nfrom torch import nn, einsum  # nn for building neural networks, einsum for Einstein summation (used in tensor operations)\nfrom torch.nn import functional as F  # functional APIs for common neural network operations\nfrom torch.utils import data  # for data loading utilities\n\n# Additional tools for tensor manipulation\nfrom einops import rearrange  # for reshaping tensors\nfrom einops_exts import rearrange_many, check_shape  # additional utilities for tensor manipulation\nfrom rotary_embedding_torch import RotaryEmbedding  # for rotary embeddings (e.g., for transformers)\n\n# Transformers library\nfrom transformers import BertModel, BertTokenizer  # for BERT models and tokenization\n\n# Miscellaneous utilities\nimport copy  # for making shallow or deep copies of objects\nfrom torch.optim import Adam  # the Adam optimizer for training models\nfrom torch.cuda.amp import autocast, GradScaler  # for automatic mixed precision (AMP) training\nimport math  # for mathematical functions\nimport colorsys  # for converting and manipulating colors\n```\n\n### Preparing the Training Data\n\nWe need to have our training data to be diverse capturing a lot of different annotated videos, **MSR-VTT** (Microsoft Research Video to Text) is a perfect choice for it. It is a large-scale dataset for the open domain video captioning, which consists of 10K video clips from 20 categories, and each video clip is annotated with English sentences by Amazon Mechanical Turks. To visualize this dataset we need to first create some functions for this task.\n\n```python\n# Function to download Kaggle dataset using Kaggle CLI\ndef download_kaggle_dataset(dataset_name: str, download_dir: str) -\u003e None:\n    \n    # Ensure the directory exists\n    Path(download_dir).mkdir(parents=True, exist_ok=True)\n    \n    # Command to download dataset using Kaggle CLI\n    command = f\"kaggle datasets download {dataset_name} -p {download_dir}\"\n    subprocess.run(command, shell=True, check=True)  # Execute the command\n\n\n# Function to unzip the downloaded zip file\ndef unzip_file(zip_path: str, extract_dir: str) -\u003e None:\n    with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n        zip_ref.extractall(extract_dir)  # Extract all contents into the directory\n\n\n# Function to visualize random video frames\ndef visualize_random_videos(videos_dir: str, num_videos: int = 8) -\u003e None:\n    # Get all video files with .mp4 extension\n    video_files = [f for f in os.listdir(videos_dir) if f.endswith('.mp4')]\n    \n    # Randomly sample 'num_videos' videos\n    random_videos = random.sample(video_files, num_videos)\n\n    # Create a subplot to display the video frames\n    fig, axes = plt.subplots(2, 4, figsize=(12, 6))\n    axes = axes.ravel()  # Flatten the axes for easy indexing\n\n    # Loop through the selected videos and display their first frame\n    for i, video_file in enumerate(random_videos):\n        video_path = os.path.join(videos_dir, video_file)\n        \n        # Load the video and extract the first 2 seconds for preview\n        clip = VideoFileClip(video_path).subclip(0, 2)\n        \n        # Get the first frame of the video\n        frame = clip.get_frame(0)\n        \n        # Display the frame on the subplot\n        axes[i].imshow(frame)\n        axes[i].axis('off')  # Hide the axes for cleaner visualization\n        axes[i].set_title(f\"Video {i+1}\")  # Set the title with the video number\n\n    # Adjust layout for better spacing\n    plt.tight_layout()\n    plt.show()\n```\nWe have just defined our three important function that will download the dataset from kaggle and visualize 8 random videos from it to see how it actually looks like.\n\n```python\n# Step 1: Download and unzip the dataset\nkaggle_dataset_name = 'vishnutheepb/msrvtt'  # Dataset name on Kaggle\ndownload_dir = './msrvtt_data'  # Directory to save the downloaded dataset\nunzip_dir = './msrvtt_data/msrvtt'  # Directory to extract the contents of the zip file\n\n# Download the dataset from Kaggle\ndownload_kaggle_dataset(kaggle_dataset_name, download_dir)\n\n# Unzip the downloaded dataset\nzip_file_path = os.path.join(download_dir, 'msrvtt.zip')\nunzip_file(zip_file_path, unzip_dir)\n\n# Step 2: Visualize 8 random videos\nvideos_dir = os.path.join(unzip_dir, 'TrainValVideo')  # Directory where videos are stored\nvisualize_random_videos(videos_dir)  # Display 8 random videos' frames\n```\n\n![MSRVTT Sample Dataset](https://cdn-images-1.medium.com/max/2000/1*br7lXKOtG7fg793EYAgPHA.gif)\n\nFor training we wont be processing .mp4 but .gif converted format videos with their corresponding txt prompt as it will reduce the memory computation while training for higher number of epochs. For example, we want our training dataset to be in this format.\n```bash\n# if your training data exist under /training_data directory\n\ntraining_data/\n|── video1.gif\n|── video1.txt\n|── video2.gif\n|── video2.txt\n...\n```\n\nWe need to convert our download data into this format, so lets do it.\n\n```python\n# Function to create the training_data directory with GIFs and text files\ndef create_training_data(videos_dir: str, output_dir: str, size=(64, 64), duration=2) -\u003e None:\n\n    Path(output_dir).mkdir(parents=True, exist_ok=True)  # Ensure the output directory exists\n\n    video_files = [f for f in os.listdir(videos_dir) if f.endswith('.mp4')]  # Get all .mp4 files\n\n    for video_file in video_files:\n        video_path = os.path.join(videos_dir, video_file)\n        base_name = os.path.splitext(video_file)[0]  # Extract filename without extension\n\n        gif_path = os.path.join(output_dir, f\"{base_name}.gif\")\n        txt_path = os.path.join(output_dir, f\"{base_name}.txt\")\n\n        # Convert video to GIF\n        clip = VideoFileClip(video_path).subclip(0, duration)  # Extract first 'duration' seconds\n        clip = clip.resize(size)  # Resize the video\n        clip.write_gif(gif_path, program='ffmpeg')  # Save as GIF\n\n        # Create a text file with the video filename\n        with open(txt_path, \"w\") as txt_file:\n            txt_file.write(f\"{base_name}\")  # Write the filename as content\n\n        print(f\"Processed: {video_file} -\u003e {base_name}.gif and {base_name}.txt\")\n\n\n# Converting MSRVTT Dataset to training data format\nvideos_dir = \"./msrvtt_data/msrvtt/TrainValVideo\"  # Path to videos\noutput_dir = \"./training_data\"  # Path to save GIFs and text files\n\n# Convert MP4 videos into training format (GIF + TXT)\ncreate_training_data(videos_dir, output_dir)\n```\nAfter running the above code we will have our training data ready in our requested format within 20 minutes.\n\n### Defining Helper Functions\n\nAfter we prepared our training dataset, we need to defines several helper functions that we’ll use in different parts of our code. They are like little handy tools that do common jobs like checking if a variable exists, setting default values, and iterating through data efficiently.\n```python\n# Check if a given value is not None\ndef exists(x: Union[None, object]) -\u003e bool:\n    return x is not None\n\n# A no-operation function that does nothing and accepts any arguments\ndef noop(*args, **kwargs) -\u003e None:\n    pass\n\n# Check if a given integer is odd\ndef is_odd(n: int) -\u003e bool:\n    return (n % 2) == 1\n\n# Return val if it exists (not None), otherwise return a default value or the result of a callable\ndef default(val: Union[None, object], d: Union[object, Callable[[], object]]) -\u003e object:\n    if exists(val):  # Check if val is not None\n        return val\n    return d() if callable(d) else d  # Call d if it's a function, otherwise return it directly\n\n# Create an infinite generator that cycles through the DataLoader\ndef cycle(dl: torch.utils.data.DataLoader) -\u003e torch.utils.data.DataLoader:\n    while True:\n        for data in dl:\n            yield data\n\n# Split a number into groups of a given divisor, with a possible remainder\n# Example: num_to_groups(10, 3) -\u003e [3, 3, 3, 1]\ndef num_to_groups(num: int, divisor: int) -\u003e List[int]:\n    groups = num // divisor  # Number of full groups\n    remainder = num % divisor  # Remaining items\n    arr = [divisor] * groups  # Create groups of full size\n    if remainder \u003e 0:\n        arr.append(remainder)  # Append the remainder if it exists\n    return arr\n\n# Generate a boolean mask tensor with the given shape, where each element has a probability 'prob' of being True\ndef prob_mask_like(shape: Tuple[int, ...], prob: float, device: torch.device) -\u003e torch.Tensor:\n    if prob == 1:\n        return torch.ones(shape, device=device, dtype=torch.bool)  # All True mask\n    elif prob == 0:\n        return torch.zeros(shape, device=device, dtype=torch.bool)  # All False mask\n    else:\n        return torch.zeros(shape, device=device).float().uniform_(0, 1) \u003c prob  # Random mask with probability\n\n# Check if a given list or tuple contains only strings\ndef is_list_str(x: Union[List[object], Tuple[object, ...]]) -\u003e bool:\n    if not isinstance(x, (list, tuple)):\n        return False  # Ensure input is a list or tuple\n    return all([type(el) == str for el in x])  # Check if all elements are strings\n```\n\nThese functions help keep our code cleaner and easier to read. It doesn’t take any input directly, it just makes these functions available for use later. And it returns the ability to call these helper functions for the subsequent code blocks.\n\n### Attention Head\n\nThe very step is to define the core attention mechanisms that allow our video generation model to selectively focus on relevant parts of its input, whether it’s spatial regions, temporal frames, or text information. The central idea of attention is to mimic how humans focus their cognitive resources, enabling the model to prioritize the most important elements when processing data. This is particularly important for handling complex video data that has a temporal dimension, and also combines text information to direct the video generation.\n\n![Attention Mechanism Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/6438/1*-5cTEP_wgN4BJWI8bILRbw.png)\n```python\nclass EinopsToAndFrom(nn.Module):\n\n    def __init__(self, from_einops: str, to_einops: str, fn: Callable[[torch.Tensor], torch.Tensor]) -\u003e None:\n        super().__init__()\n        self.from_einops = from_einops  # Store the input shape pattern\n        self.to_einops = to_einops  # Store the output shape pattern\n        self.fn = fn  # Store the function to be applied\n\n    def forward(self, x: torch.Tensor, **kwargs) -\u003e torch.Tensor:\n\n        shape = x.shape  # Store original shape\n        # Create a dictionary mapping each dimension name in `from_einops` to its size\n        reconstitute_kwargs = dict(tuple(zip(self.from_einops.split(' '), shape)))\n\n        # Rearrange input tensor according to `from_einops` -\u003e `to_einops`\n        x = rearrange(x, f'{self.from_einops} -\u003e {self.to_einops}')\n        \n        # Apply the provided function\n        x = self.fn(x, **kwargs)\n\n        # Restore original shape using `reconstitute_kwargs`\n        x = rearrange(x, f'{self.to_einops} -\u003e {self.from_einops}', **reconstitute_kwargs)\n        return x\n\n\nclass Attention(nn.Module):\n\n    def __init__(\n        self,\n        dim: int,\n        heads: int = 4,\n        dim_head: int = 32,\n        rotary_emb: Optional[nn.Module] = None\n    ) -\u003e None:\n        super().__init__()\n        self.scale = dim_head ** -0.5  # Scaling factor for dot product attention\n        self.heads = heads  # Number of attention heads\n        hidden_dim = dim_head * heads  # Total feature size after projection\n\n        self.rotary_emb = rotary_emb  # Optional rotary embeddings\n        self.to_qkv = nn.Linear(dim, hidden_dim * 3, bias=False)  # Linear projection to Q, K, V\n        self.to_out = nn.Linear(hidden_dim, dim, bias=False)  # Output projection\n\n    def forward(\n        self,\n        x: torch.Tensor,\n        pos_bias: Optional[torch.Tensor] = None,\n        focus_present_mask: Optional[torch.Tensor] = None\n    ) -\u003e torch.Tensor:\n\n        n, device = x.shape[-2], x.device  # Sequence length and device\n        qkv = self.to_qkv(x).chunk(3, dim=-1)  # Compute Q, K, V and split into 3 parts\n\n        # If focus_present_mask is fully activated, return only the values (V)\n        if exists(focus_present_mask) and focus_present_mask.all():\n            values = qkv[-1]\n            return self.to_out(values)\n\n        # Rearrange Q, K, V for multi-head attention\n        q, k, v = rearrange_many(qkv, '... n (h d) -\u003e ... h n d', h=self.heads)\n        q = q * self.scale  # Scale queries\n\n        # Apply rotary embeddings (if provided)\n        if exists(self.rotary_emb):\n            q = self.rotary_emb.rotate_queries_or_keys(q)\n            k = self.rotary_emb.rotate_queries_or_keys(k)\n\n        # Compute similarity scores (scaled dot-product attention)\n        sim = einsum('... h i d, ... h j d -\u003e ... h i j', q, k)\n\n        # Add positional bias (if provided)\n        if pos_bias is not None:\n            sim = sim + pos_bias\n\n        # Apply focus mask if needed\n        if focus_present_mask is not None and not (~focus_present_mask).all():\n            attend_all_mask = torch.ones((n, n), device=device, dtype=torch.bool)  # Full attention mask\n            attend_self_mask = torch.eye(n, device=device, dtype=torch.bool)  # Self-attention mask\n            mask = torch.where(\n                rearrange(focus_present_mask, 'b -\u003e b 1 1 1 1'),\n                rearrange(attend_self_mask, 'i j -\u003e 1 1 1 i j'),\n                rearrange(attend_all_mask, 'i j -\u003e 1 1 1 i j'),\n            )\n            sim = sim.masked_fill(~mask, -torch.finfo(sim.dtype).max)  # Apply mask\n\n        # Apply numerical stability trick before softmax\n        sim = sim - sim.amax(dim=-1, keepdim=True).detach()\n\n        # Compute attention weights\n        attn = sim.softmax(dim=-1)\n\n        # Compute weighted sum of values\n        out = einsum('... h i j, ... h j d -\u003e ... h i d', attn, v)\n\n        # Rearrange back to original shape\n        out = rearrange(out, '... h n d -\u003e ... n (h d)')\n\n        return self.to_out(out)  # Apply output projection\n```\n**EinopsToAndFrom** is a utility for reshaping tensors using the einops library. It takes an einops string (specifying input and output shapes), a transformation function, and a tensor x. It rearranges x, applies the function, and restores the original shape. This is useful for reshaping data for attention mechanisms, like converting spatial maps to sequences.\n\nWhile **Attention** implements scaled multi-head self-attention, which helps the model focus on important parts of an input sequence. It computes query, key, and value matrices, applies scaled dot-product attention, and combines values to produce the final output. It supports multiple heads (heads), feature dimensions (dim), head dimensions (dim_head), and optional rotary embeddings (rotary_emb) for positional encoding. It also handles positional bias (pos_bias) and a focus mask (focus_present_mask).\n\n**Inputs and Outputs:**\n\n* **EinopsToAndFrom**: Takes a tensor x, reshapes it, applies a transformation, and returns it in the original shape.\n\n* **Attention**: Takes a tensor x (batch size, sequence length, embedding dimension) along with optional positional bias and focus mask. It returns an output tensor with the same shape, representing weighted attention-based transformations.\n\n### Building Blocks\n\nNext, we defines the fundamental building blocks that are used throughout the U-Net architecture, creating a hierarchical structure for the video generation model. These blocks are like individual layers with specific functionalities in a deep learning model and provide a framework to transform data.\n\n![Building Blocks Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/7938/1*hxol688UqEOJnOnTmYmjcg.png)\n```python\nclass Block(nn.Module):\n    def __init__(self, dim: int, dim_out: int) -\u003e None:\n        super().__init__()\n        # Initialize 3D convolutional layer with kernel size (1, 3, 3) and padding (0, 1, 1)\n        self.proj = nn.Conv3d(dim, dim_out, (1, 3, 3), padding=(0, 1, 1))\n        \n        # Initialize layer normalization for output dimension\n        self.norm = nn.LayerNorm(dim_out)\n        \n        # Initialize SiLU activation function (Sigmoid Linear Unit)\n        self.act = nn.SiLU()\n\n    def forward(self, x: torch.Tensor, scale_shift: Optional[Tuple[torch.Tensor, torch.Tensor]] = None) -\u003e torch.Tensor:\n        # Apply 3D convolution to the input tensor\n        x = self.proj(x)  \n        \n        # Apply layer normalization to the tensor\n        x = self.norm(x)\n        \n        # If scale and shift values are provided, apply them\n        if exists(scale_shift):\n            scale, shift = scale_shift\n            # Apply scaling and shifting to the tensor\n            x = x * (scale + 1) + shift\n            \n        # Apply the SiLU activation function\n        return self.act(x)\n\n\nclass ResnetBlock(nn.Module):\n    def __init__(self, dim: int, dim_out: int, *, time_emb_dim: Optional[int] = None) -\u003e None:\n        super().__init__()\n        \n        # If time_emb_dim is specified, create an MLP to generate scale and shift values\n        self.mlp = nn.Sequential(\n            nn.SiLU(),\n            nn.Linear(time_emb_dim, dim_out * 2)\n        ) if exists(time_emb_dim) else None\n\n        # Initialize two sequential blocks of the defined Block class\n        self.block1 = Block(dim, dim_out)\n        self.block2 = Block(dim_out, dim_out)\n        \n        # If input and output dimensions differ, apply a 1x1 convolution for residual connection\n        self.res_conv = nn.Conv3d(dim, dim_out, 1) if dim != dim_out else nn.Identity()\n\n    def forward(self, x: torch.Tensor, time_emb: Optional[torch.Tensor] = None) -\u003e torch.Tensor:\n        scale_shift = None\n        \n        # If an MLP is defined, process the time embedding to generate scale and shift factors\n        if exists(self.mlp):\n            assert exists(time_emb), 'time_emb must be passed in when time_emb_dim is defined'\n            time_emb = self.mlp(time_emb)  # Pass time_emb through MLP\n            time_emb = rearrange(time_emb, 'b c -\u003e b c 1 1 1')  # Reshape to enable broadcasting\n            scale_shift = time_emb.chunk(2, dim=1)  # Split time embedding into scale and shift\n\n        # Apply the first block with optional scale/shift\n        h = self.block1(x, scale_shift=scale_shift)\n        # Apply the second block without scale/shift\n        h = self.block2(h)\n        \n        # Return the result with a residual connection\n        return h + self.res_conv(x)\n\n\nclass SpatialLinearAttention(nn.Module):\n    def __init__(self, dim: int, heads: int = 4, dim_head: int = 32) -\u003e None:\n        super().__init__()\n        \n        # Scaling factor for attention scores based on the head dimension\n        self.scale = dim_head ** -0.5  \n        self.heads = heads  # Number of attention heads\n        hidden_dim = dim_head * heads  # Total dimension for multi-head attention\n\n        # 1x1 convolution to generate query, key, and value tensors\n        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)\n        \n        # 1x1 convolution for output projection after attention computation\n        self.to_out = nn.Conv2d(hidden_dim, dim, 1)\n\n    def forward(self, x: torch.Tensor) -\u003e torch.Tensor:\n        # Get the shape of the input tensor\n        b, c, f, h, w = x.shape  \n        \n        # Rearrange the input tensor to handle it in the attention mechanism\n        x = rearrange(x, 'b c f h w -\u003e (b f) c h w')\n        \n        # Apply 1x1 convolution to compute queries, keys, and values\n        qkv = self.to_qkv(x).chunk(3, dim=1)\n        \n        # Rearrange Q, K, V for multi-head attention\n        q, k, v = rearrange_many(qkv, 'b (h c) x y -\u003e b h c (x y)', h=self.heads)\n        \n        # Apply softmax over queries (across spatial locations)\n        q = q.softmax(dim=-2)\n        \n        # Apply softmax over keys (across features)\n        k = k.softmax(dim=-1)\n        \n        # Scale queries\n        q = q * self.scale\n        \n        # Compute the context (weighted sum of values) based on keys and values\n        context = torch.einsum('b h d n, b h e n -\u003e b h d e', k, v)\n\n        # Compute the attention output by applying the queries on the context\n        out = torch.einsum('b h d e, b h d n -\u003e b h e n', context, q)\n\n        # Rearrange the output back to the original spatial format\n        out = rearrange(out, 'b h c (x y) -\u003e b (h c) x y', h=self.heads, x=h, y=w)\n        \n        # Apply the output convolution to project back to the input dimension\n        out = self.to_out(out)\n        \n        # Rearrange the output back to the original batch size and frames\n        return rearrange(out, '(b f) c h w -\u003e b c f h w', b=b)\n```\n\nWe first define a **Block** which is a basic convolutional block used as a fundamental neural network unit. It takes the input tensor x, along with dim (input feature dimension) and dim_out (output feature dimension). The block applies a convolutional operation, followed by normalization and activation.\n\nThen we define a **ResnetBlock** class which enhances the basic Block by adding a residual (skip) connection for improved training stability. It takes dim, dim_out, and optionally time_emb_dim (for incorporating time information). It applies the Block twice and then adds the skip connection. Time embeddings help the model predict noise better, improving performance on tasks like image or video generation.\n\nFinally a **SpatialLinearAttention** which performs attention across the spatial dimensions (height x width) of each frame. It takes dim, heads, and dim_head for the attention operation. It uses convolutional layers to transform the data before applying attention, helping capture spatial dependencies in the input.\n\n**Inputs and Outputs:**\n\n* **Block**: Takes an input tensor x and integers dim and dim_out. It outputs a tensor after applying convolution, normalization, and activation.\n\n* **ResnetBlock**: Takes an input tensor x, an optional time embedding time_emb, and integers for feature dimensions. It returns a transformed tensor using a residual connection.\n\n* **SpatialLinearAttention**: Takes an input tensor x of shape (batch size, channel, frames, height, width). It returns a tensor of the same shape, with spatial attention applied across each frame.\n\n### Common Components\n\nNow we need to define a collection of commonly used components that are needed across various parts of our model’s architecture, like, normalization and activation functions. These components are designed to perform specific functionalities like Exponential moving average, normalization, or time embedding.\n\n![Common Components Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/4972/1*UGzrtROH-ocwGj2uLwhAtw.png)\n\n```python\nclass EMA:\n    def __init__(self, beta: float) -\u003e None:\n        super().__init__()\n        # Store the decay factor (beta) for updating the moving average\n        self.beta = beta\n\n    def update_model_average(self, ma_model: nn.Module, current_model: nn.Module) -\u003e None:\n        # Update the moving average model using the current model's parameters\n        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):\n            # Get the old weights from the moving average model and the new weights from the current model\n            old_weight, up_weight = ma_params.data, current_params.data\n            # Update the moving average weights\n            ma_params.data = self.update_average(old_weight, up_weight)\n\n    def update_average(self, old: Optional[torch.Tensor], new: torch.Tensor) -\u003e torch.Tensor:\n        # If no old value exists, return the new value\n        if old is None:\n            return new\n        # Update the moving average based on beta and the new value\n        return old * self.beta + (1 - self.beta) * new\n\n\nclass Residual(nn.Module):\n    def __init__(self, fn: nn.Module) -\u003e None:\n        super().__init__()\n        # Store the function to be used in the residual block\n        self.fn = fn\n\n    def forward(self, x: torch.Tensor, *args, **kwargs) -\u003e torch.Tensor:\n        # Apply the function and add the input tensor to it for the residual connection\n        return self.fn(x, *args, **kwargs) + x\n\n\nclass SinusoidalPosEmb(nn.Module):\n    def __init__(self, dim: int) -\u003e None:\n        super().__init__()\n        # Store the dimension for the positional embedding\n        self.dim = dim\n\n    def forward(self, x: torch.Tensor) -\u003e torch.Tensor:\n        # Get the device for the input tensor\n        device = x.device\n        # Half the dimensionality for sine and cosine embeddings\n        half_dim = self.dim // 2\n        # Scale factor for the embedding range\n        emb_scale = math.log(10000) / (half_dim - 1)\n        # Create sinusoidal embeddings by computing exponentials of scaled arange values\n        emb = torch.exp(torch.arange(half_dim, device=device) * -emb_scale)\n        # Apply the positional encoding (sinusoidal) based on input x\n        emb = x[:, None] * emb[None, :]\n        # Concatenate sine and cosine transformations of the embeddings\n        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n        return emb\n\n\ndef Upsample(dim: int) -\u003e nn.ConvTranspose3d:\n    # Return a 3D transposed convolution layer for upsampling\n    return nn.ConvTranspose3d(dim, dim, (1, 4, 4), (1, 2, 2), (0, 1, 1))\n\n\ndef Downsample(dim: int) -\u003e nn.Conv3d:\n    # Return a 3D convolution layer for downsampling\n    return nn.Conv3d(dim, dim, (1, 4, 4), (1, 2, 2), (0, 1, 1))\n\n\nclass LayerNorm(nn.Module):\n    def __init__(self, dim: int, eps: float = 1e-5) -\u003e None:\n        super().__init__()\n        # Store epsilon for numerical stability in normalization\n        self.eps = eps\n        # Create a learnable scale parameter (gamma)\n        self.gamma = nn.Parameter(torch.ones(1, dim, 1, 1, 1))\n\n    def forward(self, x: torch.Tensor) -\u003e torch.Tensor:\n        # Compute the variance and mean of the input tensor along the channel dimension (dim=1)\n        var = torch.var(x, dim=1, unbiased=False, keepdim=True)\n        mean = torch.mean(x, dim=1, keepdim=True)\n        # Perform normalization by subtracting the mean and dividing by the variance\n        # Scale it by gamma for learnable scaling\n        return (x - mean) / (var + self.eps).sqrt() * self.gamma\n\n\nclass RMSNorm(nn.Module):\n    def __init__(self, dim: int) -\u003e None:\n        super().__init__()\n        # Compute a scaling factor based on the input dimension (dim)\n        self.scale = dim ** 0.5\n        # Create a learnable scaling parameter (gamma)\n        self.gamma = nn.Parameter(torch.ones(dim, 1, 1, 1))\n\n    def forward(self, x: torch.Tensor) -\u003e torch.Tensor:\n        # Normalize input tensor along dimension 1 (channels) and apply scaling\n        return F.normalize(x, dim=1) * self.scale * self.gamma\n\n\nclass PreNorm(nn.Module):\n    def __init__(self, dim: int, fn: nn.Module) -\u003e None:\n        super().__init__()\n        # Store the function to be used after normalization\n        self.fn = fn\n        # Initialize layer normalization with the specified dimension\n        self.norm = LayerNorm(dim)\n\n    def forward(self, x: torch.Tensor, **kwargs) -\u003e torch.Tensor:\n        # Apply normalization to the input tensor\n        x = self.norm(x)\n        # Pass the normalized tensor through the function (e.g., attention or MLP)\n        return self.fn(x, **kwargs)\n```\n\nIn **EMA** class we implements an exponential moving average (EMA) for model weights. It takes a decay factor beta as input. The update_model_average function updates the EMA model parameters based on the current model, improving generalization. **Residual** wrapper class adds a skip connection to any function fn. It returns the output of fn along with the original input. While **SinusoidalPosEmb** Generates sinusoidal positional embeddings for a 1D input x. It takes a dimension dim as input and outputs the corresponding positional encoding.\n\nWe have created two **Upsample and Downsample** functions to create 3D transposed convolution (upsample) and 3D convolution (downsample) layers with fixed kernel size and stride. They take dim as input and return the respective convolutional layer.\n\n**LayerNorm** wil act as a custom layer normalization for 3D tensors. It takes a dimension dim and a small epsilon eps to avoid division by zero, and outputs the normalized tensor and **RMSNorm** will Performs root mean square normalization. It takes a dimension dim as input and outputs the normalized tensor. While on the other side **PreNorm** applies layer normalization before a given function fn. It takes a dimension dim and a function fn as input and returns the output of fn after normalization.\n\n**Inputs and Outputs:**\n\n* **EMA**: Takes beta (decay factor) as input. update_model_average takes the EMA model and current model as inputs, updating the EMA model's parameters.\n\n* **Residual**: Takes a function fn as input and outputs fn result along with the input.\n\n* **SinusoidalPosEmb**: Takes dim (dimension) as input and outputs a positional embedding.\n\n* **Upsample**: Takes dim as input and outputs a 3D transposed convolution layer.\n\n* **Downsample**: Takes dim as input and outputs a 3D convolution layer.\n\n* **LayerNorm**: Takes dim (channel dimension) and eps (small value) as inputs and outputs the normalized tensor.\n\n* **RMSNorm**: Takes dim as input and outputs the normalized tensor.\n\n* **PreNorm**: Takes dim (channel dimension) and fn (function) as input and outputs the result of fn with applied normalization.\n\n### Relative Position Bias\n\nNow we need to code a very component **RelativePositionBias**, which is responsible for incorporating positional information into the attention mechanisms. Unlike fixed positional embeddings, relative position biases learn the relationships between different positions, potentially making the model better at generalization.\n\n![Relative Position Bias Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/5272/1*JLsHAVMKfHhlskk5erWwpQ.png)\n\n```python\nclass RelativePositionBias(nn.Module):\n    def __init__(\n        self,\n        heads: int = 8,  # Number of attention heads\n        num_buckets: int = 32,  # Number of buckets for relative position encoding\n        max_distance: int = 128  # Maximum relative distance to consider\n    ) -\u003e None:\n        super().__init__()\n        # Store the number of buckets and maximum distance for relative position bias\n        self.num_buckets = num_buckets\n        self.max_distance = max_distance\n        # Initialize the embedding layer for the relative attention bias\n        self.relative_attention_bias = nn.Embedding(num_buckets, heads)\n\n    @staticmethod\n    def _relative_position_bucket(\n        relative_position: torch.Tensor,\n        num_buckets: int = 32,\n        max_distance: int = 128\n    ) -\u003e torch.Tensor:\n        # Initialize the result variable (starting at zero)\n        ret = 0\n        # Take the negative of the relative position (to handle both directions)\n        n = -relative_position\n        # Halve the number of buckets\n        num_buckets //= 2\n        # If the position is negative, assign it to the second half of the buckets\n        ret += (n \u003c 0).long() * num_buckets\n        # Get the absolute value of the relative position\n        n = torch.abs(n)\n\n        # Half of the buckets will correspond to exact distances\n        max_exact = num_buckets // 2\n        # Flag for small distances\n        is_small = n \u003c max_exact\n        # For larger distances, compute the bucket value using a logarithmic scale\n        val_if_large = max_exact + (\n            torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact)\n        ).long()\n        # Ensure that the values for large distances do not exceed the maximum bucket index\n        val_if_large = torch.min(val_if_large, torch.full_like(val_if_large, num_buckets - 1))\n        # Update the result based on whether the distance is small or large\n        ret += torch.where(is_small, n, val_if_large)\n        return ret\n\n    def forward(self, n: int, device: torch.device) -\u003e torch.Tensor:\n        # Create a tensor of query positions (q_pos) from 0 to n-1\n        q_pos = torch.arange(n, dtype=torch.long, device=device)\n        # Create a tensor of key positions (k_pos) from 0 to n-1\n        k_pos = torch.arange(n, dtype=torch.long, device=device)\n        # Compute the relative position of each key to each query (shape: n x n)\n        rel_pos = rearrange(k_pos, 'j -\u003e 1 j') - rearrange(q_pos, 'i -\u003e i 1')\n        # Compute the relative position buckets for each pair of query and key positions\n        rp_bucket = self._relative_position_bucket(\n            rel_pos,\n            num_buckets=self.num_buckets,\n            max_distance=self.max_distance\n        )\n        # Retrieve the corresponding relative position biases from the embedding layer\n        values = self.relative_attention_bias(rp_bucket)\n        # Rearrange the values to match the expected output shape (h, i, j)\n        return rearrange(values, 'i j h -\u003e h i j')\n```\nOur **RelativePositionBias** class introduces relative positional information into attention mechanisms, which is crucial for transformers, where input sequences have no inherent order. Unlike fixed positional embeddings, relative position biases learn the relationships between positions, which can improve generalization. It takes heads (number of attention heads), num_buckets (number of relative positions considered), and max_distance (parameter to calculate bucketized relative positions). and the _relative_position_bucket method, which calculates the bucket IDs for relative positions based on the number of buckets and max distance.\n\nWhile forward method takes the sequence length n and device as input. It computes all relative positions, generates bucket indices using _relative_position_bucket, and applies an embedding layer relative_attention_bias to transform bucket IDs into relative positional bias. It returns a relative position bias tensor.\n\n**Inputs and Outputs:**\n\n* It takes heads (number of attention heads), num_buckets (number of relative positions), and max_distance as inputs.\n\n* Outputs a relative position bias tensor of shape (heads, n, n).\n\n### UNet3D Block\n\nNow we define the most important and main component for our video generation model, the Unet3D. This is a 3D U-Net architecture designed for processing video data, which is inherently three-dimensional (height, width, and time/frames). This architecture incorporates residual blocks, attention mechanisms, and time and text conditioning, allowing it to effectively process video data and generate the output.\n\n![UNET 3D Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/5540/1*A3lCtnCQUpNBW92q_EIWyw.png)\n\n```python\n# Residual block with skip connection\nclass Residual(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n\n    def forward(self, x, *args, **kwargs):\n        return self.fn(x, *args, **kwargs) + x  # Skip connection\n\n# U-Net 3D model definition\nclass Unet3D(nn.Module):\n    def __init__(self, dim: int, cond_dim: Optional[int] = None, out_dim: Optional[int] = None, \n                 dim_mults: List[int] = (1, 2, 4, 8), channels: int = 3, attn_heads: int = 8, \n                 attn_dim_head: int = 32, use_bert_text_cond: bool = True, init_dim: Optional[int] = None, \n                 init_kernel_size: int = 7, use_sparse_linear_attn: bool = True):\n        super().__init__()\n        self.channels = channels\n        rotary_emb = RotaryEmbedding(min(32, attn_dim_head))\n\n        # Temporal Attention\n        temporal_attn = lambda dim: EinopsToAndFrom('b c f h w', 'b (h w) f c', \n                                                    Attention(dim, heads=attn_heads, dim_head=attn_dim_head, rotary_emb=rotary_emb))\n\n        # Position Bias for Attention\n        self.time_rel_pos_bias = RelativePositionBias(heads=attn_heads, max_distance=32)\n\n        # Initial Conv Layer\n        init_dim = default(init_dim, dim)\n        init_padding = init_kernel_size // 2\n        self.init_conv = nn.Conv3d(channels, init_dim, (1, init_kernel_size, init_kernel_size), \n                                   padding=(0, init_padding, init_padding))\n        self.init_temporal_attn = Residual(PreNorm(init_dim, temporal_attn(init_dim)))\n\n        # Dimension setup\n        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]\n        in_out = list(zip(dims[:-1], dims[1:]))\n        time_dim = dim * 4\n        self.time_mlp = nn.Sequential(SinusoidalPosEmb(dim), nn.Linear(dim, time_dim), nn.GELU(), nn.Linear(time_dim, time_dim))\n\n        # Conditioning setup\n        self.has_cond = exists(cond_dim) or use_bert_text_cond\n        cond_dim = 768 if use_bert_text_cond else cond_dim\n        self.null_cond_emb = nn.Parameter(torch.randn(1, cond_dim)) if self.has_cond else None\n        cond_dim = time_dim + int(cond_dim or 0)\n\n        # Downsampling and Upsampling blocks\n        self.downs, self.ups = nn.ModuleList([]), nn.ModuleList([])\n        block_klass = ResnetBlock\n        block_klass_cond = partial(block_klass, time_emb_dim=cond_dim)\n\n        # Downsampling Blocks\n        for ind, (dim_in, dim_out) in enumerate(in_out):\n            is_last = ind \u003e= (len(in_out) - 1)\n            self.downs.append(nn.ModuleList([\n                block_klass_cond(dim_in, dim_out),\n                block_klass_cond(dim_out, dim_out),\n                Residual(PreNorm(dim_out, SpatialLinearAttention(dim_out, heads=attn_heads))) if use_sparse_linear_attn else nn.Identity(),\n                Residual(PreNorm(dim_out, temporal_attn(dim_out))),\n                Downsample(dim_out) if not is_last else nn.Identity()\n            ]))\n\n        # Middle Bottleneck Block\n        mid_dim = dims[-1]\n        self.mid_block1 = block_klass_cond(mid_dim, mid_dim)\n        self.mid_spatial_attn = Residual(PreNorm(mid_dim, EinopsToAndFrom('b c f h w', 'b f (h w) c', Attention(mid_dim, heads=attn_heads))))\n        self.mid_temporal_attn = Residual(PreNorm(mid_dim, temporal_attn(mid_dim)))\n        self.mid_block2 = block_klass_cond(mid_dim, mid_dim)\n\n        # Upsampling Blocks\n        for ind, (dim_in, dim_out) in enumerate(reversed(in_out)):\n            is_last = ind \u003e= (len(in_out) - 1)\n            self.ups.append(nn.ModuleList([\n                block_klass_cond(dim_out * 2, dim_in),\n                block_klass_cond(dim_in, dim_in),\n                Residual(PreNorm(dim_in, SpatialLinearAttention(dim_in, heads=attn_heads))) if use_sparse_linear_attn else nn.Identity(),\n                Residual(PreNorm(dim_in, temporal_attn(dim_in))),\n                Upsample(dim_in) if not is_last else nn.Identity()\n            ]))\n\n        # Final output layer\n        out_dim = default(out_dim, channels)\n        self.final_conv = nn.Sequential(block_klass(dim * 2, dim), nn.Conv3d(dim, out_dim, 1))\n\n    def forward(self, x: torch.Tensor, time: torch.Tensor, cond: Optional[torch.Tensor] = None, \n                null_cond_prob: float = 0., focus_present_mask: Optional[torch.Tensor] = None) -\u003e torch.Tensor:\n        assert not (self.has_cond and not exists(cond)), 'Conditioning must be provided.'\n        batch, device = x.shape[0], x.device\n        time_rel_pos_bias = self.time_rel_pos_bias(x.shape[2], device=x.device)\n\n        # Initial Conv and Temporal Attention\n        x = self.init_conv(x)\n        x = self.init_temporal_attn(x, pos_bias=time_rel_pos_bias)\n        r = x.clone()  # Save for later\n\n        # Time and Condition Embedding\n        t = self.time_mlp(time) if exists(self.time_mlp) else None\n        if self.has_cond:\n            mask = prob_mask_like((batch,), null_cond_prob, device=device)\n            cond = torch.where(rearrange(mask, 'b -\u003e b 1'), self.null_cond_emb, cond)\n            t = torch.cat((t, cond), dim=-1)\n\n        # Downsampling Path\n        h = []\n        for block1, block2, spatial_attn, temporal_attn, downsample in self.downs:\n            x = block1(x, t)\n            x = block2(x, t)\n            x = spatial_attn(x)\n            x = temporal_attn(x, pos_bias=time_rel_pos_bias)\n            h.append(x)\n            x = downsample(x)\n\n        # Mid Block (Bottleneck)\n        x = self.mid_block1(x, t)\n        x = self.mid_spatial_attn(x)\n        x = self.mid_temporal_attn(x, pos_bias=time_rel_pos_bias)\n        x = self.mid_block2(x, t)\n\n        # Upsampling Path\n        for block1, block2, spatial_attn, temporal_attn, upsample in self.ups:\n            x = torch.cat((x, h.pop()), dim=1)\n            x = block1(x, t)\n            x = block2(x, t)\n            x = spatial_attn(x)\n            x = temporal_attn(x, pos_bias=time_rel_pos_bias)\n            x = upsample(x)\n\n        # Combine and final output\n        x = torch.cat((x, r), dim=1)\n        return self.final_conv(x)\n```\nWe designed **U-Net 3D** architecture in order to process 3D video or volumetric data for tasks such as video segmentation and prediction. It accepts multiple configuration parameters, including the feature dimension (dim), embedding dimension (cond_dim), output dimension (out_dim), and dimension multipliers (dim_mults). Additionally, it supports adjustments for attention heads (heads), attention dimension (dim_head), and can conditionally use text embedding and sparse linear attention, allowing for flexible model tuning for different tasks.\n\nThe model includes several components: it starts with an initial convolutional layer (init_conv) to process the input data, followed by temporal attention via init_temporal_attn. Rotary positional embeddings (rotary_emb) are used to capture temporal relationships. A sinusoidal position embedding (SinusoidalPosEmb) further encodes time-based aspects, while relative positional bias (RelativePositionBias) accounts for relative positioning within the sequence, enhancing spatial and temporal reasoning.\n\n**Inputs and Outputs:**\n\n* The `__init__` method of the U-Net 3D model takes parameters for feature and embedding dimensions, attention configurations, and optional components like text embeddings. The forward method processes a 3D video tensor and timestep tensor, using downsampling and upsampling paths to extract and reconstruct hierarchical features.\n\n* The model applies spatial and temporal attention to capture long-range dependencies across both dimensions, producing a processed 3D tensor suitable for tasks like video segmentation, prediction, or generation.\n\n### Coding Utils\n\nWe need to designed abstract away (functions) forcommon tasks of loading, transforming, and saving video data, particularly in GIF format, which is our chosen video representation for training. These utilities are crucial for preparing the data to be used by the model. So let’s do it.\n```python\n# Dictionary mapping the number of channels to the corresponding image mode.\nCHANNELS_TO_MODE = {\n    1: 'L',       # 1 channel corresponds to grayscale ('L' mode).\n    3: 'RGB',     # 3 channels correspond to RGB color mode.\n    4: 'RGBA'     # 4 channels correspond to RGBA color mode (with alpha transparency).\n}\n\n# Generator function to seek all images from a multi-frame image (e.g., GIF).\ndef seek_all_images(img: Image.Image, channels: int = 3):\n    # Ensure the specified number of channels is valid.\n    assert channels in CHANNELS_TO_MODE, f'channels {channels} invalid'\n    # Get the corresponding mode for the number of channels.\n    mode = CHANNELS_TO_MODE[channels]\n\n    i = 0\n    while True:\n        try:\n            # Seek to the i-th frame in the image.\n            img.seek(i)\n            # Convert the image frame to the desired mode and yield it.\n            yield img.convert(mode)\n        except EOFError:\n            # End of frames (EOF), break the loop.\n            break\n        i += 1\n\n# Function to convert a video tensor into a GIF and save it to the given path.\ndef video_tensor_to_gif(\n    tensor: torch.Tensor,\n    path: str,\n    duration: int = 120,\n    loop: int = 0,\n    optimize: bool = True\n):\n    # Convert each frame from the video tensor into a PIL image.\n    images = map(T.ToPILImage(), tensor.unbind(dim=1))\n    # Unpack the first image and the rest of the images.\n    first_img, *rest_imgs = images\n    # Save the GIF with the specified parameters.\n    first_img.save(\n        path,\n        save_all=True,                # Save all frames as part of the GIF.\n        append_images=rest_imgs,      # Append the other frames to the GIF.\n        duration=duration,            # Set the duration (in ms) for each frame.\n        loop=loop,                    # Set the loop count for the GIF (0 means infinite loop).\n        optimize=optimize             # Enable optimization of the GIF file.\n    )\n    # Return the list of images as a result.\n    return images\n\n# Function to convert a GIF into a tensor (a sequence of frames).\ndef gif_to_tensor(\n    path: str,\n    channels: int = 3,\n    transform: T.Compose = T.ToTensor()\n) -\u003e torch.Tensor:\n    # Open the GIF image from the given path.\n    img = Image.open(path)\n    # Convert all frames in the GIF to tensors, applying the transform.\n    tensors = tuple(map(transform, seek_all_images(img, channels=channels)))\n    # Stack the tensors into a single tensor along the frame dimension.\n    return torch.stack(tensors, dim=1)\n\n# Identity function: returns the input tensor unchanged.\ndef identity(t, *args, **kwargs):\n    return t\n\n# Function to normalize an image tensor to the range [-1, 1].\ndef normalize_img(t: torch.Tensor) -\u003e torch.Tensor:\n    # Normalize by scaling the tensor values from [0, 1] to [-1, 1].\n    return t * 2 - 1\n\n# Function to unnormalize an image tensor back to the range [0, 1].\ndef unnormalize_img(t: torch.Tensor) -\u003e torch.Tensor:\n    # Unnormalize by scaling the tensor values from [-1, 1] to [0, 1].\n    return (t + 1) * 0.5\n\n# Function to ensure a tensor has the specified number of frames.\ndef cast_num_frames(t: torch.Tensor, *, frames: int) -\u003e torch.Tensor:\n    # Get the current number of frames in the tensor.\n    f = t.shape[1]\n    if f == frames:\n        # If the number of frames is already as required, return the tensor unchanged.\n        return t\n    if f \u003e frames:\n        # If there are more frames than needed, slice the tensor to the required number of frames.\n        return t[:, :frames]\n    # If there are fewer frames than needed, pad the tensor with zeros (no new frames).\n    return torch.nn.functional.pad(t, (0, 0, 0, 0, 0, frames - f))\n```\nSo we have coded some utilities for manipulating video and image data at the tensor level. For instance, seek_all_images helps to extract individual frames from a multi-frame image, whereas video_tensor_to_gif allows converting a video tensor into a GIF while saving it to disk. gif_to_tensor allows the reverse, converting a GIF back into a PyTorch tensor, enabling further analysis or transformation.\n\nAdditionally, functions like normalize_img and unnormalize_img help adjust image tensor values between different scales. The identity function acts as a simple pass-through for any input, ensuring that data remains unchanged, while cast_num_frames allows resizing the number of frames in a video tensor, making it adaptable to various frame-rate or frame-count requirements.\n\n### Dataset Transformation\n\nWe need to define a standard way to load data for training our diffusion model in PyTorch. The purpose of this class is to handle loading video data and text descriptions efficiently during training, so that PyTorch can load the data in batches for our training process.\n```python\n# Custom dataset class for handling GIF or video files.\nclass Dataset(data.Dataset):\n    # Initialize the dataset with the required parameters.\n    def __init__(\n        self,\n        folder: str,                   # Folder path where the dataset is stored.\n        image_size: int,               # The size to which each image is resized.\n        channels: int = 3,             # Number of color channels (default is 3 for RGB).\n        num_frames: int = 16,          # The number of frames to extract per video (default is 16).\n        horizontal_flip: bool = False, # Whether to apply horizontal flip augmentation.\n        force_num_frames: bool = True, # Whether to force the video tensors to have exactly `num_frames` frames.\n        exts: List[str] = ['gif']      # List of file extensions to look for (default is ['gif']).\n    ) -\u003e None:\n        # Call the parent constructor (from PyTorch's Dataset).\n        super().__init__()\n\n        # Initialize the dataset attributes.\n        self.folder = folder\n        self.image_size = image_size\n        self.channels = channels\n        # Get all the file paths in the folder (and subfolders) that match the given extensions.\n        self.paths = [\n            p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')\n        ]\n        # Define the function for casting the number of frames if necessary.\n        # If `force_num_frames` is True, we apply the `cast_num_frames` function, otherwise, we use the identity function.\n        self.cast_num_frames_fn = partial(cast_num_frames, frames=num_frames) if force_num_frames else identity\n\n        # Define the transformations to be applied to each image (resize, random flip, crop, and convert to tensor).\n        self.transform = T.Compose([\n            T.Resize(image_size),                           # Resize the image to the target size.\n            T.RandomHorizontalFlip() if horizontal_flip else T.Lambda(identity),  # Apply random horizontal flip if specified.\n            T.CenterCrop(image_size),                       # Center crop the image to the target size.\n            T.ToTensor()                                    # Convert the image to a PyTorch tensor.\n        ])\n\n    # Return the total number of samples in the dataset.\n    def __len__(self) -\u003e int:\n        return len(self.paths)\n\n    # Get a specific sample (image and its corresponding text, if available) by index.\n    def __getitem__(self, index: int) -\u003e Tuple[torch.Tensor, Optional[str]]:\n        # Get the file path for the sample at the given index.\n        path = self.paths[index]\n        \n        # Convert the GIF (or video) to a tensor using the `gif_to_tensor` function.\n        # Apply the transformations defined earlier.\n        tensor = gif_to_tensor(path, self.channels, transform=self.transform)\n        \n        # Cast the tensor to have the correct number of frames (if needed).\n        tensor = self.cast_num_frames_fn(tensor)\n        \n        # Check if there is a corresponding text file for this image (same name, .txt extension).\n        text_path = path.with_suffix(\".txt\")\n        if text_path.exists():\n            # If the text file exists, read its content.\n            with open(text_path, 'r') as f:\n                text = f.read()\n                # Return the tensor and the text from the file.\n                return tensor, text\n        else:\n            # If no text file exists, return the tensor with `None` for the text.\n            return tensor, None\n```\n\nThe Dataset class loads video and text data by initializing with a folder path, image size, channels, and frame count parameters. It collects video file paths, adjusts frame counts, applies transformations, and converts frames to tensors. It associates text with videos when available.\n\nIt includes three methods: `__init__` (initializes with parameters), `__len__` (returns dataset size), and `__getitem__` (fetches a video-frame tensor and optional text).\n\n### Gaussian Diffusion \n\nNow we need to define  the GaussianDiffusion class, which implements the forward and reverse diffusion processes. This process is at the heart of our model and this code handles all of the complexities of adding noise to the image and subsequently denoising using a deep neural network to sample a video.\n\n![Gaussian Diffusion Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/6704/1*T8H-NhlY1eSGfD-9njWHfQ.png)\n\n```python\n# Helper function to extract values from tensors based on time step\ndef extract(a: torch.Tensor, t: torch.Tensor, x_shape: torch.Size) -\u003e torch.Tensor:\n    b, *_ = t.shape  # Get batch size\n    out = a.gather(-1, t)  # Extract values based on time step\n    return out.reshape(b, *((1,) * (len(x_shape) - 1)))  # Reshape to match input shape\n\n# Function to create a cosine schedule for the betas\ndef cosine_beta_schedule(timesteps: int, s: float = 0.008) -\u003e torch.Tensor:\n    steps = timesteps + 1\n    x = torch.linspace(0, timesteps, steps, dtype=torch.float64)  # Create a time grid\n    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2  # Cosine function\n    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]  # Normalize\n    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])  # Calculate betas\n    return torch.clip(betas, 0, 0.9999)  # Ensure betas stay in range\n\n# Main class for Gaussian Diffusion Model\nclass GaussianDiffusion(nn.Module):\n    def __init__(self, denoise_fn: nn.Module, *, image_size: int, num_frames: int, timesteps: int = 1000):\n        super().__init__()\n        self.denoise_fn = denoise_fn\n        self.image_size = image_size\n        self.num_frames = num_frames\n\n        betas = cosine_beta_schedule(timesteps)  # Get beta schedule\n\n        # Initialize various tensors for model calculations\n        alphas = 1. - betas\n        alphas_cumprod = torch.cumprod(alphas, axis=0)\n        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.)\n\n        timesteps, = betas.shape\n        self.num_timesteps = int(timesteps)\n\n        # Register buffers (tensors that are not updated by gradient descent)\n        register_buffer = lambda name, val: self.register_buffer(name, val.to(torch.float32))\n        register_buffer('betas', betas)\n        register_buffer('alphas_cumprod', alphas_cumprod)\n        register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)\n\n        # More initialization for various coefficients (for calculating posterior and forward process)\n        register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))\n        register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))\n        register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))\n\n    # Function to calculate the mean, variance, and log variance for q distribution\n    def q_mean_variance(self, x_start: torch.Tensor, t: torch.Tensor) -\u003e tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        mean = extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start\n        variance = extract(1. - self.alphas_cumprod, t, x_start.shape)\n        log_variance = extract(self.log_one_minus_alphas_cumprod, t, x_start.shape)\n        return mean, variance, log_variance\n\n    # Function to predict the start of the image from noisy data\n    def predict_start_from_noise(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -\u003e torch.Tensor:\n        return (\n            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -\n            extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise\n        )\n\n    # Function to calculate posterior distribution\n    def q_posterior(self, x_start: torch.Tensor, x_t: torch.Tensor, t: torch.Tensor) -\u003e tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        posterior_mean = (\n            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +\n            extract(self.posterior_mean_coef2, t, x_t.shape) * x_t\n        )\n        posterior_variance = extract(self.posterior_variance, t, x_t.shape)\n        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)\n        return posterior_mean, posterior_variance, posterior_log_variance_clipped\n\n    # Function for denoising using model predictions\n    def p_mean_variance(self, x: torch.Tensor, t: torch.Tensor, clip_denoised: bool) -\u003e tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        x_recon = self.predict_start_from_noise(x, t=t, noise=self.denoise_fn(x))\n\n        if clip_denoised:  # Clip the denoised image if needed\n            x_recon = x_recon.clamp(-1., 1.)\n\n        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)\n        return model_mean, posterior_variance, posterior_log_variance\n\n    # Function for a single denoising step\n    @torch.inference_mode()\n    def p_sample(self, x: torch.Tensor, t: torch.Tensor) -\u003e torch.Tensor:\n        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=True)\n        noise = torch.randn_like(x)  # Add random noise\n        return model_mean + noise * (0.5 * model_log_variance).exp()  # Return denoised image\n\n    # Function for generating a sample (whole loop)\n    @torch.inference_mode()\n    def p_sample_loop(self, shape: torch.Size) -\u003e torch.Tensor:\n        img = torch.randn(shape, device=self.device)  # Start with random noise\n\n        for t in reversed(range(self.num_timesteps)):  # Denoise iteratively\n            img = self.p_sample(img, t)\n        return (img + 1) * 0.5  # Return final image in proper range\n\n    # Function to generate a batch of samples\n    @torch.inference_mode()\n    def sample(self, batch_size: int = 16) -\u003e torch.Tensor:\n        return self.p_sample_loop((batch_size, self.channels, self.num_frames, self.image_size, self.image_size))\n\n    # Function to calculate loss (e.g., L1 loss) between the noisy and denoised images\n    def p_losses(self, x_start: torch.Tensor, t: torch.Tensor, noise: torch.Tensor = None) -\u003e torch.Tensor:\n        x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)  # Add noise to the image\n        x_recon = self.denoise_fn(x_noisy, t)  # Use model to denoise\n\n        # Compute loss between noisy and denoised output\n        return F.l1_loss(noise, x_recon)\n\n    # The forward pass of the model\n    def forward(self, x: torch.Tensor) -\u003e torch.Tensor:\n        t = torch.randint(0, self.num_timesteps, (x.shape[0],))  # Random time steps\n        return self.p_losses(x, t)  # Compute losses\n```\n\nWe first defines a diffusion-based video generation model using a denoising neural network. It initializes with user-provided parameters such as the denoising function (denoise_fn), the number of diffusion steps, image size, number of channels, loss type, dynamic thresholding, and more. The model calculates parameters for the forward and reverse diffusion processes, including the cosine beta schedule, which determines how much noise is added at each step. These parameters are stored as module buffers, making them part of the model's parameters for training. Methods like q_mean_variance, predict_start_from_noise, q_posterior, and p_mean_variance are used to compute the necessary values for each diffusion step.\n\nThe p_sample method applies the denoising function iteratively to reduce noise in the generated image, using the precomputed parameters. The p_sample_loop method runs the entire diffusion sampling process, starting from pure noise and generating the final video output. The sample method prepares the data by converting text to embeddings (if configured), then calls p_sample_loop to generate the video. The interpolate method takes two video tensors and produces an interpolated video based on given time and parameters. The q_sample method handles the forward diffusion process by adding Gaussian noise to a given image.\n\nAdditionally, the p_losses method calculates the loss between the predicted image and the ground truth, which is used to update the model's parameters during training. The forward method serves as the main entry point, executing a full pass through the model and returning the computed loss for a given input, typically consisting of video and text data. Each method in the class serves a specific role in the overall diffusion and denoising process, contributing to the final output video generation.\n\n### Text Handler (BERT Model)\n\nWe have to code a simple way to turn text into numerical representations (i.e., embeddings), which then our diffusion model can understand and use to generate the videos. We will uses a pre-trained BERT model.\n\n![Tokenization Process Created by [Fareed Khan](undefined)](https://cdn-images-1.medium.com/max/5102/1*Of6h15X3KYfWXNT5mkPkuA.png)\n```python\n# Function to check if a value exists (not None)\ndef exists(val: Optional[Union[torch.Tensor, any]]) -\u003e bool:\n    return val is not None\n\n# Initialize the model and tokenizer variables as None\nMODEL: Optional[BertModel] = None\nTOKENIZER: Optional[BertTokenizer] = None\nBERT_MODEL_DIM: int = 768  # Dimension size for BERT model output\n\n# Function to get the tokenizer for BERT model\ndef get_tokenizer() -\u003e BertTokenizer:\n    global TOKENIZER\n    if not exists(TOKENIZER):  # If tokenizer is not already loaded\n        TOKENIZER = BertTokenizer.from_pretrained('bert-base-cased')  # Load the tokenizer\n    return TOKENIZER\n\n# Function to get the BERT model\ndef get_bert() -\u003e BertModel:\n    global MODEL\n    if not exists(MODEL):  # If the model is not already loaded\n        MODEL = BertModel.from_pretrained('bert-base-cased')  # Load the BERT model\n        if torch.cuda.is_available():  # If GPU is available\n            MODEL = MODEL.cuda()  # Move model to GPU\n    return MODEL\n\n# Function to tokenize input text (single string or a list of strings)\ndef tokenize(texts: Union[str, List[str], Tuple[str]]) -\u003e torch.Tensor:\n    if not isinstance(texts, (list, tuple)):  # If input is a single string, convert it to list\n        texts = [texts]\n\n    tokenizer = get_tokenizer()  # Get the tokenizer\n    encoding = tokenizer.batch_encode_plus(\n        texts,  # Input texts\n        add_special_tokens=True,  # Add special tokens for BERT\n        padding=True,  # Pad sequences to the same length\n        return_tensors='pt'  # Return as PyTorch tensor\n    )\n    return encoding.input_ids  # Return the token IDs (numerical representation)\n\n# Function to get the BERT embeddings (features) from token IDs\n@torch.no_grad()  # No need to track gradients for inference\ndef bert_embed(\n    token_ids: torch.Tensor,\n    return_cls_repr: bool = False,  # Whether to return only the [CLS] token representation\n    eps: float = 1e-8,  # Small value to prevent division by zero\n    pad_id: int = 0  # Padding token ID (usually 0 for BERT)\n) -\u003e torch.Tensor:\n    model = get_bert()  # Get the BERT model\n    mask = token_ids != pad_id  # Create a mask for padding tokens (to ignore them)\n\n    if torch.cuda.is_available():  # If GPU is available, move tensors to GPU\n        token_ids = token_ids.cuda()\n        mask = mask.cuda()\n\n    # Run BERT model and get the outputs (hidden states from all layers)\n    outputs = model(\n        input_ids=token_ids,\n        attention_mask=mask,  # Only pay attention to non-padding tokens\n        output_hidden_states=True  # Get hidden states from all layers\n    )\n    hidden_state = outputs.hidden_states[-1]  # Get the last hidden state (final layer)\n\n    if return_cls_repr:  # If we need the [CLS] token representation, return it\n        return hidden_state[:, 0]\n\n    # If no mask, return the mean of all hidden states\n    if not exists(mask):\n        return hidden_state.mean(dim=1)\n\n    # If there is a mask, calculate the mean ignoring the padding tokens\n    mask = mask[:, 1:]  # Remove the padding for the first token\n    mask = rearrange(mask, 'b n -\u003e b n 1')  # Rearrange for broadcasting\n    numer = (hidden_state[:, 1:] * mask).sum(dim=1)  # Sum over the masked tokens\n    denom = mask.sum(dim=1)  # Count the non-padding tokens\n    masked_mean = numer / (denom + eps)  # Compute the masked mean (avoid divide by zero)\n\n    return masked_mean  # Return the final embeddings (mean or [CLS] representation)\n```\nThe helper functions in this code are designed for working with pre-trained BERT models for natural language processing tasks. The exists function checks if a variable is not None. The get_tokenizer function loads and retrieves the pre-trained BERT tokenizer, ensuring it's loaded only once for reuse. Similarly, get_bert loads and retrieves the pre-trained BERT model, avoiding redundant loading.\n\nThe tokenize function takes a string or list of strings as input and outputs their corresponding token IDs, which are returned as a PyTorch tensor. The bert_embed function is responsible for generating text embeddings. It takes token IDs as input, with options for returning the [CLS] token's representation or calculating the average embedding of non-padded tokens using attention masks. It outputs a tensor containing the final text embeddings.\n\n**Input and Outputs:**\n\n* get_tokenizer and get_bert do not take any input and return the BERT tokenizer and model objects, respectively.\n\n* tokenize takes a string or list of strings and returns a PyTorch tensor with token IDs.\n\n* bert_embed takes token IDs, a boolean for [CLS] token usage, a small epsilon value, and a padding ID, returning the text embeddings as a tensor.\n\n### Trainer\n\nWe are getting very close to train our first model, so first we need to define the Trainer class, which orchestrates the entire training procedure for our video generation model. It encapsulates the model, optimizer, data loading, model saving, evaluation, and all the other steps required to train a deep learning model.\n```python\nclass Trainer:\n    def __init__(self, diffusion_model: nn.Module, folder: str, *, ema_decay: float = 0.995, train_batch_size: int = 32, \n                 train_lr: float = 1e-4, train_num_steps: int = 100000, gradient_accumulate_every: int = 2, amp: bool = False, \n                 save_model_every: int = 1000, results_folder: str = './results'):\n        # Initialize trainer, dataset, optimizer, and other configs\n        self.model = diffusion_model  # Diffusion model\n        self.ema = EMA(ema_decay)  # EMA model for averaging weights\n        self.ema_model = copy.deepcopy(self.model)  # Copy for EMA\n        self.batch_size = train_batch_size  # Batch size\n        self.train_num_steps = train_num_steps  # Total training steps\n        self.ds = Dataset(folder, image_size=diffusion_model.image_size)  # Dataset of videos\n        self.dl = cycle(torch.utils.data.DataLoader(self.ds, batch_size=train_batch_size, shuffle=True))  # DataLoader\n        self.opt = Adam(diffusion_model.parameters(), lr=train_lr)  # Optimizer\n        self.step = 0  # Step counter\n        self.amp = amp  # Mixed precision flag\n        self.scaler = GradScaler(enabled=amp)  # Scaler for mixed precision\n        self.results_folder = Path(results_folder)  # Folder to save results\n        self.results_folder.mkdir(exist_ok=True, parents=True)  # Create results folder if doesn't exist\n\n    def reset_parameters(self):\n        # Reset EMA model to match the model's parameters\n        self.ema_model.load_state_dict(self.model.state_dict())\n\n    def step_ema(self):\n        # Update EMA model if training steps exceed threshold\n        if self.step \u003e= 2000:  # Start updating EMA after step 2000\n            self.ema.update_model_average(self.ema_model, self.model)\n\n    def save(self, milestone: int):\n        # Save model, EMA model, and optimizer state at milestones\n        torch.save({'step': self.step, 'model': self.model.state_dict(), 'ema': self.ema_model.state_dict(), 'scaler': self.scaler.state_dict()},\n                   self.results_folder / f'model-{milestone}.pt')\n\n    def load(self, milestone: int):\n        # Load model from checkpoint\n        data = torch.load(self.results_folder / f'model-{milestone}.pt')\n        self.step = data['step']\n        self.model.load_state_dict(data['model'])\n        self.ema_model.load_state_dict(data['ema'])\n        self.scaler.load_state_dict(data['scaler'])\n\n    def train(self, log_fn: Callable[[dict], None] = noop):\n        # Training loop\n        while self.step \u003c self.train_num_steps:\n            for _ in range(self.gradient_accumulate_every):  # Accumulate gradients over multiple steps\n                data = next(self.dl)  # Load data\n                video_data, text_data = data[0].cuda(), data[1] if len(data) == 2 else None  # Move data to GPU\n                with autocast(enabled=self.amp):  # Mixed precision\n                    loss = self.model(video_data, cond=text_data)  # Forward pass\n                    self.scaler.scale(loss / self.gradient_accumulate_every).backward()  # Backpropagate loss\n                print(f'{self.step}: {loss.item()}')  # Print loss\n\n            if self.step % 10 == 0:  # Every 10 steps, update EMA\n                self.step_ema()\n\n            # Optimizer step with gradient clipping if necessary\n            self.scaler.unscale_(self.opt)\n            nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)\n            self.scaler.step(self.opt)\n            self.scaler.update()\n            self.opt.zero_grad()\n\n            # Save model every specified steps\n            if self.step % self.save_model_every == 0:\n                self.save(self.step // self.save_model_every)\n\n            log_fn({'loss': loss.item()})  # Log the loss\n            self.step += 1  # Increment step\n\n        print('Training completed.')\n```\n\nThe Trainer class is responsible for managing the training process of a diffusion model. It initializes with parameters such as the diffusion model, the folder for storing training data, EMA parameters, optimizer settings, and more. It then loads the dataset using a DataLoader, sets up the optimizer, gradient scaler, and EMA model, and creates a results folder for saving checkpoints and sampled videos. The class also includes methods for model parameter management and training progress tracking.\n\nKey methods in the class include reset_parameters, which resets the EMA model to match the current model’s parameters, and step_ema, which updates the EMA model based on the current model's state. The save method saves the model’s parameters, EMA model, optimizer state, and current training step. The load method restores the model from a checkpoint. The train method is the core of the training process, handling the forward pass, loss computation, backpropagation, optimizer steps, model saving, and video sampling as specified in the configuration.\n\nThe input and output for each method are as follows:\n\n* The `__init__` method takes the diffusion model, data folder path, and training parameters as input.\n\n* reset_parameters doesn’t take any input and resets the EMA model’s parameters.\n\n* step_ema doesn’t take input and updates the EMA parameters based on the current model.\n\n* save takes the current training step as input and saves the model checkpoint.\n\n* load takes a training step as input and loads the corresponding model checkpoint.\n\n* train takes training-related parameters and a logging function as input and returns the trained model.\n\n### Configs\n\nNow that we have coded our entire model architecture, we need to define configs that will be used for training and generating videos using our trained model. So let’s defined all of our training and inferencing parameters.\n```yaml\n# Content: Default configuration file for training the TinySora model with text-to-video diffusion\n\ntraining_data_dir: \"./training_data\"  # Directory containing training data (text and video frames)\n\nmodel:\n  dim: 64  # Dimensionality of the model (embedding size)\n  use_bert_text_cond: True  # Enables BERT-based text conditioning for model input\n  dim_mults: [1, 2, 4, 8]  # Scaling factors for each model block (increasing depth)\n  init_dim: null  # Initial dimension, not specified (default is None)\n  init_kernel_size: 7  # Kernel size for initial layers (often used for convolutional layers)\n  use_sparse_linear_attn: True  # Enables sparse attention mechanism for efficiency\n  block_type: \"basic\"  # Type of model block (e.g., 'basic', 'resnet', etc.)\n\ndiffusion:\n  image_size: 32    # Height and width of the video frames\n  num_frames: 5     # Number of frames in the video (sequence length)\n  timesteps: 10   # Number of diffusion timesteps used during training\n  loss_type: \"l1\"   # Loss function for optimization ('l1' means L1 loss)\n  use_dynamic_thres: False  # Whether to use dynamic thresholding during training\n  dynamic_thres_percentile: 0.9  # Threshold percentile used for dynamic thresholding\n\ntrainer:\n  ema_decay: 0.995  # Exponential moving average decay rate for model weights\n  train_batch_size: 2  # Number of samples per batch during training\n  train_lr: 0.0001  # Learning rate for training\n  train_num_steps: 10000  # Total number of training steps (epochs)\n  gradient_accumulate_every: 1  # How often to accumulate gradients (1 means no accumulation)\n  amp: False  # Whether to use automatic mixed precision for training (default: False)\n  step_start_ema: 2000  # Step at which to start applying EMA smoothing\n  update_ema_every: 10  # Frequency of updating EMA weights (every 10 steps)\n  save_model_every: 10  # Save model every 10 steps\n  results_folder: \"./saved_models\"  # Folder where results (models, samples) are saved\n  num_sample_rows: 4  # Number of rows to display during sampling (visualization)\n  max_grad_norm: null  # Maximum gradient norm for clipping (null means no clipping)\n```\n\n### Model instantiation and Training\n\nWe have coded everything its time to define our diffusion model and start training.\n```python\n# Initialize the 3D U-Net model with the configuration parameters for the model.\n# This model is moved to the GPU (cuda).\nmodel = Unet3D(**config['model']).cuda()\n\n# Initialize the GaussianDiffusion model with the U-Net model as the denoising function.\n# Additional configuration parameters for the diffusion process are loaded from the `config['diffusion']`.\n# The model is moved to the GPU (cuda).\ndiffusion = GaussianDiffusion(\n    denoise_fn = model,  # The model will be used to denoise the noisy images during the diffusion process.\n    **config['diffusion']  # Additional diffusion settings, such as timesteps, noise schedules, etc.\n).cuda()\n\n# Initialize the Trainer class with the diffusion model, training configuration, and the folder containing the training data.\n# This is also moved to the GPU.\ntrainer = Trainer(\n    diffusion_model = diffusion,  # The diffusion model to train.\n    **config['trainer'],          # Configuration settings for the training process (e.g., learning rate, batch size).\n    folder = config['training_data_dir']  # Directory where the training data is stored.\n)\n\n# Start the training\ntrainer.train()\n```\n\nIt will start printing the loss after each iteration like this.\n\n```bash\n    0: 0.9512512\n    1: 0.5235211\n    ...\n```\n\nOnce the training gets completed our trained model will be saved under the ./saved_models directory. Now we can use the trained model to generate new output.\n\n### Generating Videos\n\nBefore generating videos we need to code two function that will correct the structure of model output so we can easily visualize the generated video/gif.\n\n```python\ndef generate_video(diffusion: GaussianDiffusion, text: str, batch_size: int, cond_scale: float) -\u003e torch.Tensor:\n    \"\"\"Generate a video using the trained diffusion model.\"\"\"\n    with torch.no_grad():\n        video = diffusion.sample(cond=[text], batch_size=batch_size, cond_scale=cond_scale)\n    return video\n\ndef save_video_as_gif_pil(video_tensor: torch.Tensor, output_path: str) -\u003e None:\n    video_np = (video_tensor.squeeze(0).permute(1, 2, 3, 0).cpu().numpy() * 255).astype(np.uint8)\n    frames = [Image.fromarray(frame) for frame in video_np]\n    frames[0].save(output_path, save_all=True, append_images=frames[1:], duration=100, loop=0)\n    print(f\"Saved GIF: {output_path}\")\n\nLet’s generate a video of our first promptnews_reporter just talking.\n\n# For inference, we should load pre-trained model\nDEFAULT_MODEL_PATH = \"./saved_models\"\nDEFAULT_OUTPUT_DIR = \"./results\"\n\n# Find the latest model checkpoint\nmodel_path = DEFAULT_MODEL_PATH\nif os.path.isdir(model_path):\n        checkpoint_files = [f for f in os.listdir(model_path) if f.endswith(\".pt\")]\n        if not checkpoint_files:\n            raise FileNotFoundError(f\"No model checkpoint found in {model_path}\")\n        checkpoint_files.sort()\n        model_path = os.path.join(model_path, checkpoint_files[-1])\nprint('Loading Model from path: ', model_path)\n\n\ntrainer.load(milestone=-1) # Load the latest model\n\n# Generate video\ntext_prompt = \"News Reporter talking\"\nbatch_size = 1\ncond_scale = 2.0\ngenerated_video = generate_video(diffusion, text_prompt, batch_size, cond_scale)\n\n# Save video\ngif_filename = sanitize_filename(text_prompt) + \".gif\"\noutput_path = os.path.join(DEFAULT_OUTPUT_DIR, gif_filename)\n\n# Create output directory if it doesn't exist\nPath(DEFAULT_OUTPUT_DIR).mkdir(parents=True, exist_ok=True)\n\nsave_video_as_gif_pil(generated_video, output_path)\n```\n\nFollowing are the generated output while running the above code 4 times.\n\n![News Reporter talking prompt](https://cdn-images-1.medium.com/max/2048/1*_2np6scXbzHHkpUeJfX6Fw.gif)\n\nBelow are some outputs of our trained model for 10K Epochs (4 Days 😭)\n\n![Model Result trained for 10K Epochs](https://cdn-images-1.medium.com/max/2560/1*oV4Z__FPants9VwwAfYPRw.gif)\n\n\n## Usage\n\nTO use the training scripts you can simply clone the repository and install the required dependencies. The following steps will guide you through the process:\n```bash\ngit clone https://github.com/FareedKhan-dev/train-text2video-scratch.git\ncd train-text2video-scratch\n```\nIf you encounter any issues with the import paths, you can set the PYTHONPATH environment variable to include the project directory. This will allow Python to find the modules and packages within the project directory.\n\n```bash\nexport PYTHONPATH=\"${PYTHONPATH}:/path/to/text2video-from-scratch\"\n\n# or if you are already in the text2video-from-scratch directory\nexport PYTHONPATH=\"$PYTHONPATH:.\"\n```\n\ncreate a virtual environment and activate it:\n```bash\npython3 -m venv env\nsource env/bin/activate\n```\n\nInstall the required dependencies:\n```bash\npip install -r requirements.txt\n```\n\nYou can modify the architecture under `src/architecture/` and the training parameters under `configs/default.yaml` to suit your needs. The default configuration file contains hyperparameters for training the model.\n\nWe already have a data_generation script that generates synthetic object data or loads the MSRVTT dataset. You can use this script to generate training data for the model. The data generation script is located at `data_generation/synthetic_object.py` and `data_generation/msrvtt.py`.\n\ndata_generation script will generate training data under the default directory `training_data/`. You can modify the script to generate data in a different directory.\n\nOnce you have the training data, you can start training the model using the `train.py` script:\n```bash\npython train.py\n```\n\nThe script will load the configuration from `configs/default.yaml`, create the model, diffusion process, and trainer, and start the training loop. The model checkpoints and logs will be saved in the default directory `./saved_models/`.\n\nYou can modify the configuration file to change the training parameters, such as the learning rate, batch size, number of epochs, etc.\n\nTo generate videos from text prompts using the trained model, you can use the `generate.py` script:\n```bash\npython generate.py --text \"A cat is dancing\"\n```\n\nThe script will load the same configuration used for training, create the model, diffusion process, and generate a video from the input text prompt. The generated video will be saved in the default directory `./results/`.\n\nFollowing parameters can be passed to the `generate.py` script:\n- `--model_path`: Path to the trained model checkpoint (.pt file).\n- `--output_dir`: Directory to save the output GIF.\n- `--text`: Text prompt for video generation.\n- `--batch_size`: Batch size for video generation.\n- `--cond_scale`: Conditioning scale for diffusion sampling.\n\n\nFor the model architecture, the following parameters can be modified:\n\n| Parameter                  | Value                          | Description |\n|----------------------------|--------------------------------|-------------|\n| `dim`                      | 64                             | Dimensionality of the model (embedding size) |\n| `use_bert_text_cond`       | True                           | Enables BERT-based text conditioning for model input |\n| `dim_mults`                | `[1, 2, 4, 8]`                 | Scaling factors for each model block (increasing depth) |\n| `init_dim`                 | `null`                         | Initial dimension, not specified |\n| `init_kernel_size`         | 7                              | Kernel size for initial layers |\n| `use_sparse_linear_attn`   | True                           | Enables sparse attention mechanism for efficiency |\n| `block_type`               | \"basic\"                        | Type of model block (e.g., 'basic', 'resnet', etc.) |\n\nFor the diffusion process, the following parameters can be modified:\n\n| Parameter                 | Value      | Description |\n|---------------------------|------------|-------------|\n| `image_size`              | 32         | Height and width of the video frames |\n| `num_frames`              | 5          | Number of frames in the video sequence |\n| `timesteps`               | 10         | Number of diffusion timesteps used during training |\n| `loss_type`               | \"l1\"       | Loss function for optimization ('l1' means L1 loss) |\n| `use_dynamic_thres`       | False      | Whether to use dynamic thresholding during training |\n| `dynamic_thres_percentile`| 0.9        | Threshold percentile used for dynamic thresholding |\n\nFor the training process, the following parameters can be modified:\n\n| Parameter                   | Value        | Description |\n|-----------------------------|-------------|-------------|\n| `ema_decay`                 | 0.995       | Exponential moving average decay rate for model weights |\n| `train_batch_size`          | 2           | Number of samples per batch during training |\n| `train_lr`                  | 0.0001      | Learning rate for training |\n| `train_num_steps`           | 10000       | Total number of training steps (epochs) |\n| `gradient_accumulate_every` | 1           | How often to accumulate gradients (1 means no accumulation) |\n| `amp`                       | False       | Whether to use automatic mixed precision for training |\n| `step_start_ema`            | 2000        | Step at which to start applying EMA smoothing |\n| `update_ema_every`          | 10          | Frequency of updating EMA weights (every 10 steps) |\n| `save_model_every`          | 10          | Save model every 10 steps |\n| `results_folder`            | \"./saved_models\" | Folder where results (models, samples) are saved |\n| `num_sample_rows`           | 4           | Number of rows to display during sampling (visualization) |\n| `max_grad_norm`             | `null`      | Maximum gradient norm for clipping (null means no clipping) |\n\n\n## References\n\n- [OpenAI SORA Technical Report](https://openai.com/index/video-generation-models-as-world-simulators/)\n- [MSRVTT Dataset](https://paperswithcode.com/dataset/msr-vtt)\n- [PyTorch](https://pytorch.org/)\n- [Hugging Face Transformers](https://huggingface.co/transformers/)\n- [Lightning AI For GPU Resources](https://lightning.ai/)\n\n\n## Contributing\n\nContributions are welcome! Please feel free to submit a pull request or open an issue if you encounter any problems or have any suggestions for improvement.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffareedkhan-dev%2Ftext2video-from-scratch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ffareedkhan-dev%2Ftext2video-from-scratch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffareedkhan-dev%2Ftext2video-from-scratch/lists"}