Add app

app.py

@@ -0,0 +1,218 @@
+from functools import partial
+from typing import Any, Callable, Dict
+
+import clip
+import gradio as gr
+from gradio_rerun import Rerun
+import numpy as np
+from pytorch3d.renderer import TexturesVertex
+from pytorch3d.structures import Meshes
+import rerun as rr
+import torch
+
+from utils.common_viz import init, get_batch
+from utils.random_utils import set_random_seed
+from utils.rerun import log_sample
+from src.diffuser import Diffuser
+from src.datasets.multimodal_dataset import MultimodalDataset
+
+# ------------------------------------------------------------------------------------- #
+
+batch_size, num_cams, num_verts = None, None, None
+
+SAMPLE_IDS = [
+    "2011_KAeAqaA0Llg_00005_00001",
+    "2011_F_EuMeT2wBo_00014_00001",
+    "2011_MCkKihQrNA4_00014_00000",
+]
+LABEL_TO_IDS = {
+    "right": 0,
+    "static": 1,
+    "complex": 2,
+}
+EXAMPLES = [
+    "While the character moves right, the camera trucks right.",
+    "While the character moves right, the camera performs a push in.",
+    "While the character moves right, the camera performs a pull out.",
+    "While the character stays static, the camera performs a boom bottom.",
+    "While the character stays static, the camera performs a boom top.",
+    "While the character moves to the right, the camera trucks right alongside them. Once the character comes to a stop, the camera remains static.",  # noqa
+    "While the character moves to the right, the camera remains static. Once the character comes to a stop, the camera pushes in.",  # noqa
+]
+DEFAULT_TEXT = [
+    "While the character moves right, the camera [...].",
+    "While the character remains static, [...].",
+    "While the character moves to the right, the camera [...]. "
+    "Once the character comes to a stop, the camera [...].",
+]
+
+HEADER = """
+
+<div align="center">
+<h1 style='text-align: center'>E.T. the Exceptional Trajectories</h2>
+<a href="https://robincourant.github.io/info/"><strong>Robin Courant</strong></a>
+·
+<a href="https://nicolas-dufour.github.io/"><strong>Nicolas Dufour</strong></a>
+·
+<a href="https://triocrossing.github.io/"><strong>Xi Wang</strong></a>
+·
+<a href="http://people.irisa.fr/Marc.Christie/"><strong>Marc Christie</strong></a>
+·
+<a href="https://vicky.kalogeiton.info/"><strong>Vicky Kalogeiton</strong></a>
+</div>
+
+
+<div align="center">
+    <a href="https://www.lix.polytechnique.fr/vista/projects/2024_et_courant/" class="button"><b>[Webpage]</b></a> &nbsp;&nbsp;&nbsp;&nbsp;
+    <a href="https://github.com/robincourant/DIRECTOR" class="button"><b>[DIRECTOR]</b></a> &nbsp;&nbsp;&nbsp;&nbsp;
+    <a href="https://github.com/robincourant/CLaTr" class="button"><b>[CLaTr]</b></a> &nbsp;&nbsp;&nbsp;&nbsp;
+    <a href="https://github.com/robincourant/the-exceptional-trajectories" class="button"><b>[Data]</b></a> &nbsp;&nbsp;&nbsp;&nbsp;
+</div>
+
+<br/>
+"""
+
+# ------------------------------------------------------------------------------------- #
+
+
+def get_normals(vertices: torch.Tensor, faces: torch.Tensor) -> torch.Tensor:
+    num_frames, num_faces = vertices.shape[0], faces.shape[-2]
+    faces = faces.expand(num_frames, num_faces, 3)
+
+    verts_rgb = torch.ones_like(vertices)
+    verts_rgb[:, :, 1] = 0
+    textures = TexturesVertex(verts_features=verts_rgb)
+    meshes = Meshes(verts=vertices, faces=faces, textures=textures)
+    normals = meshes.verts_normals_padded()
+
+    return normals, meshes
+
+
+def generate(
+    prompt: str,
+    seed: int,
+    guidance_weight: float,
+    sample_label: str,
+    # ----------------------- ß#
+    dataset: MultimodalDataset,
+    device: torch.device,
+    diffuser: Diffuser,
+    clip_model: clip.model.CLIP,
+) -> Dict[str, Any]:
+    # Set arguments
+    set_random_seed(seed)
+    diffuser.gen_seeds = np.array([seed])
+    diffuser.guidance_weight = guidance_weight
+
+    # Inference
+    sample_id = SAMPLE_IDS[LABEL_TO_IDS[sample_label]]
+    seq_feat = diffuser.net.model.clip_sequential
+    batch = get_batch(prompt, sample_id, clip_model, dataset, seq_feat, device)
+    with torch.no_grad():
+        out = diffuser.predict_step(batch, 0)
+
+    # Run visualization
+    padding_mask = out["padding_mask"][0].to(bool).cpu()
+    padded_traj = out["gen_samples"][0].cpu()
+    traj = padded_traj[padding_mask]
+    padded_vertices = out["char_raw"]["char_vertices"][0]
+    vertices = padded_vertices[padding_mask]
+    faces = out["char_raw"]["char_faces"][0]
+    normals, meshes = get_normals(vertices, faces)
+    fx, fy, cx, cy = out["intrinsics"][0].cpu().numpy()
+    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+    caption = out["caption_raw"][0]
+
+    rr.init(f"{sample_id}")
+    rr.save(".tmp_gr.rrd")
+    log_sample(
+        root_name="world",
+        traj=traj.numpy(),
+        K=K,
+        vertices=vertices.numpy(),
+        faces=faces.numpy(),
+        normals=normals.numpy(),
+        caption=caption,
+        mesh_masks=None,
+    )
+    return "./.tmp_gr.rrd"
+
+
+# ------------------------------------------------------------------------------------- #
+
+
+def main(gen_fn: Callable):
+    theme = gr.themes.Default(primary_hue="blue", secondary_hue="gray")
+
+    with gr.Blocks(theme=theme) as demo:
+        gr.Markdown(HEADER)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Column(scale=2):
+                    sample_str = gr.Dropdown(
+                        choices=["static", "right", "complex"],
+                        label="Character trajectory",
+                        value="right",
+                        interactive=True,
+                    )
+                    text = gr.Textbox(
+                        placeholder="Type the camera motion you want to generate",
+                        show_label=True,
+                        label="Text prompt",
+                        value=DEFAULT_TEXT[LABEL_TO_IDS[sample_str.value]],
+                    )
+                    seed = gr.Number(value=33, label="Seed")
+                    guidance = gr.Slider(0, 10, value=1.4, label="Guidance", step=0.1)
+
+                with gr.Column(scale=1):
+                    btn = gr.Button("Generate", variant="primary")
+
+            with gr.Column(scale=2):
+                examples = gr.Examples(
+                    examples=[[x, None, None] for x in EXAMPLES],
+                    inputs=[text],
+                )
+
+        with gr.Row():
+            output = Rerun()
+
+        def load_example(example_id):
+            processed_example = examples.non_none_processed_examples[example_id]
+            return gr.utils.resolve_singleton(processed_example)
+
+        def change_fn(change):
+            sample_index = LABEL_TO_IDS[change]
+            return gr.update(value=DEFAULT_TEXT[sample_index])
+
+        sample_str.change(fn=change_fn, inputs=[sample_str], outputs=[text])
+
+        inputs = [text, seed, guidance, sample_str]
+        examples.dataset.click(
+            load_example,
+            inputs=[examples.dataset],
+            outputs=examples.inputs_with_examples,
+            show_progress=False,
+            postprocess=False,
+            queue=False,
+        ).then(fn=gen_fn, inputs=inputs, outputs=[output])
+        btn.click(fn=gen_fn, inputs=inputs, outputs=[output])
+        text.submit(fn=gen_fn, inputs=inputs, outputs=[output])
+    demo.launch(share=False)
+
+
+# ------------------------------------------------------------------------------------- #
+
+
+if __name__ == "__main__":
+    # Initialize the models and dataset
+    diffuser, clip_model, dataset, device = init("config")
+    generate_sample = partial(
+        generate,
+        dataset=dataset,
+        device=device,
+        diffuser=diffuser,
+        clip_model=clip_model,
+    )
+
+    main(generate_sample)

configs/compnode/cpu.yaml

@@ -0,0 +1,3 @@
+device: cpu
+num_gpus: 1
+num_workers: 8

configs/config.yaml

@@ -0,0 +1,17 @@
+defaults:
+  - dataset: traj+caption+char
+  - diffuser: rn_director_edm
+  - compnode: cpu
+  - _self_
+
+dataset:
+  char:
+    load_vertices: true
+
+checkpoint_path: 'checkpoints/ca-mixed-e449.ckpt'
+batch_size: 128
+data_dir: data
+
+hydra:
+  run:
+    dir: ./${results_dir}/${xp_name}/${timestamp}

configs/dataset/caption/caption.yaml

@@ -0,0 +1,14 @@
+_target_: src.datasets.modalities.caption_dataset.CaptionDataset
+
+name: caption
+
+dataset_dir: ${dataset.dataset_dir}
+segment_dir: ${dataset.dataset_dir}/cam_segments
+raw_caption_dir: ${dataset.dataset_dir}/caption
+feat_caption_dir: ${dataset.dataset_dir}/caption_clip
+
+num_segments: 27
+num_feats: 512
+num_cams: ${dataset.standardization.num_cams}
+sequential: ${diffuser.network.module.clip_sequential}
+max_feat_length: 77

configs/dataset/char/char.yaml

@@ -0,0 +1,15 @@
+_target_: src.datasets.modalities.char_dataset.CharacterDataset
+
+name: char
+dataset_dir: ${dataset.dataset_dir}
+num_cams: ${dataset.num_cams}
+num_raw_feats: 3
+num_frequencies: 10
+min_freq: 0
+max_freq: 4
+num_encoding: 3 # ${eval:'2 * ${dataset.char.num_frequencies} * ${dataset.char.num_raw_feats}'}
+sequential: ${diffuser.network.module.cond_sequential}
+num_feats: ${eval:'${dataset.char.num_encoding} if ${dataset.char.sequential} else ${dataset.num_cams} * ${dataset.char.num_encoding}'}
+standardize: ${dataset.trajectory.standardize}
+standardization: ${dataset.standardization}
+load_vertices: ${diffuser.do_projection}

configs/dataset/standardization/0300.yaml

@@ -0,0 +1,15 @@
+name: '0300'
+num_interframes: 0
+num_cams: 300
+num_total_frames: ${eval:'${dataset.standardization.num_interframes} * (${dataset.standardization.num_cams} - 1) + ${dataset.standardization.num_cams} '}
+
+norm_mean: [7.93987673e-05, -9.98621393e-05, 4.12940653e-04]
+norm_std: [0.027841, 0.01819818, 0.03138536]
+
+shift_mean: [0.00201079, -0.27488501, -1.23616805]
+shift_std: [1.13433516, 1.19061042, 1.58744263]
+
+norm_mean_h: [6.676e-05, -5.084e-05, -7.782e-04]
+norm_std_h: [0.0105, 0.006958, 0.01145]
+
+velocity: true

configs/dataset/traj+caption+char.yaml

@@ -0,0 +1,18 @@
+_target_: src.datasets.multimodal_dataset.MultimodalDataset
+
+defaults:
+  - _self_
+  - trajectory: rot6d_trajectory
+  - char: char
+  - caption: caption
+  - standardization: '0300'
+
+name: "${dataset.standardization.name}-t:${dataset.trajectory.name}|c:${dataset.caption.name}|h:${dataset.char.name}"
+dataset_name: ${dataset.standardization.name}
+dataset_dir: ${data_dir}
+
+num_rawfeats: 12
+num_cams: ${dataset.standardization.num_cams}
+feature_type: ${dataset.trajectory.name}
+num_feats: ${dataset.trajectory.num_feats}
+num_cond_feats: ['${dataset.char.num_feats}','${dataset.caption.num_feats}']

configs/dataset/trajectory/rot6d_trajectory.yaml

@@ -0,0 +1,10 @@
+_target_: src.datasets.modalities.trajectory_dataset.TrajectoryDataset
+
+name: rot6d
+set_name: null
+dataset_dir: ${dataset.dataset_dir}
+num_feats: 9
+num_rawfeats: ${dataset.num_rawfeats}
+num_cams: ${dataset.num_cams}
+standardize: true
+standardization: ${dataset.standardization}

configs/diffuser/network/module/ca_director.yaml

@@ -0,0 +1,17 @@
+_target_: src.models.modules.director.CrossAttentionDirector
+name: ca_director
+num_feats: ${dataset.num_feats}
+num_rawfeats: ${dataset.num_rawfeats}
+num_cams: ${dataset.num_cams}
+num_cond_feats: ${dataset.num_cond_feats}
+latent_dim: 512
+mlp_multiplier: 4
+num_layers: 8
+num_heads: 16
+dropout: 0.1
+stochastic_depth: 0.1
+label_dropout: 0.1
+num_text_registers: 16
+clip_sequential: True
+cond_sequential: True
+device: ${compnode.device}

configs/diffuser/network/rn_director.yaml

@@ -0,0 +1,7 @@
+_target_: src.models.networks.RnEDMPrecond
+
+defaults:
+  - module: ca_director
+
+name: rn_director
+sigma_data: 0.5

configs/diffuser/rn_director_edm.yaml

@@ -0,0 +1,24 @@
+_target_: src.diffuser.Diffuser
+
+defaults:
+  - _self_
+  - network: rn_director
+
+guidance_weight: 1.4
+edm2_normalization: true
+
+# EMA
+ema_kwargs:
+  beta: 0.9999
+  update_every: 1
+
+# Sampling
+sampling_kwargs:
+  num_steps: 10
+  sigma_min: 0.002
+  sigma_max: 80
+  rho: 40
+  S_churn: 0
+  S_min: 0
+  S_max: inf
+  S_noise: 1

src/datasets/datamodule.py

@@ -0,0 +1,68 @@
+from lightning import LightningDataModule
+from torch.utils.data import Dataset, DataLoader
+
+
+class Datamodule(LightningDataModule):
+    def __init__(
+        self,
+        train_dataset: Dataset,
+        eval_dataset: Dataset,
+        batch_train_size: int,
+        num_workers: int,
+        eval_batch_size: int = None,
+    ):
+        super().__init__()
+
+        self.train_dataset = train_dataset
+        self.eval_dataset = eval_dataset
+
+        self.batch_train_size = batch_train_size
+        self.eval_batch_size = (
+            eval_batch_size if eval_batch_size is not None else batch_train_size
+        )
+
+        self.num_workers = num_workers
+
+    def train_dataloader(self) -> DataLoader:
+        """Load train set loader."""
+        persistent_workers = True if self.num_workers > 0 else False
+
+        dataloader = DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_train_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            persistent_workers=persistent_workers,
+        )
+        return dataloader
+
+    def val_dataloader(self) -> DataLoader:
+        """Load val set loader."""
+        persistent_workers = True if self.num_workers > 0 else False
+
+        dataloader = DataLoader(
+            self.eval_dataset,
+            batch_size=self.eval_batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            persistent_workers=persistent_workers,
+        )
+        return dataloader
+
+    def predict_dataloader(self) -> DataLoader:
+        """Load predict set loader."""
+        dataloader = DataLoader(
+            self.eval_dataset,
+            batch_size=self.eval_batch_size,
+            num_workers=self.num_workers,
+        )
+        return dataloader
+
+    def test_dataloader(self) -> DataLoader:
+        """Load test set loader."""
+        dataloader = DataLoader(
+            self.eval_dataset,
+            batch_size=self.eval_batch_size,
+            num_workers=self.num_workers,
+        )
+        return dataloader

src/datasets/modalities/caption_dataset.py

@@ -0,0 +1,107 @@
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+
+from utils.file_utils import load_txt
+
+
+class CaptionDataset(Dataset):
+    def __init__(
+        self,
+        name: str,
+        dataset_dir: str,
+        num_cams: int,
+        num_feats: int,
+        num_segments: int,
+        sequential: bool,
+        **kwargs,
+    ):
+        super().__init__()
+        self.modality = name
+        self.name = name
+        self.dataset_dir = Path(dataset_dir)
+        # Set data paths (segments, captions, etc...)
+        for name, field in kwargs.items():
+            if isinstance(field, str):
+                field = Path(field)
+            if name == "feat_caption_dir":
+                field = field / "seq" if sequential else field / "token"
+            setattr(self, name, field)
+
+        self.filenames = None
+
+        self.clip_seq_dir = self.dataset_dir / "caption_clip" / "seq"  # For CLaTrScore
+        self.num_cams = num_cams
+        self.num_feats = num_feats
+        self.num_segments = num_segments
+        self.sequential = sequential
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, index):
+        filename = self.filenames[index]
+
+        # Load data
+        if hasattr(self, "segment_dir"):
+            raw_segments = torch.from_numpy(
+                np.load((self.segment_dir / (filename + ".npy")))
+            )
+            padded_raw_segments = F.pad(
+                raw_segments,
+                (0, self.num_cams - len(raw_segments)),
+                value=self.num_segments,
+            )
+        if hasattr(self, "raw_caption_dir"):
+            raw_caption = load_txt(self.raw_caption_dir / (filename + ".txt"))
+        if hasattr(self, "feat_caption_dir"):
+            feat_caption = torch.from_numpy(
+                np.load((self.feat_caption_dir / (filename + ".npy")))
+            )
+            if self.sequential:
+                feat_caption = F.pad(
+                    feat_caption.to(torch.float32),
+                    (0, 0, 0, self.max_feat_length - feat_caption.shape[0]),
+                )
+
+        if self.modality == "caption":
+            raw_data = {"caption": raw_caption, "segments": padded_raw_segments}
+            feat_data = (
+                feat_caption.permute(1, 0) if feat_caption.dim() == 2 else feat_caption
+            )
+        elif self.modality == "segments":
+            raw_data = {"segments": padded_raw_segments}
+            # Shift by one for padding
+            feat_data = F.one_hot(
+                padded_raw_segments, num_classes=self.num_segments + 1
+            ).to(torch.float32)
+            if self.sequential:
+                feat_data = feat_data.permute(1, 0)
+            else:
+                feat_data = feat_data.reshape(-1)
+        elif self.modality == "class":
+            raw_data = {"segments": padded_raw_segments}
+            most_frequent_segment = Counter(raw_segments).most_common(1)[0][0]
+            feat_data = F.one_hot(
+                torch.tensor(most_frequent_segment), num_classes=self.num_segments
+            ).to(torch.float32)
+        else:
+            raise ValueError(f"Modality {self.modality} not supported")
+
+        clip_seq_caption = torch.from_numpy(
+            np.load((self.clip_seq_dir / (filename + ".npy")))
+        )
+        padding_mask = torch.ones((self.max_feat_length))
+        padding_mask[clip_seq_caption.shape[0] :] = 0
+        clip_seq_caption = F.pad(
+            clip_seq_caption.to(torch.float32),
+            (0, 0, 0, self.max_feat_length - clip_seq_caption.shape[0]),
+        )
+        raw_data["clip_seq_caption"] = clip_seq_caption
+        raw_data["clip_seq_mask"] = padding_mask
+
+        return filename, feat_data, raw_data

src/datasets/modalities/char_dataset.py

@@ -0,0 +1,120 @@
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+
+# ------------------------------------------------------------------------------------- #
+
+num_frequencies = None
+
+# ------------------------------------------------------------------------------------- #
+
+
+class CharacterDataset(Dataset):
+    def __init__(
+        self,
+        name: str,
+        dataset_dir: str,
+        standardize: bool,
+        num_feats: int,
+        num_cams: int,
+        sequential: bool,
+        num_frequencies: int,
+        min_freq: int,
+        max_freq: int,
+        load_vertices: bool,
+        **kwargs,
+    ):
+        super().__init__()
+        self.modality = "char"
+        self.name = name
+        self.dataset_dir = Path(dataset_dir)
+        self.traj_dir = self.dataset_dir / "traj"
+        self.data_dir = self.dataset_dir / self.name
+        self.vert_dir = self.dataset_dir / "vert_raw"
+        self.center_dir = self.dataset_dir / "char_raw"
+
+        self.filenames = None
+        self.standardize = standardize
+        if self.standardize:
+            mean_std = kwargs["standardization"]
+            self.norm_mean = torch.Tensor(mean_std["norm_mean_h"])[:, None]
+            self.norm_std = torch.Tensor(mean_std["norm_std_h"])[:, None]
+            self.velocity = mean_std["velocity"]
+
+        self.num_cams = num_cams
+        self.num_feats = num_feats
+        self.sequential = sequential
+        self.num_frequencies = num_frequencies
+        self.min_freq = min_freq
+        self.max_freq = max_freq
+
+        self.load_vertices = load_vertices
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, index):
+        filename = self.filenames[index]
+
+        char_filename = filename + ".npy"
+        char_path = self.data_dir / char_filename
+
+        raw_char_feature = torch.from_numpy(np.load((char_path))).to(torch.float32)
+        padding_size = self.num_cams - raw_char_feature.shape[0]
+        padded_raw_char_feature = F.pad(
+            raw_char_feature, (0, 0, 0, padding_size)
+        ).permute(1, 0)
+
+        center_path = self.center_dir / char_filename  # Center to offset mesh
+        center_offset = torch.from_numpy(np.load(center_path)[0]).to(torch.float32)
+        if self.load_vertices:
+            vert_path = self.vert_dir / char_filename
+            raw_verts = np.load(vert_path, allow_pickle=True)[()]
+            if raw_verts["vertices"] is None:
+                num_frames = raw_char_feature.shape[0]
+                verts = torch.zeros((num_frames, 6890, 3), dtype=torch.float32)
+                padded_verts = torch.zeros(
+                    (self.num_cams, 6890, 3), dtype=torch.float32
+                )
+                faces = torch.zeros((13776, 3), dtype=torch.int16)
+            else:
+                verts = torch.from_numpy(raw_verts["vertices"]).to(torch.float32)
+                verts -= center_offset
+                padded_verts = F.pad(verts, (0, 0, 0, 0, 0, padding_size))
+                faces = torch.from_numpy(raw_verts["faces"]).to(torch.int16)
+
+        char_feature = raw_char_feature.clone()
+        if self.velocity:
+            velocity = char_feature[1:].clone() - char_feature[:-1].clone()
+            char_feature = torch.cat([raw_char_feature[0][None], velocity])
+
+        if self.standardize:
+            # Normalize the first frame (orgin) and the rest (velocity) separately
+            if len(self.norm_mean) == 6:
+                char_feature[0] -= self.norm_mean[:3, 0].to(raw_char_feature.device)
+                char_feature[0] /= self.norm_std[:3, 0].to(raw_char_feature.device)
+                char_feature[1:] -= self.norm_mean[3:, 0].to(raw_char_feature.device)
+                char_feature[1:] /= self.norm_std[3:, 0].to(raw_char_feature.device)
+            # Normalize all in one
+            else:
+                char_feature -= self.norm_mean[:, 0].to(raw_char_feature.device)
+                char_feature /= self.norm_std[:, 0].to(raw_char_feature.device)
+        padded_char_feature = F.pad(
+            char_feature,
+            (0, 0, 0, self.num_cams - char_feature.shape[0]),
+        )
+
+        if self.sequential:
+            padded_char_feature = padded_char_feature.permute(1, 0)
+        else:
+            padded_char_feature = padded_char_feature.reshape(-1)
+
+        raw_feats = {"char_raw_feat": padded_raw_char_feature}
+        if self.load_vertices:
+            raw_feats["char_vertices"] = padded_verts
+            raw_feats["char_faces"] = faces
+
+        return char_filename, padded_char_feature, raw_feats

src/datasets/modalities/trajectory_dataset.py

@@ -0,0 +1,152 @@
+from pathlib import Path
+
+from evo.tools.file_interface import read_kitti_poses_file
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from torchtyping import TensorType
+import torch.nn.functional as F
+from typing import Tuple
+
+from utils.file_utils import load_txt
+from utils.rotation_utils import compute_rotation_matrix_from_ortho6d
+
+num_cams = None
+
+
+# ------------------------------------------------------------------------------------- #
+
+
+class TrajectoryDataset(Dataset):
+    def __init__(
+        self,
+        name: str,
+        set_name: str,
+        dataset_dir: str,
+        num_rawfeats: int,
+        num_feats: int,
+        num_cams: int,
+        standardize: bool,
+        **kwargs,
+    ):
+        super().__init__()
+        self.name = name
+        self.set_name = set_name
+        self.dataset_dir = Path(dataset_dir)
+        if name == "relative":
+            self.data_dir = self.dataset_dir / "traj_raw"
+            self.relative_dir = self.dataset_dir / "relative"
+        else:
+            self.data_dir = self.dataset_dir / "traj"
+        self.intrinsics_dir = self.dataset_dir / "intrinsics"
+
+        self.num_rawfeats = num_rawfeats
+        self.num_feats = num_feats
+        self.num_cams = num_cams
+
+        self.augmentation = None
+        self.standardize = standardize
+        if self.standardize:
+            mean_std = kwargs["standardization"]
+            self.norm_mean = torch.Tensor(mean_std["norm_mean"])
+            self.norm_std = torch.Tensor(mean_std["norm_std"])
+            self.shift_mean = torch.Tensor(mean_std["shift_mean"])
+            self.shift_std = torch.Tensor(mean_std["shift_std"])
+            self.velocity = mean_std["velocity"]
+
+    # --------------------------------------------------------------------------------- #
+
+    def set_split(self, split: str, train_rate: float = 1.0):
+        self.split = split
+        split_path = Path(self.dataset_dir) / f"{split}_split.txt"
+        split_traj = load_txt(split_path).split("\n")
+        self.filenames = sorted(split_traj)
+
+        return self
+
+    # --------------------------------------------------------------------------------- #
+
+    def get_feature(
+        self, raw_matrix_trajectory: TensorType["num_cams", 4, 4]
+    ) -> TensorType[9, "num_cams"]:
+        matrix_trajectory = torch.clone(raw_matrix_trajectory)
+
+        raw_trans = torch.clone(matrix_trajectory[:, :3, 3])
+        if self.velocity:
+            velocity = raw_trans[1:] - raw_trans[:-1]
+            raw_trans = torch.cat([raw_trans[0][None], velocity])
+        if self.standardize:
+            raw_trans[0] -= self.shift_mean
+            raw_trans[0] /= self.shift_std
+            raw_trans[1:] -= self.norm_mean
+            raw_trans[1:] /= self.norm_std
+
+        # Compute the 6D continuous rotation
+        raw_rot = matrix_trajectory[:, :3, :3]
+        rot6d = raw_rot[:, :, :2].permute(0, 2, 1).reshape(-1, 6)
+
+        # Stack rotation 6D and translation
+        rot6d_trajectory = torch.hstack([rot6d, raw_trans]).permute(1, 0)
+
+        return rot6d_trajectory
+
+    def get_matrix(
+        self, raw_rot6d_trajectory: TensorType[9, "num_cams"]
+    ) -> TensorType["num_cams", 4, 4]:
+        rot6d_trajectory = torch.clone(raw_rot6d_trajectory)
+        device = rot6d_trajectory.device
+
+        num_cams = rot6d_trajectory.shape[1]
+        matrix_trajectory = torch.eye(4, device=device)[None].repeat(num_cams, 1, 1)
+
+        raw_trans = rot6d_trajectory[6:].permute(1, 0)
+        if self.standardize:
+            raw_trans[0] *= self.shift_std.to(device)
+            raw_trans[0] += self.shift_mean.to(device)
+            raw_trans[1:] *= self.norm_std.to(device)
+            raw_trans[1:] += self.norm_mean.to(device)
+        if self.velocity:
+            raw_trans = torch.cumsum(raw_trans, dim=0)
+        matrix_trajectory[:, :3, 3] = raw_trans
+
+        rot6d = rot6d_trajectory[:6].permute(1, 0)
+        raw_rot = compute_rotation_matrix_from_ortho6d(rot6d)
+        matrix_trajectory[:, :3, :3] = raw_rot
+
+        return matrix_trajectory
+
+    # --------------------------------------------------------------------------------- #
+
+    def __getitem__(self, index: int) -> Tuple[str, TensorType["num_cams", 4, 4]]:
+        filename = self.filenames[index]
+
+        trajectory_filename = filename + ".txt"
+        trajectory_path = self.data_dir / trajectory_filename
+
+        trajectory = read_kitti_poses_file(trajectory_path)
+        matrix_trajectory = torch.from_numpy(np.array(trajectory.poses_se3)).to(
+            torch.float32
+        )
+
+        trajectory_feature = self.get_feature(matrix_trajectory)
+
+        padded_trajectory_feature = F.pad(
+            trajectory_feature, (0, self.num_cams - trajectory_feature.shape[1])
+        )
+        # Padding mask: 1 for valid cams, 0 for padded cams
+        padding_mask = torch.ones((self.num_cams))
+        padding_mask[trajectory_feature.shape[1] :] = 0
+
+        intrinsics_filename = filename + ".npy"
+        intrinsics_path = self.intrinsics_dir / intrinsics_filename
+        intrinsics = np.load(intrinsics_path)
+
+        return (
+            trajectory_filename,
+            padded_trajectory_feature,
+            padding_mask,
+            intrinsics
+        )
+
+    def __len__(self):
+        return len(self.filenames)

src/datasets/multimodal_dataset.py

@@ -0,0 +1,88 @@
+from copy import deepcopy as dp
+from pathlib import Path
+
+from torch.utils.data import Dataset
+
+
+class MultimodalDataset(Dataset):
+    def __init__(
+        self,
+        name,
+        dataset_name,
+        dataset_dir,
+        trajectory,
+        feature_type,
+        num_rawfeats,
+        num_feats,
+        num_cams,
+        num_cond_feats,
+        standardization,
+        augmentation=None,
+        **modalities,
+    ):
+        self.dataset_dir = Path(dataset_dir)
+        self.name = name
+        self.dataset_name = dataset_name
+        self.feature_type = feature_type
+        self.num_rawfeats = num_rawfeats
+        self.num_feats = num_feats
+        self.num_cams = num_cams
+        self.trajectory_dataset = trajectory
+        self.standardization = standardization
+        self.modality_datasets = modalities
+
+        if augmentation is not None:
+            self.augmentation = True
+            self.augmentation_rate = augmentation.rate
+            self.trajectory_dataset.set_augmentation(augmentation.trajectory)
+            if hasattr(augmentation, "modalities"):
+                for modality, augments in augmentation.modalities:
+                    self.modality_datasets[modality].set_augmentation(augments)
+        else:
+            self.augmentation = False
+
+    # --------------------------------------------------------------------------------- #
+
+    def set_split(self, split: str, train_rate: float = 1.0):
+        self.split = split
+
+        # Get trajectory split
+        self.trajectory_dataset = dp(self.trajectory_dataset).set_split(
+            split, train_rate
+        )
+        self.root_filenames = self.trajectory_dataset.filenames
+
+        # Get modality split
+        for modality_name in self.modality_datasets.keys():
+            self.modality_datasets[modality_name].filenames = self.root_filenames
+
+        self.get_feature = self.trajectory_dataset.get_feature
+        self.get_matrix = self.trajectory_dataset.get_matrix
+
+        return self
+
+    # --------------------------------------------------------------------------------- #
+
+    def __getitem__(self, index):
+        traj_out = self.trajectory_dataset[index]
+        traj_filename, traj_feature, padding_mask, intrinsics = traj_out
+
+        out = {
+            "traj_filename": traj_filename,
+            "traj_feat": traj_feature,
+            "padding_mask": padding_mask,
+            "intrinsics": intrinsics,
+        }
+
+        for modality_name, modality_dataset in self.modality_datasets.items():
+            modality_filename, modality_feature, modality_raw = modality_dataset[index]
+            assert traj_filename.split(".")[0] == modality_filename.split(".")[0]
+            out[f"{modality_name}_filename"] = modality_filename
+            out[f"{modality_name}_feat"] = modality_feature
+            out[f"{modality_name}_raw"] = modality_raw
+            out[f"{modality_name}_padding_mask"] = padding_mask
+
+        return out
+
+    def __len__(self):
+        return len(self.trajectory_dataset)

src/diffuser.py

@@ -0,0 +1,221 @@
+from omegaconf.dictconfig import DictConfig
+from typing import List, Tuple
+
+from ema_pytorch import EMA
+import numpy as np
+import torch
+from torchtyping import TensorType
+import torch.nn as nn
+import lightning as L
+
+from utils.random_utils import StackedRandomGenerator
+
+# ------------------------------------------------------------------------------------- #
+
+batch_size, num_samples = None, None
+num_feats, num_rawfeats, num_cams = None, None, None
+RawTrajectory = TensorType["num_samples", "num_rawfeats", "num_cams"]
+
+# ------------------------------------------------------------------------------------- #
+
+
+class Diffuser(L.LightningModule):
+    def __init__(
+        self,
+        network: nn.Module,
+        guidance_weight: float,
+        ema_kwargs: DictConfig,
+        sampling_kwargs: DictConfig,
+        edm2_normalization: bool,
+        **kwargs,
+    ):
+        super().__init__()
+
+        # Network and EMA
+        self.net = network
+        self.ema = EMA(self.net, **ema_kwargs)
+        self.guidance_weight = guidance_weight
+        self.edm2_normalization = edm2_normalization
+        self.sigma_data = network.sigma_data
+
+        # Sampling
+        self.num_steps = sampling_kwargs.num_steps
+        self.sigma_min = sampling_kwargs.sigma_min
+        self.sigma_max = sampling_kwargs.sigma_max
+        self.rho = sampling_kwargs.rho
+        self.S_churn = sampling_kwargs.S_churn
+        self.S_noise = sampling_kwargs.S_noise
+        self.S_min = sampling_kwargs.S_min
+        self.S_max = (
+            sampling_kwargs.S_max
+            if isinstance(sampling_kwargs.S_max, float)
+            else float("inf")
+        )
+
+    # ---------------------------------------------------------------------------------- #
+
+    def on_predict_start(self):
+        eval_dataset = self.trainer.datamodule.eval_dataset
+        self.modalities = list(eval_dataset.modality_datasets.keys())
+
+        self.get_matrix = self.trainer.datamodule.train_dataset.get_matrix
+        self.v_get_matrix = self.trainer.datamodule.eval_dataset.get_matrix
+
+    def predict_step(self, batch, batch_idx):
+        ref_samples, mask = batch["traj_feat"], batch["padding_mask"]
+
+        if len(self.modalities) > 0:
+            cond_k = [x for x in batch.keys() if "traj" not in x and "feat" in x]
+            cond_data = [batch[cond] for cond in cond_k]
+            conds = {}
+            for cond in cond_k:
+                cond_name = cond.replace("_feat", "")
+                if isinstance(batch[f"{cond_name}_raw"], dict):
+                    for cond_name_, x in batch[f"{cond_name}_raw"].items():
+                        conds[cond_name_] = x
+                else:
+                    conds[cond_name] = batch[f"{cond_name}_raw"]
+            batch["conds"] = conds
+        else:
+            cond_data = None
+
+        # cf edm2 sigma_data normalization / https://arxiv.org/pdf/2312.02696.pdf
+        if self.edm2_normalization:
+            ref_samples *= self.sigma_data
+        _, gen_samples = self.sample(self.ema.ema_model, ref_samples, cond_data, mask)
+
+        batch["ref_samples"] = torch.stack([self.v_get_matrix(x) for x in ref_samples])
+        batch["gen_samples"] = torch.stack([self.get_matrix(x) for x in gen_samples])
+
+        return batch
+
+    # --------------------------------------------------------------------------------- #
+
+    def sample(
+        self,
+        net: torch.nn.Module,
+        traj_samples: RawTrajectory,
+        cond_samples: TensorType["num_samples", "num_feats"],
+        mask: TensorType["num_samples", "num_feats"],
+        external_seeds: List[int] = None,
+    ) -> Tuple[RawTrajectory, RawTrajectory]:
+        # Pick latents
+        num_samples = traj_samples.shape[0]
+        seeds = self.gen_seeds if hasattr(self, "gen_seeds") else range(num_samples)
+        rnd = StackedRandomGenerator(self.device, seeds)
+
+        sz = [num_samples, self.net.num_feats, self.net.num_cams]
+        latents = rnd.randn_rn(sz, device=self.device)
+        # Generate trajectories.
+        generations = self.edm_sampler(
+            net,
+            latents,
+            class_labels=cond_samples,
+            mask=mask,
+            randn_like=rnd.randn_like,
+            guidance_weight=self.guidance_weight,
+            # ----------------------------------- #
+            num_steps=self.num_steps,
+            sigma_min=self.sigma_min,
+            sigma_max=self.sigma_max,
+            rho=self.rho,
+            S_churn=self.S_churn,
+            S_min=self.S_min,
+            S_max=self.S_max,
+            S_noise=self.S_noise,
+        )
+
+        return latents, generations
+
+    @staticmethod
+    def edm_sampler(
+        net,
+        latents,
+        class_labels=None,
+        mask=None,
+        guidance_weight=2.0,
+        randn_like=torch.randn_like,
+        num_steps=18,
+        sigma_min=0.002,
+        sigma_max=80,
+        rho=7,
+        S_churn=0,
+        S_min=0,
+        S_max=float("inf"),
+        S_noise=1,
+    ):
+        # Time step discretization.
+        step_indices = torch.arange(num_steps, device=latents.device)
+        t_steps = (
+            sigma_max ** (1 / rho)
+            + step_indices
+            / (num_steps - 1)
+            * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))
+        ) ** rho
+        t_steps = torch.cat(
+            [torch.as_tensor(t_steps), torch.zeros_like(t_steps[:1])]
+        )  # t_N = 0
+
+        # Main sampling loop.
+        bool_mask = ~mask.to(bool)
+        x_next = latents * t_steps[0]
+        bs = latents.shape[0]
+        for i, (t_cur, t_next) in enumerate(
+            zip(t_steps[:-1], t_steps[1:])
+        ):  # 0, ..., N-1
+            x_cur = x_next
+
+            # Increase noise temporarily.
+            gamma = (
+                min(S_churn / num_steps, np.sqrt(2) - 1)
+                if S_min <= t_cur <= S_max
+                else 0
+            )
+            t_hat = torch.as_tensor(t_cur + gamma * t_cur)
+            x_hat = x_cur + (t_hat**2 - t_cur**2).sqrt() * S_noise * randn_like(x_cur)
+
+            # Euler step.
+            if class_labels is not None:
+                class_label_knot = [torch.zeros_like(label) for label in class_labels]
+                x_hat_both = torch.cat([x_hat, x_hat], dim=0)
+                y_label_both = [
+                    torch.cat([y, y_knot], dim=0)
+                    for y, y_knot in zip(class_labels, class_label_knot)
+                ]
+
+                bool_mask_both = torch.cat([bool_mask, bool_mask], dim=0)
+                t_hat_both = torch.cat([t_hat.expand(bs), t_hat.expand(bs)], dim=0)
+                cond_denoised, denoised = net(
+                    x_hat_both, t_hat_both, y=y_label_both, mask=bool_mask_both
+                ).chunk(2, dim=0)
+                denoised = denoised + (cond_denoised - denoised) * guidance_weight
+            else:
+                denoised = net(x_hat, t_hat.expand(bs), mask=bool_mask)
+            d_cur = (x_hat - denoised) / t_hat
+            x_next = x_hat + (t_next - t_hat) * d_cur
+
+            # Apply 2nd order correction.
+            if i < num_steps - 1:
+                if class_labels is not None:
+                    class_label_knot = [
+                        torch.zeros_like(label) for label in class_labels
+                    ]
+                    x_next_both = torch.cat([x_next, x_next], dim=0)
+                    y_label_both = [
+                        torch.cat([y, y_knot], dim=0)
+                        for y, y_knot in zip(class_labels, class_label_knot)
+                    ]
+                    bool_mask_both = torch.cat([bool_mask, bool_mask], dim=0)
+                    t_next_both = torch.cat(
+                        [t_next.expand(bs), t_next.expand(bs)], dim=0
+                    )
+                    cond_denoised, denoised = net(
+                        x_next_both, t_next_both, y=y_label_both, mask=bool_mask_both
+                    ).chunk(2, dim=0)
+                    denoised = denoised + (cond_denoised - denoised) * guidance_weight
+                else:
+                    denoised = net(x_next, t_next.expand(bs), mask=bool_mask)
+                d_prime = (x_next - denoised) / t_next
+                x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime)
+
+        return x_next

src/models/modules/director.py

@@ -0,0 +1,1154 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+import numpy as np
+from einops import rearrange
+
+from typing import Optional, List
+from torchtyping import TensorType
+from einops._torch_specific import allow_ops_in_compiled_graph  # requires einops>=0.6.1
+
+allow_ops_in_compiled_graph()
+
+batch_size, num_cond_feats = None, None
+
+
+class FusedMLP(nn.Sequential):
+    def __init__(
+        self,
+        dim_model: int,
+        dropout: float,
+        activation: nn.Module,
+        hidden_layer_multiplier: int = 4,
+        bias: bool = True,
+    ):
+        super().__init__(
+            nn.Linear(dim_model, dim_model * hidden_layer_multiplier, bias=bias),
+            activation(),
+            nn.Dropout(dropout),
+            nn.Linear(dim_model * hidden_layer_multiplier, dim_model, bias=bias),
+        )
+
+
+def _cast_if_autocast_enabled(tensor):
+    if torch.is_autocast_enabled():
+        if tensor.device.type == "cuda":
+            dtype = torch.get_autocast_gpu_dtype()
+        elif tensor.device.type == "cpu":
+            dtype = torch.get_autocast_cpu_dtype()
+        else:
+            raise NotImplementedError()
+        return tensor.to(dtype=dtype)
+    return tensor
+
+
+class LayerNorm16Bits(torch.nn.LayerNorm):
+    """
+    16-bit friendly version of torch.nn.LayerNorm
+    """
+
+    def __init__(
+        self,
+        normalized_shape,
+        eps=1e-06,
+        elementwise_affine=True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+            device=device,
+            dtype=dtype,
+        )
+
+    def forward(self, x):
+        module_device = x.device
+        downcast_x = _cast_if_autocast_enabled(x)
+        downcast_weight = (
+            _cast_if_autocast_enabled(self.weight)
+            if self.weight is not None
+            else self.weight
+        )
+        downcast_bias = (
+            _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
+        )
+        with torch.autocast(enabled=False, device_type=module_device.type):
+            return nn.functional.layer_norm(
+                downcast_x,
+                self.normalized_shape,
+                downcast_weight,
+                downcast_bias,
+                self.eps,
+            )
+
+
+class StochatichDepth(nn.Module):
+    def __init__(self, p: float):
+        super().__init__()
+        self.survival_prob = 1.0 - p
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.training and self.survival_prob < 1:
+            mask = (
+                torch.empty(x.shape[0], 1, 1, device=x.device).uniform_()
+                + self.survival_prob
+            )
+            mask = mask.floor()
+            if self.survival_prob > 0:
+                mask = mask / self.survival_prob
+            return x * mask
+        else:
+            return x
+
+
+class CrossAttentionOp(nn.Module):
+    def __init__(
+        self, attention_dim, num_heads, dim_q, dim_kv, use_biases=True, is_sa=False
+    ):
+        super().__init__()
+        self.dim_q = dim_q
+        self.dim_kv = dim_kv
+        self.attention_dim = attention_dim
+        self.num_heads = num_heads
+        self.use_biases = use_biases
+        self.is_sa = is_sa
+        if self.is_sa:
+            self.qkv = nn.Linear(dim_q, attention_dim * 3, bias=use_biases)
+        else:
+            self.q = nn.Linear(dim_q, attention_dim, bias=use_biases)
+            self.kv = nn.Linear(dim_kv, attention_dim * 2, bias=use_biases)
+        self.out = nn.Linear(attention_dim, dim_q, bias=use_biases)
+
+    def forward(self, x_to, x_from=None, attention_mask=None):
+        if x_from is None:
+            x_from = x_to
+        if self.is_sa:
+            q, k, v = self.qkv(x_to).chunk(3, dim=-1)
+        else:
+            q = self.q(x_to)
+            k, v = self.kv(x_from).chunk(2, dim=-1)
+        q = rearrange(q, "b n (h d) -> b h n d", h=self.num_heads)
+        k = rearrange(k, "b n (h d) -> b h n d", h=self.num_heads)
+        v = rearrange(v, "b n (h d) -> b h n d", h=self.num_heads)
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1)
+        x = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=attention_mask
+        )
+        x = rearrange(x, "b h n d -> b n (h d)")
+        x = self.out(x)
+        return x
+
+
+class CrossAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim_q: int,
+        dim_kv: int,
+        num_heads: int,
+        attention_dim: int = 0,
+        mlp_multiplier: int = 4,
+        dropout: float = 0.0,
+        stochastic_depth: float = 0.0,
+        use_biases: bool = True,
+        retrieve_attention_scores: bool = False,
+        use_layernorm16: bool = True,
+    ):
+        super().__init__()
+        layer_norm = (
+            nn.LayerNorm
+            if not use_layernorm16 or retrieve_attention_scores
+            else LayerNorm16Bits
+        )
+        self.retrieve_attention_scores = retrieve_attention_scores
+        self.initial_to_ln = layer_norm(dim_q, eps=1e-6)
+        attention_dim = min(dim_q, dim_kv) if attention_dim == 0 else attention_dim
+        self.ca = CrossAttentionOp(
+            attention_dim, num_heads, dim_q, dim_kv, is_sa=False, use_biases=use_biases
+        )
+        self.ca_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.middle_ln = layer_norm(dim_q, eps=1e-6)
+        self.ffn = FusedMLP(
+            dim_model=dim_q,
+            dropout=dropout,
+            activation=nn.GELU,
+            hidden_layer_multiplier=mlp_multiplier,
+            bias=use_biases,
+        )
+        self.ffn_stochastic_depth = StochatichDepth(stochastic_depth)
+
+    def forward(
+        self,
+        to_tokens: Tensor,
+        from_tokens: Tensor,
+        to_token_mask: Optional[Tensor] = None,
+        from_token_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        if to_token_mask is None and from_token_mask is None:
+            attention_mask = None
+        else:
+            if to_token_mask is None:
+                to_token_mask = torch.ones(
+                    to_tokens.shape[0],
+                    to_tokens.shape[1],
+                    dtype=torch.bool,
+                    device=to_tokens.device,
+                )
+            if from_token_mask is None:
+                from_token_mask = torch.ones(
+                    from_tokens.shape[0],
+                    from_tokens.shape[1],
+                    dtype=torch.bool,
+                    device=from_tokens.device,
+                )
+            attention_mask = from_token_mask.unsqueeze(1) * to_token_mask.unsqueeze(2)
+        attention_output = self.ca(
+            self.initial_to_ln(to_tokens),
+            from_tokens,
+            attention_mask=attention_mask,
+        )
+        to_tokens = to_tokens + self.ca_stochastic_depth(attention_output)
+        to_tokens = to_tokens + self.ffn_stochastic_depth(
+            self.ffn(self.middle_ln(to_tokens))
+        )
+        return to_tokens
+
+
+class SelfAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim_qkv: int,
+        num_heads: int,
+        attention_dim: int = 0,
+        mlp_multiplier: int = 4,
+        dropout: float = 0.0,
+        stochastic_depth: float = 0.0,
+        use_biases: bool = True,
+        use_layer_scale: bool = False,
+        layer_scale_value: float = 0.0,
+        use_layernorm16: bool = True,
+    ):
+        super().__init__()
+        layer_norm = LayerNorm16Bits if use_layernorm16 else nn.LayerNorm
+        self.initial_ln = layer_norm(dim_qkv, eps=1e-6)
+        attention_dim = dim_qkv if attention_dim == 0 else attention_dim
+        self.sa = CrossAttentionOp(
+            attention_dim,
+            num_heads,
+            dim_qkv,
+            dim_qkv,
+            is_sa=True,
+            use_biases=use_biases,
+        )
+        self.sa_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.middle_ln = layer_norm(dim_qkv, eps=1e-6)
+        self.ffn = FusedMLP(
+            dim_model=dim_qkv,
+            dropout=dropout,
+            activation=nn.GELU,
+            hidden_layer_multiplier=mlp_multiplier,
+            bias=use_biases,
+        )
+        self.ffn_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_1 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+            self.layer_scale_2 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        token_mask: Optional[torch.Tensor] = None,
+    ):
+        if token_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = token_mask.unsqueeze(1) * torch.ones(
+                tokens.shape[0],
+                tokens.shape[1],
+                1,
+                dtype=torch.bool,
+                device=tokens.device,
+            )
+        attention_output = self.sa(
+            self.initial_ln(tokens),
+            attention_mask=attention_mask,
+        )
+        if self.use_layer_scale:
+            tokens = tokens + self.sa_stochastic_depth(
+                self.layer_scale_1 * attention_output
+            )
+            tokens = tokens + self.ffn_stochastic_depth(
+                self.layer_scale_2 * self.ffn(self.middle_ln(tokens))
+            )
+        else:
+            tokens = tokens + self.sa_stochastic_depth(attention_output)
+            tokens = tokens + self.ffn_stochastic_depth(
+                self.ffn(self.middle_ln(tokens))
+            )
+        return tokens
+
+
+class AdaLNSABlock(nn.Module):
+    def __init__(
+        self,
+        dim_qkv: int,
+        dim_cond: int,
+        num_heads: int,
+        attention_dim: int = 0,
+        mlp_multiplier: int = 4,
+        dropout: float = 0.0,
+        stochastic_depth: float = 0.0,
+        use_biases: bool = True,
+        use_layer_scale: bool = False,
+        layer_scale_value: float = 0.1,
+        use_layernorm16: bool = True,
+    ):
+        super().__init__()
+        layer_norm = LayerNorm16Bits if use_layernorm16 else nn.LayerNorm
+        self.initial_ln = layer_norm(dim_qkv, eps=1e-6, elementwise_affine=False)
+        attention_dim = dim_qkv if attention_dim == 0 else attention_dim
+        self.adaln_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim_cond, dim_qkv * 6, bias=use_biases),
+        )
+        # Zero init
+        nn.init.zeros_(self.adaln_modulation[1].weight)
+        nn.init.zeros_(self.adaln_modulation[1].bias)
+
+        self.sa = CrossAttentionOp(
+            attention_dim,
+            num_heads,
+            dim_qkv,
+            dim_qkv,
+            is_sa=True,
+            use_biases=use_biases,
+        )
+        self.sa_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.middle_ln = layer_norm(dim_qkv, eps=1e-6, elementwise_affine=False)
+        self.ffn = FusedMLP(
+            dim_model=dim_qkv,
+            dropout=dropout,
+            activation=nn.GELU,
+            hidden_layer_multiplier=mlp_multiplier,
+            bias=use_biases,
+        )
+        self.ffn_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_1 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+            self.layer_scale_2 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        cond: torch.Tensor,
+        token_mask: Optional[torch.Tensor] = None,
+    ):
+        if token_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = token_mask.unsqueeze(1) * torch.ones(
+                tokens.shape[0],
+                tokens.shape[1],
+                1,
+                dtype=torch.bool,
+                device=tokens.device,
+            )
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaln_modulation(cond).chunk(6, dim=-1)
+        )
+        attention_output = self.sa(
+            modulate_shift_and_scale(self.initial_ln(tokens), shift_msa, scale_msa),
+            attention_mask=attention_mask,
+        )
+        if self.use_layer_scale:
+            tokens = tokens + self.sa_stochastic_depth(
+                gate_msa.unsqueeze(1) * self.layer_scale_1 * attention_output
+            )
+            tokens = tokens + self.ffn_stochastic_depth(
+                gate_mlp.unsqueeze(1)
+                * self.layer_scale_2
+                * self.ffn(
+                    modulate_shift_and_scale(
+                        self.middle_ln(tokens), shift_mlp, scale_mlp
+                    )
+                )
+            )
+        else:
+            tokens = tokens + gate_msa.unsqueeze(1) * self.sa_stochastic_depth(
+                attention_output
+            )
+            tokens = tokens + self.ffn_stochastic_depth(
+                gate_mlp.unsqueeze(1)
+                * self.ffn(
+                    modulate_shift_and_scale(
+                        self.middle_ln(tokens), shift_mlp, scale_mlp
+                    )
+                )
+            )
+        return tokens
+
+
+class CrossAttentionSABlock(nn.Module):
+    def __init__(
+        self,
+        dim_qkv: int,
+        dim_cond: int,
+        num_heads: int,
+        attention_dim: int = 0,
+        mlp_multiplier: int = 4,
+        dropout: float = 0.0,
+        stochastic_depth: float = 0.0,
+        use_biases: bool = True,
+        use_layer_scale: bool = False,
+        layer_scale_value: float = 0.0,
+        use_layernorm16: bool = True,
+    ):
+        super().__init__()
+        layer_norm = LayerNorm16Bits if use_layernorm16 else nn.LayerNorm
+        attention_dim = dim_qkv if attention_dim == 0 else attention_dim
+        self.ca = CrossAttentionOp(
+            attention_dim,
+            num_heads,
+            dim_qkv,
+            dim_cond,
+            is_sa=False,
+            use_biases=use_biases,
+        )
+        self.ca_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.ca_ln = layer_norm(dim_qkv, eps=1e-6)
+
+        self.initial_ln = layer_norm(dim_qkv, eps=1e-6)
+        attention_dim = dim_qkv if attention_dim == 0 else attention_dim
+
+        self.sa = CrossAttentionOp(
+            attention_dim,
+            num_heads,
+            dim_qkv,
+            dim_qkv,
+            is_sa=True,
+            use_biases=use_biases,
+        )
+        self.sa_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.middle_ln = layer_norm(dim_qkv, eps=1e-6)
+        self.ffn = FusedMLP(
+            dim_model=dim_qkv,
+            dropout=dropout,
+            activation=nn.GELU,
+            hidden_layer_multiplier=mlp_multiplier,
+            bias=use_biases,
+        )
+        self.ffn_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_1 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+            self.layer_scale_2 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        cond: torch.Tensor,
+        token_mask: Optional[torch.Tensor] = None,
+        cond_mask: Optional[torch.Tensor] = None,
+    ):
+        if cond_mask is None:
+            cond_attention_mask = None
+        else:
+            cond_attention_mask = torch.ones(
+                cond.shape[0],
+                1,
+                cond.shape[1],
+                dtype=torch.bool,
+                device=tokens.device,
+            ) * token_mask.unsqueeze(2)
+        if token_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = token_mask.unsqueeze(1) * torch.ones(
+                tokens.shape[0],
+                tokens.shape[1],
+                1,
+                dtype=torch.bool,
+                device=tokens.device,
+            )
+        ca_output = self.ca(
+            self.ca_ln(tokens),
+            cond,
+            attention_mask=cond_attention_mask,
+        )
+        ca_output = torch.nan_to_num(
+            ca_output, nan=0.0, posinf=0.0, neginf=0.0
+        )  # Needed as some tokens get attention from no token so Nan
+        tokens = tokens + self.ca_stochastic_depth(ca_output)
+        attention_output = self.sa(
+            self.initial_ln(tokens),
+            attention_mask=attention_mask,
+        )
+        if self.use_layer_scale:
+            tokens = tokens + self.sa_stochastic_depth(
+                self.layer_scale_1 * attention_output
+            )
+            tokens = tokens + self.ffn_stochastic_depth(
+                self.layer_scale_2 * self.ffn(self.middle_ln(tokens))
+            )
+        else:
+            tokens = tokens + self.sa_stochastic_depth(attention_output)
+            tokens = tokens + self.ffn_stochastic_depth(
+                self.ffn(self.middle_ln(tokens))
+            )
+        return tokens
+
+
+class CAAdaLNSABlock(nn.Module):
+    def __init__(
+        self,
+        dim_qkv: int,
+        dim_cond: int,
+        num_heads: int,
+        attention_dim: int = 0,
+        mlp_multiplier: int = 4,
+        dropout: float = 0.0,
+        stochastic_depth: float = 0.0,
+        use_biases: bool = True,
+        use_layer_scale: bool = False,
+        layer_scale_value: float = 0.1,
+        use_layernorm16: bool = True,
+    ):
+        super().__init__()
+        layer_norm = LayerNorm16Bits if use_layernorm16 else nn.LayerNorm
+        self.ca = CrossAttentionOp(
+            attention_dim,
+            num_heads,
+            dim_qkv,
+            dim_cond,
+            is_sa=False,
+            use_biases=use_biases,
+        )
+        self.ca_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.ca_ln = layer_norm(dim_qkv, eps=1e-6)
+        self.initial_ln = layer_norm(dim_qkv, eps=1e-6)
+        attention_dim = dim_qkv if attention_dim == 0 else attention_dim
+        self.adaln_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim_cond, dim_qkv * 6, bias=use_biases),
+        )
+        # Zero init
+        nn.init.zeros_(self.adaln_modulation[1].weight)
+        nn.init.zeros_(self.adaln_modulation[1].bias)
+
+        self.sa = CrossAttentionOp(
+            attention_dim,
+            num_heads,
+            dim_qkv,
+            dim_qkv,
+            is_sa=True,
+            use_biases=use_biases,
+        )
+        self.sa_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.middle_ln = layer_norm(dim_qkv, eps=1e-6)
+        self.ffn = FusedMLP(
+            dim_model=dim_qkv,
+            dropout=dropout,
+            activation=nn.GELU,
+            hidden_layer_multiplier=mlp_multiplier,
+            bias=use_biases,
+        )
+        self.ffn_stochastic_depth = StochatichDepth(stochastic_depth)
+        self.use_layer_scale = use_layer_scale
+        if use_layer_scale:
+            self.layer_scale_1 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+            self.layer_scale_2 = nn.Parameter(
+                torch.ones(dim_qkv) * layer_scale_value, requires_grad=True
+            )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        cond_1: torch.Tensor,
+        cond_2: torch.Tensor,
+        cond_1_mask: Optional[torch.Tensor] = None,
+        token_mask: Optional[torch.Tensor] = None,
+    ):
+        if token_mask is None and cond_1_mask is None:
+            cond_attention_mask = None
+        elif token_mask is None:
+            cond_attention_mask = cond_1_mask.unsqueeze(1) * torch.ones(
+                cond_1.shape[0],
+                cond_1.shape[1],
+                1,
+                dtype=torch.bool,
+                device=cond_1.device,
+            )
+        elif cond_1_mask is None:
+            cond_attention_mask = torch.ones(
+                tokens.shape[0],
+                1,
+                tokens.shape[1],
+                dtype=torch.bool,
+                device=tokens.device,
+            ) * token_mask.unsqueeze(2)
+        else:
+            cond_attention_mask = cond_1_mask.unsqueeze(1) * token_mask.unsqueeze(2)
+        if token_mask is None:
+            attention_mask = None
+        else:
+            attention_mask = token_mask.unsqueeze(1) * torch.ones(
+                tokens.shape[0],
+                tokens.shape[1],
+                1,
+                dtype=torch.bool,
+                device=tokens.device,
+            )
+        ca_output = self.ca(
+            self.ca_ln(tokens),
+            cond_1,
+            attention_mask=cond_attention_mask,
+        )
+        ca_output = torch.nan_to_num(ca_output, nan=0.0, posinf=0.0, neginf=0.0)
+        tokens = tokens + self.ca_stochastic_depth(ca_output)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.adaln_modulation(cond_2).chunk(6, dim=-1)
+        )
+        attention_output = self.sa(
+            modulate_shift_and_scale(self.initial_ln(tokens), shift_msa, scale_msa),
+            attention_mask=attention_mask,
+        )
+        if self.use_layer_scale:
+            tokens = tokens + self.sa_stochastic_depth(
+                gate_msa.unsqueeze(1) * self.layer_scale_1 * attention_output
+            )
+            tokens = tokens + self.ffn_stochastic_depth(
+                gate_mlp.unsqueeze(1)
+                * self.layer_scale_2
+                * self.ffn(
+                    modulate_shift_and_scale(
+                        self.middle_ln(tokens), shift_mlp, scale_mlp
+                    )
+                )
+            )
+        else:
+            tokens = tokens + gate_msa.unsqueeze(1) * self.sa_stochastic_depth(
+                attention_output
+            )
+            tokens = tokens + self.ffn_stochastic_depth(
+                gate_mlp.unsqueeze(1)
+                * self.ffn(
+                    modulate_shift_and_scale(
+                        self.middle_ln(tokens), shift_mlp, scale_mlp
+                    )
+                )
+            )
+        return tokens
+
+
+class PositionalEmbedding(nn.Module):
+    """
+    Taken from https://github.com/NVlabs/edm
+    """
+
+    def __init__(self, num_channels, max_positions=10000, endpoint=False):
+        super().__init__()
+        self.num_channels = num_channels
+        self.max_positions = max_positions
+        self.endpoint = endpoint
+        freqs = torch.arange(start=0, end=self.num_channels // 2, dtype=torch.float32)
+        freqs = 2 * freqs / self.num_channels
+        freqs = (1 / self.max_positions) ** freqs
+        self.register_buffer("freqs", freqs)
+
+    def forward(self, x):
+        x = torch.outer(x, self.freqs)
+        out = torch.cat([x.cos(), x.sin()], dim=1)
+        return out.to(x.dtype)
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.0, max_len=10000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # not used in the final model
+        x = x + self.pe[:, : x.shape[1], :]
+        return self.dropout(x)
+
+
+class TimeEmbedder(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        time_scaling: float,
+        expansion: int = 4,
+    ):
+        super().__init__()
+        self.encode_time = PositionalEmbedding(num_channels=dim, endpoint=True)
+
+        self.time_scaling = time_scaling
+        self.map_time = nn.Sequential(
+            nn.Linear(dim, dim * expansion),
+            nn.SiLU(),
+            nn.Linear(dim * expansion, dim * expansion),
+        )
+
+    def forward(self, t: Tensor) -> Tensor:
+        time = self.encode_time(t * self.time_scaling)
+        time_mean = time.mean(dim=-1, keepdim=True)
+        time_std = time.std(dim=-1, keepdim=True)
+        time = (time - time_mean) / time_std
+        return self.map_time(time)
+
+
+def modulate_shift_and_scale(x: Tensor, shift: Tensor, scale: Tensor) -> Tensor:
+    return x * (1 + scale).unsqueeze(1) + shift.unsqueeze(1)
+
+
+# ------------------------------------------------------------------------------------- #
+
+
+class BaseDirector(nn.Module):
+    def __init__(
+        self,
+        name: str,
+        num_feats: int,
+        num_cond_feats: int,
+        num_cams: int,
+        latent_dim: int,
+        mlp_multiplier: int,
+        num_layers: int,
+        num_heads: int,
+        dropout: float,
+        stochastic_depth: float,
+        label_dropout: float,
+        num_rawfeats: int,
+        clip_sequential: bool = False,
+        cond_sequential: bool = False,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        super().__init__()
+        self.name = name
+        self.label_dropout = label_dropout
+        self.num_rawfeats = num_rawfeats
+        self.num_feats = num_feats
+        self.num_cams = num_cams
+        self.clip_sequential = clip_sequential
+        self.cond_sequential = cond_sequential
+        self.use_layernorm16 = device == "cuda"
+
+        self.input_projection = nn.Sequential(
+            nn.Linear(num_feats, latent_dim),
+            PositionalEncoding(latent_dim),
+        )
+        self.time_embedding = TimeEmbedder(latent_dim // 4, time_scaling=1000)
+        self.init_conds_mappings(num_cond_feats, latent_dim)
+        self.init_backbone(
+            num_layers, latent_dim, mlp_multiplier, num_heads, dropout, stochastic_depth
+        )
+        self.init_output_projection(num_feats, latent_dim)
+
+    def forward(
+        self,
+        x: Tensor,
+        timesteps: Tensor,
+        y: List[Tensor] = None,
+        mask: Tensor = None,
+    ) -> Tensor:
+        mask = mask.logical_not() if mask is not None else None
+        x = rearrange(x, "b c n -> b n c")
+        x = self.input_projection(x)
+        t = self.time_embedding(timesteps)
+        if y is not None:
+            y = self.mask_cond(y)
+            y = self.cond_mapping(y, mask, t)
+
+        x = self.backbone(x, y, mask)
+        x = self.output_projection(x, y)
+        return rearrange(x, "b n c -> b c n")
+
+    def init_conds_mappings(self, num_cond_feats, latent_dim):
+        raise NotImplementedError(
+            "This method should be implemented in the derived class"
+        )
+
+    def init_backbone(self):
+        raise NotImplementedError(
+            "This method should be implemented in the derived class"
+        )
+
+    def cond_mapping(self, cond: List[Tensor], mask: Tensor, t: Tensor) -> Tensor:
+        raise NotImplementedError(
+            "This method should be implemented in the derived class"
+        )
+
+    def backbone(self, x: Tensor, y: Tensor, mask: Tensor) -> Tensor:
+        raise NotImplementedError(
+            "This method should be implemented in the derived class"
+        )
+
+    def mask_cond(
+        self, cond: List[TensorType["batch_size", "num_cond_feats"]]
+    ) -> TensorType["batch_size", "num_cond_feats"]:
+        bs = cond[0].shape[0]
+        if self.training and self.label_dropout > 0.0:
+            # 1-> use null_cond, 0-> use real cond
+            prob = torch.ones(bs, device=cond[0].device) * self.label_dropout
+            masked_cond = []
+            common_mask = torch.bernoulli(prob)  # Common to all modalities
+            for _cond in cond:
+                modality_mask = torch.bernoulli(prob)  # Modality only
+                mask = torch.clip(common_mask + modality_mask, 0, 1)
+                mask = mask.view(bs, 1, 1) if _cond.dim() == 3 else mask.view(bs, 1)
+                masked_cond.append(_cond * (1.0 - mask))
+            return masked_cond
+        else:
+            return cond
+
+    def init_output_projection(self, num_feats, latent_dim):
+        raise NotImplementedError(
+            "This method should be implemented in the derived class"
+        )
+
+    def output_projection(self, x: Tensor, y: Tensor) -> Tensor:
+        raise NotImplementedError(
+            "This method should be implemented in the derived class"
+        )
+
+
+class AdaLNDirector(BaseDirector):
+    def __init__(
+        self,
+        name: str,
+        num_feats: int,
+        num_cond_feats: int,
+        num_cams: int,
+        latent_dim: int,
+        mlp_multiplier: int,
+        num_layers: int,
+        num_heads: int,
+        dropout: float,
+        stochastic_depth: float,
+        label_dropout: float,
+        num_rawfeats: int,
+        clip_sequential: bool = False,
+        cond_sequential: bool = False,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            num_feats=num_feats,
+            num_cond_feats=num_cond_feats,
+            num_cams=num_cams,
+            latent_dim=latent_dim,
+            mlp_multiplier=mlp_multiplier,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            dropout=dropout,
+            stochastic_depth=stochastic_depth,
+            label_dropout=label_dropout,
+            num_rawfeats=num_rawfeats,
+            clip_sequential=clip_sequential,
+            cond_sequential=cond_sequential,
+            device=device,
+        )
+        assert not (clip_sequential and cond_sequential)
+
+    def init_conds_mappings(self, num_cond_feats, latent_dim):
+        self.joint_cond_projection = nn.Linear(sum(num_cond_feats), latent_dim)
+
+    def cond_mapping(self, cond: List[Tensor], mask: Tensor, t: Tensor) -> Tensor:
+        c_emb = torch.cat(cond, dim=-1)
+        return self.joint_cond_projection(c_emb) + t
+
+    def init_backbone(
+        self,
+        num_layers,
+        latent_dim,
+        mlp_multiplier,
+        num_heads,
+        dropout,
+        stochastic_depth,
+    ):
+        self.backbone_module = nn.ModuleList(
+            [
+                AdaLNSABlock(
+                    dim_qkv=latent_dim,
+                    dim_cond=latent_dim,
+                    num_heads=num_heads,
+                    mlp_multiplier=mlp_multiplier,
+                    dropout=dropout,
+                    stochastic_depth=stochastic_depth,
+                    use_layernorm16=self.use_layernorm16,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def backbone(self, x: Tensor, y: Tensor, mask: Tensor) -> Tensor:
+        for block in self.backbone_module:
+            x = block(x, y, mask)
+        return x
+
+    def init_output_projection(self, num_feats, latent_dim):
+        layer_norm = LayerNorm16Bits if self.use_layernorm16 else nn.LayerNorm
+
+        self.final_norm = layer_norm(latent_dim, eps=1e-6, elementwise_affine=False)
+        self.final_linear = nn.Linear(latent_dim, num_feats, bias=True)
+        self.final_adaln = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(latent_dim, latent_dim * 2, bias=True),
+        )
+        # Zero init
+        nn.init.zeros_(self.final_adaln[1].weight)
+        nn.init.zeros_(self.final_adaln[1].bias)
+
+    def output_projection(self, x: Tensor, y: Tensor) -> Tensor:
+        shift, scale = self.final_adaln(y).chunk(2, dim=-1)
+        x = modulate_shift_and_scale(self.final_norm(x), shift, scale)
+        return self.final_linear(x)
+
+
+class CrossAttentionDirector(BaseDirector):
+    def __init__(
+        self,
+        name: str,
+        num_feats: int,
+        num_cond_feats: int,
+        num_cams: int,
+        latent_dim: int,
+        mlp_multiplier: int,
+        num_layers: int,
+        num_heads: int,
+        dropout: float,
+        stochastic_depth: float,
+        label_dropout: float,
+        num_rawfeats: int,
+        num_text_registers: int,
+        clip_sequential: bool = True,
+        cond_sequential: bool = True,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        self.num_text_registers = num_text_registers
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.mlp_multiplier = mlp_multiplier
+        self.stochastic_depth = stochastic_depth
+        super().__init__(
+            name=name,
+            num_feats=num_feats,
+            num_cond_feats=num_cond_feats,
+            num_cams=num_cams,
+            latent_dim=latent_dim,
+            mlp_multiplier=mlp_multiplier,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            dropout=dropout,
+            stochastic_depth=stochastic_depth,
+            label_dropout=label_dropout,
+            num_rawfeats=num_rawfeats,
+            clip_sequential=clip_sequential,
+            cond_sequential=cond_sequential,
+            device=device,
+        )
+        assert clip_sequential and cond_sequential
+
+    def init_conds_mappings(self, num_cond_feats, latent_dim):
+        self.cond_projection = nn.ModuleList(
+            [nn.Linear(num_cond_feat, latent_dim) for num_cond_feat in num_cond_feats]
+        )
+        self.cond_registers = nn.Parameter(
+            torch.randn(self.num_text_registers, latent_dim), requires_grad=True
+        )
+        nn.init.trunc_normal_(self.cond_registers, std=0.02, a=-2 * 0.02, b=2 * 0.02)
+        self.cond_sa = nn.ModuleList(
+            [
+                SelfAttentionBlock(
+                    dim_qkv=latent_dim,
+                    num_heads=self.num_heads,
+                    mlp_multiplier=self.mlp_multiplier,
+                    dropout=self.dropout,
+                    stochastic_depth=self.stochastic_depth,
+                    use_layernorm16=self.use_layernorm16,
+                )
+                for _ in range(2)
+            ]
+        )
+        self.cond_positional_embedding = PositionalEncoding(latent_dim, max_len=10000)
+
+    def cond_mapping(self, cond: List[Tensor], mask: Tensor, t: Tensor) -> Tensor:
+        batch_size = cond[0].shape[0]
+        cond_emb = [
+            cond_proj(rearrange(c, "b c n -> b n c"))
+            for cond_proj, c in zip(self.cond_projection, cond)
+        ]
+        cond_emb = [
+            self.cond_registers.unsqueeze(0).expand(batch_size, -1, -1),
+            t.unsqueeze(1),
+        ] + cond_emb
+        cond_emb = torch.cat(cond_emb, dim=1)
+        cond_emb = self.cond_positional_embedding(cond_emb)
+        for block in self.cond_sa:
+            cond_emb = block(cond_emb)
+        return cond_emb
+
+    def init_backbone(
+        self,
+        num_layers,
+        latent_dim,
+        mlp_multiplier,
+        num_heads,
+        dropout,
+        stochastic_depth,
+    ):
+        self.backbone_module = nn.ModuleList(
+            [
+                CrossAttentionSABlock(
+                    dim_qkv=latent_dim,
+                    dim_cond=latent_dim,
+                    num_heads=num_heads,
+                    mlp_multiplier=mlp_multiplier,
+                    dropout=dropout,
+                    stochastic_depth=stochastic_depth,
+                    use_layernorm16=self.use_layernorm16,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def backbone(self, x: Tensor, y: Tensor, mask: Tensor) -> Tensor:
+        for block in self.backbone_module:
+            x = block(x, y, mask, None)
+        return x
+
+    def init_output_projection(self, num_feats, latent_dim):
+        layer_norm = LayerNorm16Bits if self.use_layernorm16 else nn.LayerNorm
+
+        self.final_norm = layer_norm(latent_dim, eps=1e-6)
+        self.final_linear = nn.Linear(latent_dim, num_feats, bias=True)
+
+    def output_projection(self, x: Tensor, y: Tensor) -> Tensor:
+        return self.final_linear(self.final_norm(x))
+
+
+class InContextDirector(BaseDirector):
+    def __init__(
+        self,
+        name: str,
+        num_feats: int,
+        num_cond_feats: int,
+        num_cams: int,
+        latent_dim: int,
+        mlp_multiplier: int,
+        num_layers: int,
+        num_heads: int,
+        dropout: float,
+        stochastic_depth: float,
+        label_dropout: float,
+        num_rawfeats: int,
+        clip_sequential: bool = False,
+        cond_sequential: bool = False,
+        device: str = "cuda",
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            num_feats=num_feats,
+            num_cond_feats=num_cond_feats,
+            num_cams=num_cams,
+            latent_dim=latent_dim,
+            mlp_multiplier=mlp_multiplier,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            dropout=dropout,
+            stochastic_depth=stochastic_depth,
+            label_dropout=label_dropout,
+            num_rawfeats=num_rawfeats,
+            clip_sequential=clip_sequential,
+            cond_sequential=cond_sequential,
+            device=device,
+        )
+
+    def init_conds_mappings(self, num_cond_feats, latent_dim):
+        self.cond_projection = nn.ModuleList(
+            [nn.Linear(num_cond_feat, latent_dim) for num_cond_feat in num_cond_feats]
+        )
+
+    def cond_mapping(self, cond: List[Tensor], mask: Tensor, t: Tensor) -> Tensor:
+        for i in range(len(cond)):
+            if cond[i].dim() == 3:
+                cond[i] = rearrange(cond[i], "b c n -> b n c")
+        cond_emb = [cond_proj(c) for cond_proj, c in zip(self.cond_projection, cond)]
+        cond_emb = [c.unsqueeze(1) if c.dim() == 2 else cond_emb for c in cond_emb]
+        cond_emb = torch.cat([t.unsqueeze(1)] + cond_emb, dim=1)
+        return cond_emb
+
+    def init_backbone(
+        self,
+        num_layers,
+        latent_dim,
+        mlp_multiplier,
+        num_heads,
+        dropout,
+        stochastic_depth,
+    ):
+        self.backbone_module = nn.ModuleList(
+            [
+                SelfAttentionBlock(
+                    dim_qkv=latent_dim,
+                    num_heads=num_heads,
+                    mlp_multiplier=mlp_multiplier,
+                    dropout=dropout,
+                    stochastic_depth=stochastic_depth,
+                    use_layernorm16=self.use_layernorm16,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+    def backbone(self, x: Tensor, y: Tensor, mask: Tensor) -> Tensor:
+        bs, n_y, _ = y.shape
+        mask = torch.cat([torch.ones(bs, n_y, device=y.device), mask], dim=1)
+        x = torch.cat([y, x], dim=1)
+        for block in self.backbone_module:
+            x = block(x, mask)
+        return x
+
+    def init_output_projection(self, num_feats, latent_dim):
+        layer_norm = LayerNorm16Bits if self.use_layernorm16 else nn.LayerNorm
+
+        self.final_norm = layer_norm(latent_dim, eps=1e-6)
+        self.final_linear = nn.Linear(latent_dim, num_feats, bias=True)
+
+    def output_projection(self, x: Tensor, y: Tensor) -> Tensor:
+        num_y = y.shape[1]
+        x = x[:, num_y:]
+        return self.final_linear(self.final_norm(x))

src/models/networks.py

@@ -0,0 +1,33 @@
+import torch.nn as nn
+
+
+# ----------------------------------------------------------------------------
+# Improved preconditioning proposed in the paper "Elucidating the Design
+# Space of Diffusion-Based Generative Models" (EDM).
+
+
+class RnEDMPrecond(nn.Module):
+    def __init__(self, sigma_data: float = 0.5, module: nn.Module = None, **kwargs):
+        super().__init__()
+        self.sigma_data = sigma_data
+
+        self.model = module
+        self.num_rawfeats = module.num_rawfeats
+        self.num_feats = module.num_feats
+        self.num_cams = module.num_cams
+
+    def forward(self, x, sigma, y=None, mask=None):
+        """
+        x: [batch_size, num_feats, max_frames], denoted x_t in the paper
+        sigma: [batch_size] (int)
+        """
+        sigma = sigma.reshape(-1, 1, 1)
+        c_skip = self.sigma_data**2 / (sigma**2 + self.sigma_data**2)
+        c_out = sigma * self.sigma_data / (sigma**2 + self.sigma_data**2).sqrt()
+        c_in = 1 / (self.sigma_data**2 + sigma**2).sqrt()
+        c_noise = sigma.log() / 4
+
+        F_x = self.model(c_in * x, c_noise.flatten(), y=y, mask=mask)
+        D_x = c_skip * x + c_out * F_x
+
+        return D_x

utils/common_viz.py

@@ -0,0 +1,136 @@
+from typing import Any, Dict, List, Tuple
+
+import clip
+from hydra import compose, initialize
+from hydra.utils import instantiate
+from omegaconf import OmegaConf
+import torch
+from torchtyping import TensorType
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
+from src.diffuser import Diffuser
+from src.datasets.multimodal_dataset import MultimodalDataset
+
+# ------------------------------------------------------------------------------------- #
+
+batch_size, context_length = None, None
+collate_fn = DataLoader([]).collate_fn
+
+# ------------------------------------------------------------------------------------- #
+
+
+def to_device(batch: Dict[str, Any], device: torch.device) -> Dict[str, Any]:
+    for key, value in batch.items():
+        if isinstance(value, torch.Tensor):
+            batch[key] = value.to(device)
+    return batch
+
+
+def load_clip_model(version: str, device: str) -> clip.model.CLIP:
+    model, _ = clip.load(version, device=device, jit=False)
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad = False
+    return model
+
+
+def encode_text(
+    caption_raws: List[str],  # batch_size
+    clip_model: clip.model.CLIP,
+    max_token_length: int,
+    device: str,
+) -> TensorType["batch_size", "context_length"]:
+    if max_token_length is not None:
+        default_context_length = 77
+        context_length = max_token_length + 2  # start_token + 20 + end_token
+        assert context_length < default_context_length
+        # [bs, context_length] # if n_tokens > context_length -> will truncate
+        texts = clip.tokenize(
+            caption_raws, context_length=context_length, truncate=True
+        )
+        zero_pad = torch.zeros(
+            [texts.shape[0], default_context_length - context_length],
+            dtype=texts.dtype,
+            device=texts.device,
+        )
+        texts = torch.cat([texts, zero_pad], dim=1)
+    else:
+        # [bs, context_length] # if n_tokens > 77 -> will truncate
+        texts = clip.tokenize(caption_raws, truncate=True)
+
+    # [batch_size, n_ctx, d_model]
+    x = clip_model.token_embedding(texts.to(device)).type(clip_model.dtype)
+    x = x + clip_model.positional_embedding.type(clip_model.dtype)
+    x = x.permute(1, 0, 2)  # NLD -> LND
+    x = clip_model.transformer(x)
+    x = x.permute(1, 0, 2)  # LND -> NLD
+    x = clip_model.ln_final(x).type(clip_model.dtype)
+    # x.shape = [batch_size, n_ctx, transformer.width]
+    # take features from the eot embedding (eot_token is the highest in each sequence)
+    x_tokens = x[torch.arange(x.shape[0]), texts.argmax(dim=-1)].float()
+    x_seq = [x[k, : (m + 1)].float() for k, m in enumerate(texts.argmax(dim=-1))]
+
+    return x_seq, x_tokens
+
+
+def get_batch(
+    prompt: str,
+    sample_id: str,
+    clip_model: clip.model.CLIP,
+    dataset: MultimodalDataset,
+    seq_feat: bool,
+    device: torch.device,
+) -> Dict[str, Any]:
+    # Get base batch
+    sample_index = dataset.root_filenames.index(sample_id)
+    raw_batch = dataset[sample_index]
+    batch = collate_fn([to_device(raw_batch, device)])
+
+    # Encode text
+    caption_seq, caption_tokens = encode_text([prompt], clip_model, None, device)
+
+    if seq_feat:
+        caption_feat = caption_seq[0]
+        caption_feat = F.pad(caption_feat, (0, 0, 0, 77 - caption_feat.shape[0]))
+        caption_feat = caption_feat.unsqueeze(0).permute(0, 2, 1)
+    else:
+        caption_feat = caption_tokens
+
+    # Update batch
+    batch["caption_raw"] = [prompt]
+    batch["caption_feat"] = caption_feat
+
+    return batch
+
+
+def init(
+    config_name: str,
+) -> Tuple[Diffuser, clip.model.CLIP, MultimodalDataset, torch.device]:
+    with initialize(version_base="1.3", config_path="../configs"):
+        config = compose(config_name=config_name)
+
+    OmegaConf.register_new_resolver("eval", eval)
+
+    # Initialize model
+    device = torch.device(config.compnode.device)
+    diffuser = instantiate(config.diffuser)
+    state_dict = torch.load(config.checkpoint_path, map_location=device)["state_dict"]
+    state_dict["ema.initted"] = diffuser.ema.initted
+    state_dict["ema.step"] = diffuser.ema.step
+    diffuser.load_state_dict(state_dict, strict=False)
+    diffuser.to(device).eval()
+
+    # Initialize CLIP model
+    clip_model = load_clip_model("ViT-B/32", device)
+
+    # Initialize dataset
+    config.dataset.char.load_vertices = True
+    config.batch_size = 1
+    dataset = instantiate(config.dataset)
+    dataset.set_split("demo")
+    diffuser.modalities = list(dataset.modality_datasets.keys())
+    diffuser.get_matrix = dataset.get_matrix
+    diffuser.v_get_matrix = dataset.get_matrix
+
+    return diffuser, clip_model, dataset, device

utils/file_utils.py

@@ -0,0 +1,117 @@
+import json
+import os
+import os.path as osp
+import pickle
+import subprocess
+from typing import Any
+
+import h5py
+import numpy as np
+import pandas as pd
+import torch
+import torchaudio
+from torchtyping import TensorType
+
+num_channels, num_frames, height, width = None, None, None, None
+
+
+def create_dir(dir_name: str):
+    """Create a directory if it does not exist yet."""
+    if not osp.exists(dir_name):
+        os.makedirs(dir_name)
+
+
+def move_files(source_path: str, destpath: str):
+    """Move files from `source_path` to `dest_path`."""
+    subprocess.call(["mv", source_path, destpath])
+
+
+def load_pickle(pickle_path: str) -> Any:
+    """Load a pickle file."""
+    with open(pickle_path, "rb") as f:
+        data = pickle.load(f)
+    return data
+
+
+def load_hdf5(hdf5_path: str) -> Any:
+    with h5py.File(hdf5_path, "r") as h5file:
+        data = {key: np.array(value) for key, value in h5file.items()}
+    return data
+
+
+def save_hdf5(data: Any, hdf5_path: str):
+    with h5py.File(hdf5_path, "w") as h5file:
+        for key, value in data.items():
+            h5file.create_dataset(key, data=value)
+
+
+def save_pickle(data: Any, pickle_path: str):
+    """Save data in a pickle file."""
+    with open(pickle_path, "wb") as f:
+        pickle.dump(data, f, protocol=4)
+
+
+def load_txt(txt_path: str):
+    """Load a txt file."""
+    with open(txt_path, "r") as f:
+        data = f.read()
+    return data
+
+
+def save_txt(data: str, txt_path: str):
+    """Save data in a txt file."""
+    with open(txt_path, "w") as f:
+        f.write(data)
+
+
+def load_pth(pth_path: str) -> Any:
+    """Load a pth (PyTorch) file."""
+    data = torch.load(pth_path)
+    return data
+
+
+def save_pth(data: Any, pth_path: str):
+    """Save data in a pth (PyTorch) file."""
+    torch.save(data, pth_path)
+
+
+def load_csv(csv_path: str, header: Any = None) -> pd.DataFrame:
+    """Load a csv file."""
+    try:
+        data = pd.read_csv(csv_path, header=header)
+    except pd.errors.EmptyDataError:
+        data = pd.DataFrame()
+    return data
+
+
+def save_csv(data: Any, csv_path: str):
+    """Save data in a csv file."""
+    pd.DataFrame(data).to_csv(csv_path, header=False, index=False)
+
+
+def load_json(json_path: str, header: Any = None) -> pd.DataFrame:
+    """Load a json file."""
+    with open(json_path, "r") as f:
+        data = json.load(f)
+    return data
+
+
+def save_json(data: Any, json_path: str):
+    """Save data in a json file."""
+    with open(json_path, "w") as json_file:
+        json.dump(data, json_file)
+
+
+def load_audio(audio_path: str, **kwargs):
+    """Load an audio file."""
+    waveform, sample_rate = torchaudio.load(audio_path, **kwargs)
+    return waveform, sample_rate
+
+
+def save_audio(
+    data: TensorType["num_channels", "num_frames"],
+    audio_path: str,
+    sample_rate: int = 44100,
+):
+    """Save data in an audio file."""
+    torchaudio.save(audio_path, data, sample_rate)

utils/random_utils.py

@@ -0,0 +1,46 @@
+import numpy as np
+import random
+import torch
+
+
+def set_random_seed(seed: int):
+    torch.manual_seed((seed) % (1 << 31))
+    torch.cuda.manual_seed((seed) % (1 << 31))
+    torch.cuda.manual_seed_all((seed) % (1 << 31))
+    np.random.seed((seed) % (1 << 31))
+    random.seed((seed) % (1 << 31))
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+class StackedRandomGenerator:
+    """
+    Wrapper for torch.Generator that allows specifying a different random seed for each
+    sample in a minibatch.
+    """
+
+    def __init__(self, device, seeds):
+        super().__init__()
+        self.generators = [
+            torch.Generator(device).manual_seed(int(seed) % (1 << 31)) for seed in seeds
+        ]
+
+    def randn_rn(self, size, **kwargs):
+        assert size[0] == len(self.generators)
+        return torch.stack(
+            [torch.randn(size[1:], generator=gen, **kwargs) for gen in self.generators]
+        )
+
+    def randn_like(self, input):
+        return self.randn_rn(
+            input.shape, dtype=input.dtype, layout=input.layout, device=input.device
+        )
+
+    def randint(self, *args, size, **kwargs):
+        assert size[0] == len(self.generators)
+        return torch.stack(
+            [
+                torch.randint(*args, size=size[1:], generator=gen, **kwargs)
+                for gen in self.generators
+            ]
+        )

utils/rerun.py

@@ -0,0 +1,74 @@
+import numpy as np
+from matplotlib import colormaps
+import rerun as rr
+from rerun.components import Material
+from scipy.spatial import transform
+
+
+def color_fn(x, cmap="tab10"):
+    return colormaps[cmap](x % colormaps[cmap].N)
+
+
+def log_sample(
+    root_name: str,
+    traj: np.ndarray,
+    K: np.ndarray,
+    vertices: np.ndarray,
+    faces: np.ndarray,
+    normals: np.ndarray,
+    caption: str,
+    mesh_masks: np.ndarray,
+):
+    num_cameras = traj.shape[0]
+
+    rr.log(root_name, rr.ViewCoordinates.RIGHT_HAND_Y_DOWN, timeless=True)
+    rr.log(
+        f"{root_name}/trajectory/points",
+        rr.Points3D(traj[:, :3, 3]),
+        timeless=True,
+    )
+    rr.log(
+        f"{root_name}/trajectory/line",
+        rr.LineStrips3D(
+            np.stack((traj[:, :3, 3][:-1], traj[:, :3, 3][1:]), axis=1),
+            colors=[(1.0, 0.0, 1.0, 1.0)],
+        ),
+        timeless=True,
+    )
+    for k in range(num_cameras):
+        rr.set_time_sequence("frame_idx", k)
+
+        translation = traj[k][:3, 3]
+        rotation_q = transform.Rotation.from_matrix(traj[k][:3, :3]).as_quat()
+        rr.log(
+            f"{root_name}/camera/image",
+            rr.Pinhole(
+                image_from_camera=K,
+                width=K[0, -1] * 2,
+                height=K[1, -1] * 2,
+            ),
+        )
+        rr.log(
+            f"{root_name}/camera",
+            rr.Transform3D(
+                translation=translation,
+                rotation=rr.Quaternion(xyzw=rotation_q),
+            ),
+        )
+        rr.set_time_sequence("image", k)
+
+        # Null vertices
+        if vertices[k].sum() == 0:
+            rr.log(f"{root_name}/char/char", rr.Clear(recursive=False))
+            rr.log(f"{root_name}/camera/image/bbox", rr.Clear(recursive=False))
+            continue
+
+        rr.log(
+            f"{root_name}/char/char",
+            rr.Mesh3D(
+                vertex_positions=vertices[k],
+                indices=faces,
+                vertex_normals=normals[k],
+                mesh_material=Material(albedo_factor=color_fn(0)),
+            ),
+        )

utils/rotation_utils.py

@@ -0,0 +1,178 @@
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+import torch
+from torchtyping import TensorType
+from itertools import product
+
+num_samples, num_cams = None, None
+
+
+def rotvec_to_matrix(rotvec):
+    return R.from_rotvec(rotvec).as_matrix()
+
+
+def matrix_to_rotvec(mat):
+    return R.from_matrix(mat).as_rotvec()
+
+
+def compose_rotvec(r1, r2):
+    """
+    #TODO: adapt to torch
+    Compose two rotation euler vectors.
+    """
+    r1 = r1.cpu().numpy() if isinstance(r1, torch.Tensor) else r1
+    r2 = r2.cpu().numpy() if isinstance(r2, torch.Tensor) else r2
+
+    R1 = rotvec_to_matrix(r1)
+    R2 = rotvec_to_matrix(r2)
+    cR = np.einsum("...ij,...jk->...ik", R1, R2)
+    return torch.from_numpy(matrix_to_rotvec(cR))
+
+
+def quat_to_rotvec(quat, eps=1e-6):
+    # w > 0 to ensure 0 <= angle <= pi
+    flip = (quat[..., :1] < 0).float()
+    quat = (-1 * quat) * flip + (1 - flip) * quat
+
+    angle = 2 * torch.atan2(torch.linalg.norm(quat[..., 1:], dim=-1), quat[..., 0])
+
+    angle2 = angle * angle
+    small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
+    large_angle_scales = angle / torch.sin(angle / 2 + eps)
+
+    small_angles = (angle <= 1e-3).float()
+    rot_vec_scale = (
+        small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
+    )
+    rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
+    return rot_vec
+
+
+# batch*n
+def normalize_vector(v, return_mag=False):
+    batch = v.shape[0]
+    v_mag = torch.sqrt(v.pow(2).sum(1))  # batch
+    v_mag = torch.max(
+        v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8])).to(v.device)
+    )
+    v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
+    v = v / v_mag
+    if return_mag is True:
+        return v, v_mag[:, 0]
+    else:
+        return v
+
+
+# u, v batch*n
+def cross_product(u, v):
+    batch = u.shape[0]
+    i = u[:, 1] * v[:, 2] - u[:, 2] * v[:, 1]
+    j = u[:, 2] * v[:, 0] - u[:, 0] * v[:, 2]
+    k = u[:, 0] * v[:, 1] - u[:, 1] * v[:, 0]
+
+    out = torch.cat(
+        (i.view(batch, 1), j.view(batch, 1), k.view(batch, 1)), 1
+    )  # [batch, 6]
+
+    return out
+
+
+def compute_rotation_matrix_from_ortho6d(ortho6d):
+    x_raw = ortho6d[:, 0:3]  # [batch, 6]
+    y_raw = ortho6d[:, 3:6]  # [batch, 6]
+
+    x = normalize_vector(x_raw)  # [batch, 6]
+    z = cross_product(x, y_raw)  # [batch, 6]
+    z = normalize_vector(z)  # [batch, 6]
+    y = cross_product(z, x)  # [batch, 6]
+
+    x = x.view(-1, 3, 1)
+    y = y.view(-1, 3, 1)
+    z = z.view(-1, 3, 1)
+    matrix = torch.cat((x, y, z), 2)  # [batch, 3, 3]
+    return matrix
+
+
+def invert_rotvec(rotvec: TensorType["num_samples", 3]):
+    angle = torch.norm(rotvec, dim=-1)
+    axis = rotvec / (angle.unsqueeze(-1) + 1e-6)
+    inverted_rotvec = -angle.unsqueeze(-1) * axis
+    return inverted_rotvec
+
+
+def are_rotations(matrix: TensorType["num_samples", 3, 3]) -> TensorType["num_samples"]:
+    """Check if a matrix is a rotation matrix."""
+    # Check if the matrix is orthogonal
+    identity = torch.eye(3, device=matrix.device)
+    is_orthogonal = (
+        torch.isclose(torch.bmm(matrix, matrix.transpose(1, 2)), identity, atol=1e-6)
+        .all(dim=1)
+        .all(dim=1)
+    )
+
+    # Check if the determinant is 1
+    determinant = torch.det(matrix)
+    is_determinant_one = torch.isclose(
+        determinant, torch.tensor(1.0, device=matrix.device), atol=1e-6
+    )
+
+    return torch.logical_and(is_orthogonal, is_determinant_one)
+
+
+def project_so3(
+    matrix: TensorType["num_samples", 4, 4]
+) -> TensorType["num_samples", 4, 4]:
+    # Project rotation matrix to SO(3)
+    # TODO: use torch
+    rot = R.from_matrix(matrix[:, :3, :3].cpu().numpy()).as_matrix()
+
+    projection = torch.eye(4).unsqueeze(0).repeat(matrix.shape[0], 1, 1).to(matrix)
+    projection[:, :3, :3] = torch.from_numpy(rot).to(matrix)
+    projection[:, :3, 3] = matrix[:, :3, 3]
+
+    return projection
+
+
+def pairwise_geodesic(
+    R_x: TensorType["num_samples", "num_cams", 3, 3],
+    R_y: TensorType["num_samples", "num_cams", 3, 3],
+    reduction: str = "mean",
+    block_size: int = 200,
+):
+    def arange(start, stop, step, endpoint=True):
+        arr = torch.arange(start, stop, step)
+        if endpoint and arr[-1] != stop - 1:
+            arr = torch.cat((arr, torch.tensor([stop - 1], dtype=arr.dtype)))
+        return arr
+
+    # Geodesic distance
+    # https://math.stackexchange.com/questions/2113634/comparing-two-rotation-matrices
+    num_samples, num_cams, _, _ = R_x.shape
+
+    C = torch.zeros(num_samples, num_samples, device=R_x.device)
+    chunk_indices = arange(0, num_samples + 1, block_size, endpoint=True)
+    for i, j in product(
+        range(chunk_indices.shape[0] - 1), range(chunk_indices.shape[0] - 1)
+    ):
+        start_x, stop_x = chunk_indices[i], chunk_indices[i + 1]
+        start_y, stop_y = chunk_indices[j], chunk_indices[j + 1]
+        r_x, r_y = R_x[start_x:stop_x], R_y[start_y:stop_y]
+
+        # Compute rotations between each pair of cameras of each sample
+        r_xy = torch.einsum("anjk,bnlk->abnjl", r_x, r_y)  # b, b, N, 3, 3
+
+        # Compute axis-angle representations: angle is the geodesic distance
+        traces = r_xy.diagonal(dim1=-2, dim2=-1).sum(-1)
+        c = torch.acos(torch.clamp((traces - 1) / 2, -1, 1)) / torch.pi
+
+        # Average distance between cameras over samples
+        if reduction == "mean":
+            C[start_x:stop_x, start_y:stop_y] = c.mean(-1)
+        elif reduction == "sum":
+            C[start_x:stop_x, start_y:stop_y] = c.sum(-1)
+
+        # Check for NaN values in traces
+        if torch.isnan(c).any():
+            raise ValueError("NaN values detected in traces")
+
+    return C