upload

README.md

@@ -5,7 +5,7 @@ colorFrom: purple
 colorTo: blue
 sdk: gradio
 sdk_version: 5.20.1
-app_file: app.py
+app_file: gen_demo.py
 pinned: false
 license: mit
 ---

configs/autoregressive_l.yaml

@@ -0,0 +1,68 @@
+trainer:
+  target: semanticist.engine.gpt_trainer.GPTTrainer
+  params:
+    num_epoch: 400
+    blr: 1e-4
+    cosine_lr: False
+    warmup_epochs: 100
+    batch_size: 16
+    num_workers: 8
+    pin_memory: True
+    grad_accum_steps: 1
+    precision: 'bf16'
+    max_grad_norm: 1.0
+    enable_ema: True
+    save_every: 10000
+    sample_every: 5000
+    fid_every: 50000
+    eval_fid: False
+    result_folder: "./output/autoregressive"
+    log_dir: "./output/autoregressive/logs"
+    ae_cfg: 1.0
+    cfg: 6.0
+    cfg_schedule: "linear"
+    train_num_slots: 32
+    test_num_slots: 32
+    compile: True
+    enable_cache_latents: False
+    ae_model:
+      target: semanticist.stage1.diffuse_slot.DiffuseSlot
+      params:
+        encoder: 'vit_base_patch16'
+        enc_img_size: 256
+        enc_causal: True
+        num_slots: 256
+        slot_dim: 16
+        norm_slots: True
+        cond_method: 'token'
+        dit_model: 'DiT-L-2'
+        vae: 'xwen99/mar-vae-kl16'
+        num_sampling_steps: '250'
+        # ckpt_path: ./output/tokenizer/models_l/step250000/custom_checkpoint_1.pkl
+        ckpt_path: /mnt/ceph_rbd/mnt_pvc_vid_data/zbc/cache/semanticist_tok_L.pkl
+
+    gpt_model:
+      target: GPT-L
+      params:
+        num_slots: 32
+        slot_dim: 16
+        num_classes: 1000
+        cls_token_num: 1
+        resid_dropout_p: 0.1
+        ffn_dropout_p: 0.1
+        diffloss_d: 12
+        diffloss_w: 1536
+        num_sampling_steps: '100'
+        diffusion_batch_mul: 4
+        use_si: True
+        cond_method: 'concat'
+        ckpt_path: None
+
+    dataset:
+      target: semanticist.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: train
+        # aug:  tencrop_cached # or centercrop_cached
+        aug: randcrop
+        img_size: 256

configs/autoregressive_xl.yaml

@@ -0,0 +1,66 @@
+trainer:
+  target: semanticist.engine.gpt_trainer.GPTTrainer
+  params:
+    num_epoch: 400
+    blr: 1e-4
+    cosine_lr: False
+    warmup_epochs: 100
+    batch_size: 256
+    num_workers: 8
+    pin_memory: True
+    grad_accum_steps: 1
+    precision: 'bf16'
+    max_grad_norm: 1.0
+    enable_ema: True
+    save_every: 10000
+    sample_every: 5000
+    fid_every: 50000
+    eval_fid: False
+    result_folder: "./output/autoregressive"
+    log_dir: "./output/autoregressive/logs"
+    ae_cfg: 1.0
+    cfg: 5.0
+    cfg_schedule: "linear"
+    train_num_slots: 32
+    test_num_slots: 32
+    compile: True
+    enable_cache_latents: True
+    ae_model:
+      target: semanticist.stage1.diffuse_slot.DiffuseSlot
+      params:
+        encoder: 'vit_base_patch16'
+        enc_img_size: 256
+        enc_causal: True
+        num_slots: 256
+        slot_dim: 16
+        norm_slots: True
+        cond_method: 'token'
+        dit_model: 'DiT-XL-2'
+        vae: 'xwen99/mar-vae-kl16'
+        num_sampling_steps: '250'
+        ckpt_path: ./output/tokenizer/models_xl/step250000/custom_checkpoint_1.pkl
+
+    gpt_model:
+      target: GPT-L
+      params:
+        num_slots: 32
+        slot_dim: 16
+        num_classes: 1000
+        cls_token_num: 1
+        resid_dropout_p: 0.1
+        ffn_dropout_p: 0.1
+        diffloss_d: 12
+        diffloss_w: 1536
+        num_sampling_steps: '100'
+        diffusion_batch_mul: 4
+        use_si: True
+        cond_method: 'concat'
+        ckpt_path: None
+
+    dataset:
+      target: semanticist.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: train
+        aug:  tencrop_cached # or centercrop_cached
+        img_size: 256

configs/onenode_config.yaml

@@ -0,0 +1,11 @@
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: MULTI_GPU
+fsdp_config: {}
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+num_machines: 1
+num_processes: 1
+use_cpu: false

configs/tokenizer_l.yaml

@@ -0,0 +1,55 @@
+trainer:
+  target: semanticist.engine.diffusion_trainer.DiffusionTrainer
+  params:
+    num_epoch: 400
+    valid_size: 64
+    blr: 2.5e-5
+    cosine_lr: True
+    warmup_epochs: 1
+    batch_size: 64
+    num_workers: 16
+    pin_memory: True
+    grad_accum_steps: 1
+    precision: 'bf16'
+    max_grad_norm: 3.0
+    enable_ema: True
+    save_every: 10000
+    sample_every: 5000
+    fid_every: 50000
+    result_folder: "./output/tokenizer/models_l"
+    log_dir: "./output/tokenizer/models_l/logs"
+    cfg: 3.0
+    compile: True
+    model:
+      target: semanticist.stage1.diffuse_slot.DiffuseSlot
+      params:
+        encoder: 'vit_base_patch16'
+        enc_img_size: 256
+        enc_causal: True
+        enc_use_mlp: False
+        num_slots: 256
+        slot_dim: 16
+        norm_slots: True
+        dit_model: 'DiT-L-2'
+        vae: 'xwen99/mar-vae-kl16'
+        enable_nest: False
+        enable_nest_after: 50
+        use_repa: True
+        eval_fid: True
+        fid_stats: 'fid_stats/adm_in256_stats.npz'
+        num_sampling_steps: '250'
+        ckpt_path: None
+
+    dataset:
+      target: semanticist.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: train
+        img_size: 256
+
+    test_dataset:
+      target: semanticist.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: val
+        img_size: 256

configs/tokenizer_xl.yaml

@@ -0,0 +1,55 @@
+trainer:
+  target: semanticist.engine.diffusion_trainer.DiffusionTrainer
+  params:
+    num_epoch: 400
+    valid_size: 64
+    blr: 2.5e-5
+    cosine_lr: True
+    warmup_epochs: 100
+    batch_size: 256
+    num_workers: 16
+    pin_memory: True
+    grad_accum_steps: 1
+    precision: 'bf16'
+    max_grad_norm: 3.0
+    enable_ema: True
+    save_every: 10000
+    sample_every: 5000
+    fid_every: 50000
+    result_folder: "./output/tokenizer/models_xl"
+    log_dir: "./output/tokenizer/models_xl/logs"
+    cfg: 3.0
+    compile: True
+    model:
+      target: semanticist.stage1.diffuse_slot.DiffuseSlot
+      params:
+        encoder: 'vit_base_patch16'
+        enc_img_size: 256
+        enc_causal: True
+        enc_use_mlp: False
+        num_slots: 256
+        slot_dim: 16
+        norm_slots: True
+        dit_model: 'DiT-XL-2'
+        vae: 'xwen99/mar-vae-kl16'
+        enable_nest: False
+        enable_nest_after: 50
+        use_repa: True
+        eval_fid: True
+        fid_stats: 'fid_stats/adm_in256_stats.npz'
+        num_sampling_steps: '250'
+        ckpt_path: None
+
+    dataset:
+      target: semanticist.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: train
+        img_size: 256
+
+    test_dataset:
+      target: semanticist.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: val
+        img_size: 256

gen_demo.py

@@ -0,0 +1,262 @@
+import gradio as gr
+import numpy as np
+from PIL import Image
+import os.path as osp
+import torch
+import matplotlib.pyplot as plt
+from omegaconf import OmegaConf
+from tqdm import tqdm
+from huggingface_hub import hf_hub_download
+from semanticist.engine.trainer_utils import instantiate_from_config
+from semanticist.stage1.diffuse_slot import DiffuseSlot
+from semanticist.stage2.gpt import GPT_models
+from semanticist.stage2.generate import generate
+from safetensors import safe_open
+from semanticist.utils.datasets import vae_transforms
+from PIL import Image
+from imagenet_classes import imagenet_classes
+
+transform = vae_transforms('test')
+
+
+def norm_ip(img, low, high):
+    img.clamp_(min=low, max=high)
+    img.sub_(low).div_(max(high - low, 1e-5))
+
+def norm_range(t, value_range):
+    if value_range is not None:
+        norm_ip(t, value_range[0], value_range[1])
+    else:
+        norm_ip(t, float(t.min()), float(t.max()))
+
+from PIL import Image
+def convert_np(img):
+    ndarr = img.mul(255).add_(0.5).clamp_(0, 255)\
+            .permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    return ndarr
+def convert_PIL(img):
+    ndarr = img.mul(255).add_(0.5).clamp_(0, 255)\
+            .permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+    img = Image.fromarray(ndarr)
+    return img
+
+def norm_slots(slots):
+    mean = torch.mean(slots, dim=-1, keepdim=True)
+    std = torch.std(slots, dim=-1, keepdim=True)
+    return (slots - mean) / std
+
+def load_state_dict(state_dict, model):
+    """Helper to load a state dict with proper prefix handling."""
+    if 'state_dict' in state_dict:
+        state_dict = state_dict['state_dict']
+    # Remove '_orig_mod' prefix if present
+    state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+    missing, unexpected = model.load_state_dict(
+        state_dict, strict=False
+    )
+    # print(f"Loaded model. Missing: {missing}, Unexpected: {unexpected}")
+
+def load_safetensors(path, model):
+    """Helper to load a safetensors checkpoint."""
+    from safetensors.torch import safe_open
+    with safe_open(path, framework="pt", device="cpu") as f:
+        state_dict = {k: f.get_tensor(k) for k in f.keys()}
+    load_state_dict(state_dict, model)
+
+def load_checkpoint(ckpt_path, model):
+    if ckpt_path is None or not osp.exists(ckpt_path):
+        return
+
+    if osp.isdir(ckpt_path):
+        # ckpt_path is something like 'path/to/models/step10/'
+        model_path = osp.join(ckpt_path, "model.safetensors")
+        if osp.exists(model_path):
+            load_safetensors(model_path, model)
+    else:
+        # ckpt_path is something like 'path/to/models/step10.pt'
+        if ckpt_path.endswith(".safetensors"):
+            load_safetensors(ckpt_path, model)
+        else:
+            state_dict = torch.load(ckpt_path, map_location="cpu")
+            load_state_dict(state_dict, model)
+
+    print(f"Loaded checkpoint from {ckpt_path}")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+if device == 'cuda':
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
+ckpt_path = hf_hub_download(repo_id='tennant/semanticist', filename="semanticist_ar_gen_L.pkl", cache_dir='/mnt/ceph_rbd/mnt_pvc_vid_data/zbc/cache/')
+config_path = 'configs/autoregressive_xl.yaml'
+
+cfg = OmegaConf.load(config_path)
+params = cfg.trainer.params
+
+ae_model = instantiate_from_config(params.ae_model).to(device)
+ae_model_path = hf_hub_download(repo_id='tennant/semanticist', filename="semanticist_tok_XL.pkl", cache_dir='/mnt/ceph_rbd/mnt_pvc_vid_data/zbc/cache/')
+load_checkpoint(ae_model_path, ae_model)
+ae_model.eval()
+
+gpt_model = GPT_models[params.gpt_model.target](**params.gpt_model.params).to(device)
+load_checkpoint(ckpt_path, gpt_model)
+gpt_model.eval();
+
+def viz_diff_slots(model, slots, nums, cfg=1.0, return_figs=False):
+    n_slots_inf = []
+    for num_slots_to_inference in nums:
+        drop_mask = model.nested_sampler(slots.shape[0], device, num_slots_to_inference)
+        recon_n = model.sample(slots, drop_mask=drop_mask, cfg=cfg)
+        n_slots_inf.append(recon_n)
+    return [convert_np(n_slots_inf[i][0]) for i in range(len(n_slots_inf))]
+
+num_slots = params.ae_model.params.num_slots
+slot_dim = params.ae_model.params.slot_dim
+dtype = torch.bfloat16
+# the model is trained with only 32 tokens.
+num_slots_to_gen = 32
+
+# Function to generate image from class
+def generate_from_class(class_id, cfg_scale):
+    with torch.no_grad():
+        dtype = torch.bfloat16
+        num_slots_to_gen = 32
+        with torch.autocast(device, dtype=dtype):
+            slots_gen = generate(
+                gpt_model, 
+                torch.tensor([class_id]).to(device), 
+                num_slots_to_gen, 
+                cfg_scale=cfg_scale, 
+                cfg_schedule="linear"
+            )
+        if num_slots_to_gen < num_slots:
+            null_slots = ae_model.dit.null_cond.expand(slots_gen.shape[0], -1, -1)
+            null_slots = null_slots[:, num_slots_to_gen:, :]
+            slots_gen = torch.cat([slots_gen, null_slots], dim=1)
+    return slots_gen
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        # First column - Input and configs
+        with gr.Column(scale=1):
+            gr.Markdown("## Input")
+            
+            # Replace image input with ImageNet class selection
+            imagenet_classes = {k: v for k, v in enumerate(imagenet_classes)}
+            class_choices = [f"{id}: {name}" for id, name in imagenet_classes.items()]
+            
+            # Dropdown for class selection
+            class_dropdown = gr.Dropdown(
+                choices=class_choices[:20],  # Limit for demonstration
+                label="Select ImageNet Class",
+                value=class_choices[0] if class_choices else None
+            )
+            
+            # Option to enter class ID directly
+            class_id_input = gr.Number(
+                label="Or enter class ID directly (0-999)",
+                value=0,
+                minimum=0,
+                maximum=999,
+                step=1
+            )
+            
+            with gr.Group():
+                gr.Markdown("### Configuration")
+                show_gallery = gr.Checkbox(label="Show Gallery", value=True)
+                slider = gr.Slider(minimum=0.1, maximum=20.0, value=4.0, label="CFG value")
+                labels_input = gr.Textbox(
+                    label="Number of tokens to reconstruct (comma-separated)", 
+                    value="1, 2, 4, 8, 16", 
+                    placeholder="Enter comma-separated numbers for the number of slots to use"
+                )
+        
+        # Second column - Output (conditionally rendered)
+        with gr.Column(scale=1):
+            gr.Markdown("## Output")
+            
+            # Container for conditional rendering
+            with gr.Group(visible=True) as gallery_container:
+                gallery = gr.Gallery(label="Result Gallery", columns=3, height="auto", show_label=True)
+            
+            # Always visible output image
+            output_image = gr.Image(label="Generated Image", type="numpy")
+    
+    # Handle form submission
+    submit_btn = gr.Button("Generate")
+    
+    # Define the processing logic
+    def update_outputs(class_selection, class_id, show_gallery_value, slider_value, labels_text):
+        # Determine which class to use - either from dropdown or direct input
+        if class_selection:
+            # Extract class ID from the dropdown selection
+            selected_class_id = int(class_selection.split(":")[0])
+        else:
+            selected_class_id = int(class_id)
+        
+        # Update the visibility of the gallery container
+        gallery_container.visible = show_gallery_value
+        
+        try:
+            # Parse the labels from the text input
+            if labels_text and "," in labels_text:
+                labels = [int(label.strip()) for label in labels_text.split(",")]
+            else:
+                # Default labels if none provided or in wrong format
+                labels = [1, 4, 16, 64, 256]
+        except:
+            labels = [1, 4, 16, 64, 256]
+        
+        while len(labels) < 3:
+            labels.append(256)
+        
+        # Generate the image based on the selected class
+        slots_gen = generate_from_class(selected_class_id, cfg_scale=slider_value)
+        
+        recon = viz_diff_slots(ae_model, slots_gen, [32], cfg=slider_value)[0]
+        
+        # Always generate the model decomposition for potential gallery display
+        model_decompose = viz_diff_slots(ae_model, slots_gen, labels, cfg=slider_value)
+        
+        if not show_gallery_value:
+            # If only the image should be shown, return just the processed image
+            return gallery_container, [], recon
+        else:
+            # Create image variations and pair them with labels
+            gallery_images = [
+                (recon, f'Generated from class {selected_class_id}'),
+            ] + [(img, 'Gen. with ' + str(label) + ' tokens') for img, label in zip(model_decompose, labels)]
+            return gallery_container, gallery_images, recon
+    
+    # Connect the inputs and outputs
+    submit_btn.click(
+        fn=update_outputs,
+        inputs=[class_dropdown, class_id_input, show_gallery, slider, labels_input],
+        outputs=[gallery_container, gallery, output_image]
+    )
+    
+    # Also update when checkbox changes
+    show_gallery.change(
+        fn=lambda value: gr.update(visible=value),
+        inputs=[show_gallery],
+        outputs=[gallery_container]
+    )
+
+    # Add examples
+    examples = [
+        # ["0: tench, Tinca tinca", 0, True, 4.0, "1,2,4,8,16"],
+        ["1: goldfish, Carassius auratus", 1, True, 4.0, "1,2,4,8,16"],
+        # ["2: great white shark, white shark", 2, True, 4.0, "1,2,4,8,16"],
+    ]
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[class_dropdown, class_id_input, show_gallery, slider, labels_input],
+        outputs=[gallery_container, gallery, output_image],
+        fn=update_outputs,
+        cache_examples=False
+    )
+
+# Launch the demo
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,16 @@
+numpy==1.26.4
+accelerate
+diffusers[torch]
+transformers
+safetensors
+omegaconf
+tensorboard
+huggingface-hub
+einops
+timm
+scipy
+scikit-learn
+scikit-image
+git+https://github.com/LTH14/torch-fidelity.git@master#egg=torch-fidelity
+opencv-python-headless
+torchmetrics

semanticist/engine/diffusion_trainer.py

@@ -0,0 +1,488 @@
+import os, torch
+import os.path as osp
+import shutil
+from tqdm.auto import tqdm
+from einops import rearrange
+from accelerate import Accelerator
+from torchvision.utils import make_grid, save_image
+from torch.utils.data import DataLoader, random_split, DistributedSampler
+from semanticist.utils.logger import SmoothedValue, MetricLogger, empty_cache
+from accelerate.utils import DistributedDataParallelKwargs
+from torchmetrics.functional.image import (
+    peak_signal_noise_ratio as psnr,
+    structural_similarity_index_measure as ssim
+)
+from semanticist.engine.trainer_utils import (
+    instantiate_from_config, concat_all_gather,
+    save_img_batch, get_fid_stats,
+    EMAModel, PaddedDataset, create_scheduler, load_state_dict,
+    load_safetensors, setup_result_folders, create_optimizer
+)
+
+class DiffusionTrainer:
+    def __init__(
+        self,
+        model,
+        dataset,
+        test_dataset=None,
+        test_only=False,
+        num_epoch=400,
+        valid_size=32,
+        blr=1e-4,
+        cosine_lr=True,
+        lr_min=0,
+        warmup_epochs=100,
+        warmup_steps=None,
+        warmup_lr_init=0,
+        decay_steps=None,
+        batch_size=32,
+        eval_bs=32,
+        test_bs=64,
+        num_workers=8,
+        pin_memory=False,
+        max_grad_norm=None,
+        grad_accum_steps=1,
+        precision='bf16',
+        save_every=10000,
+        sample_every=1000,
+        fid_every=50000,
+        result_folder=None,
+        log_dir="./log",
+        cfg=3.0,
+        test_num_slots=None,
+        eval_fid=False,
+        fid_stats=None,
+        enable_ema=False,
+        compile=False,
+    ):
+        kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+        self.accelerator = Accelerator(
+            kwargs_handlers=[kwargs],
+            mixed_precision=precision,
+            gradient_accumulation_steps=grad_accum_steps,
+            log_with="tensorboard",
+            project_dir=log_dir,
+        )
+        
+        self.model = instantiate_from_config(model)
+        self.num_slots = model.params.num_slots
+
+        if test_dataset is not None:
+            test_dataset = instantiate_from_config(test_dataset)
+            self.test_ds = test_dataset
+            
+            # Calculate padded dataset size to ensure even distribution
+            total_size = len(test_dataset)
+            world_size = self.accelerator.num_processes
+            padding_size = world_size * test_bs - (total_size % (world_size * test_bs))
+            self.test_dataset_size = total_size
+            
+            self.test_ds = PaddedDataset(self.test_ds, padding_size)
+            self.test_dl = DataLoader(
+                self.test_ds,
+                batch_size=test_bs,
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                shuffle=False,
+                drop_last=True,
+            )
+            if self.accelerator.is_main_process:
+                print(f"test dataset size: {len(test_dataset)}, test batch size: {test_bs}")
+        else:
+            self.test_dl = None
+        self.test_only = test_only
+
+        if not test_only:
+            dataset = instantiate_from_config(dataset)
+            train_size = len(dataset) - valid_size
+            self.train_ds, self.valid_ds = random_split(
+                dataset,
+                [train_size, valid_size],
+                generator=torch.Generator().manual_seed(42),
+            )
+            if self.accelerator.is_main_process:
+                print(f"train dataset size: {train_size}, valid dataset size: {valid_size}")
+
+            sampler = DistributedSampler(
+                self.train_ds,
+                num_replicas=self.accelerator.num_processes,
+                rank=self.accelerator.process_index,
+                shuffle=True,
+            )
+            self.train_dl = DataLoader(
+                self.train_ds,
+                batch_size=batch_size,
+                sampler=sampler,
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                drop_last=True,
+            )
+            self.valid_dl = DataLoader(
+                self.valid_ds,
+                batch_size=eval_bs,
+                shuffle=False,
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                drop_last=False,
+            )
+
+            effective_bs = batch_size * grad_accum_steps * self.accelerator.num_processes
+            lr = blr * effective_bs / 256
+            if self.accelerator.is_main_process:
+                print(f"Effective batch size is {effective_bs}")
+
+            self.g_optim = create_optimizer(self.model, weight_decay=0.05, learning_rate=lr,) # accelerator=self.accelerator)
+            
+            if warmup_epochs is not None:
+                warmup_steps = warmup_epochs * len(self.train_dl)
+            
+            self.g_sched = create_scheduler(
+                self.g_optim,
+                num_epoch,
+                len(self.train_dl),
+                lr_min,
+                warmup_steps,
+                warmup_lr_init,
+                decay_steps,
+                cosine_lr
+            )
+            self.accelerator.register_for_checkpointing(self.g_sched)
+            self.model, self.g_optim, self.g_sched = self.accelerator.prepare(self.model, self.g_optim, self.g_sched)
+        else:
+           self.model, self.test_dl = self.accelerator.prepare(self.model, self.test_dl)
+
+        self.steps = 0
+        self.loaded_steps = -1
+
+        if compile:
+            _model = self.accelerator.unwrap_model(self.model)
+            _model.vae = torch.compile(_model.vae, mode="reduce-overhead")
+            _model.dit = torch.compile(_model.dit, mode="reduce-overhead")
+            # _model.encoder = torch.compile(_model.encoder, mode="reduce-overhead") # nan loss when compiled together with dit, no idea why
+            _model.encoder2slot = torch.compile(_model.encoder2slot, mode="reduce-overhead")
+
+        self.enable_ema = enable_ema
+        if self.enable_ema and not self.test_only: # when testing, we directly load the ema dict and skip here
+            self.ema_model = EMAModel(self.accelerator.unwrap_model(self.model), self.device)
+            self.accelerator.register_for_checkpointing(self.ema_model)
+
+        self._load_checkpoint(model.params.ckpt_path)
+        if self.test_only:
+            self.steps = self.loaded_steps
+
+        self.num_epoch = num_epoch
+        self.save_every = save_every
+        self.sample_every = sample_every
+        self.fid_every = fid_every
+        self.max_grad_norm = max_grad_norm
+
+        self.cfg = cfg
+        self.test_num_slots = test_num_slots
+        if self.test_num_slots is not None:
+            self.test_num_slots = min(self.test_num_slots, self.num_slots)
+        else:
+            self.test_num_slots = self.num_slots
+        eval_fid = eval_fid or model.params.eval_fid # legacy
+        self.eval_fid = eval_fid
+        if eval_fid:
+            if fid_stats is None:
+                fid_stats = model.params.fid_stats # legacy
+            assert fid_stats is not None
+            assert test_dataset is not None
+        self.fid_stats = fid_stats
+
+        self.result_folder = result_folder
+        self.model_saved_dir, self.image_saved_dir = setup_result_folders(result_folder)
+
+    @property
+    def device(self):
+        return self.accelerator.device
+
+    def _load_checkpoint(self, ckpt_path=None):
+        if ckpt_path is None or not osp.exists(ckpt_path):
+            return
+        
+        model = self.accelerator.unwrap_model(self.model)
+
+        if osp.isdir(ckpt_path):
+            # ckpt_path is something like 'path/to/models/step10/'
+            self.loaded_steps = int(
+                ckpt_path.split("step")[-1].split("/")[0]
+            )
+            if not self.test_only:
+                self.accelerator.load_state(ckpt_path)
+            else:
+                if self.enable_ema:
+                    model_path = osp.join(ckpt_path, "custom_checkpoint_1.pkl")
+                    if osp.exists(model_path):
+                        state_dict = torch.load(model_path, map_location="cpu")
+                        load_state_dict(state_dict, model)
+                        if self.accelerator.is_main_process:
+                            print(f"Loaded ema model from {model_path}")
+                else:
+                    model_path = osp.join(ckpt_path, "model.safetensors")
+                    if osp.exists(model_path):
+                        load_safetensors(model_path, model)
+        else:
+            # ckpt_path is something like 'path/to/models/step10.pt'
+            if ckpt_path.endswith(".safetensors"):
+                load_safetensors(ckpt_path, model)
+            else:
+                state_dict = torch.load(ckpt_path, map_location="cpu")
+                load_state_dict(state_dict, model)
+        if self.accelerator.is_main_process:
+            print(f"Loaded checkpoint from {ckpt_path}")
+
+    def train(self, config=None):
+        n_parameters = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        if self.accelerator.is_main_process:
+            print(f"number of learnable parameters: {n_parameters//1e6}M")
+        if config is not None:
+            # save the config
+            from omegaconf import OmegaConf
+            if isinstance(config, str) and osp.exists(config):
+                # If it's a path, copy the file to config.yaml
+                shutil.copy(config, osp.join(self.result_folder, "config.yaml"))
+            else:
+                # If it's an OmegaConf object, dump it
+                config_save_path = osp.join(self.result_folder, "config.yaml")
+                OmegaConf.save(config, config_save_path)
+
+        self.accelerator.init_trackers("semanticist")
+
+        if self.test_only:
+            empty_cache()
+            self.evaluate()
+            self.accelerator.wait_for_everyone()
+            empty_cache()
+            return
+
+        for epoch in range(self.num_epoch):
+            if ((epoch + 1) * len(self.train_dl)) <= self.loaded_steps:
+                if self.accelerator.is_main_process:
+                    print(f"Epoch {epoch} is skipped because it is loaded from ckpt")
+                self.steps += len(self.train_dl)
+                continue
+
+            if self.steps < self.loaded_steps:
+                for _ in self.train_dl:
+                    self.steps += 1
+                    if self.steps >= self.loaded_steps:
+                        break
+            
+            
+            self.accelerator.unwrap_model(self.model).current_epoch = epoch
+            self.model.train()  # Set model to training mode
+
+            logger = MetricLogger(delimiter="  ")
+            logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))
+            header = 'Epoch: [{}/{}]'.format(epoch, self.num_epoch)
+            print_freq = 20
+            for data_iter_step, batch in enumerate(logger.log_every(self.train_dl, print_freq, header)):
+                img, _ = batch
+                img = img.to(self.device, non_blocking=True)
+                self.steps += 1
+
+                with self.accelerator.accumulate(self.model):
+                    with self.accelerator.autocast():
+                        if self.steps == 1:
+                            print(f"Training batch size: {img.size(0)}")
+                            print(f"Hello from index {self.accelerator.local_process_index}")
+                        losses = self.model(img, epoch=epoch)
+                        # combine
+                        loss = sum([v for _, v in losses.items()])
+
+                    self.accelerator.backward(loss)
+                    if self.accelerator.sync_gradients and self.max_grad_norm is not None:
+                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+                    self.g_optim.step()
+                    if self.g_sched is not None:
+                        self.g_sched.step_update(self.steps)
+                    self.g_optim.zero_grad()
+
+                self.accelerator.wait_for_everyone()
+
+                # update ema with state dict
+                if self.enable_ema:
+                    self.ema_model.update(self.accelerator.unwrap_model(self.model))
+
+                for key, value in losses.items():
+                    logger.update(**{key: value.item()})
+                logger.update(lr=self.g_optim.param_groups[0]["lr"])
+
+                if self.steps % self.save_every == 0:
+                    self.save()
+
+                if (self.steps % self.sample_every == 0) or (self.steps % self.fid_every == 0):
+                    empty_cache()
+                    self.evaluate()
+                    self.accelerator.wait_for_everyone()
+                    empty_cache()
+
+                write_dict = dict(epoch=epoch)
+                for key, value in losses.items(): # omitted all_gather here
+                    write_dict.update(**{key: value.item()})
+                write_dict.update(lr=self.g_optim.param_groups[0]["lr"])
+                self.accelerator.log(write_dict, step=self.steps)
+
+            logger.synchronize_between_processes()
+            if self.accelerator.is_main_process:
+                print("Averaged stats:", logger)
+
+        self.accelerator.end_training()
+        self.save()
+        if self.accelerator.is_main_process:
+            print("Train finished!")
+
+    def save(self):
+        self.accelerator.wait_for_everyone()
+        self.accelerator.save_state(
+            os.path.join(self.model_saved_dir, f"step{self.steps}")
+        )
+
+    @torch.no_grad()
+    def evaluate(self):
+        self.model.eval()
+        if not self.test_only:
+            with tqdm(
+                self.valid_dl,
+                dynamic_ncols=True,
+                disable=not self.accelerator.is_main_process,
+            ) as valid_dl:
+                for batch_i, batch in enumerate(valid_dl):
+                    if isinstance(batch, tuple) or isinstance(batch, list):
+                        img, targets = batch[0], batch[1]
+                    else:
+                        img = batch
+
+                    with self.accelerator.autocast():
+                        rec = self.model(img, sample=True, inference_with_n_slots=self.test_num_slots, cfg=1.0)
+                    imgs_and_recs = torch.stack((img.to(rec.device), rec), dim=0)
+                    imgs_and_recs = rearrange(imgs_and_recs, "r b ... -> (b r) ...")
+                    imgs_and_recs = imgs_and_recs.detach().cpu().float()
+
+                    grid = make_grid(
+                        imgs_and_recs, nrow=6, normalize=True, value_range=(0, 1)
+                    )
+                    if self.accelerator.is_main_process:
+                        save_image(
+                            grid,
+                            os.path.join(
+                                self.image_saved_dir, f"step_{self.steps}_slots{self.test_num_slots}_{batch_i}.jpg"
+                            ),
+                        )
+
+                    if self.cfg != 1.0:
+                        with self.accelerator.autocast():
+                            rec = self.model(img, sample=True, inference_with_n_slots=self.test_num_slots, cfg=self.cfg)
+
+                        imgs_and_recs = torch.stack((img.to(rec.device), rec), dim=0)
+                        imgs_and_recs = rearrange(imgs_and_recs, "r b ... -> (b r) ...")
+                        imgs_and_recs = imgs_and_recs.detach().cpu().float()
+
+                        grid = make_grid(
+                            imgs_and_recs, nrow=6, normalize=True, value_range=(0, 1)
+                        )
+                        if self.accelerator.is_main_process:
+                            save_image(
+                                grid,
+                                os.path.join(
+                                    self.image_saved_dir, f"step_{self.steps}_cfg_{self.cfg}_slots{self.test_num_slots}_{batch_i}.jpg"
+                                ),
+                            )
+        if (self.eval_fid and self.test_dl is not None) and (self.test_only or (self.steps % self.fid_every == 0)):
+            real_dir = "./dataset/imagenet/val256"
+            rec_dir = os.path.join(self.image_saved_dir, f"rec_step{self.steps}_slots{self.test_num_slots}")
+            os.makedirs(rec_dir, exist_ok=True)
+            
+            if self.cfg != 1.0:
+                rec_cfg_dir = os.path.join(self.image_saved_dir, f"rec_step{self.steps}_cfg_{self.cfg}_slots{self.test_num_slots}")
+                os.makedirs(rec_cfg_dir, exist_ok=True)
+
+            def process_batch(cfg_value, save_dir, header):
+                logger = MetricLogger(delimiter="  ")
+                print_freq = 5
+                psnr_values = []
+                ssim_values = []
+                total_processed = 0
+                
+                for batch_i, batch in enumerate(logger.log_every(self.test_dl, print_freq, header)):
+                    imgs, targets = (batch[0], batch[1]) if isinstance(batch, (tuple, list)) else (batch, None)
+                    
+                    # Skip processing if we've already processed all real samples
+                    if total_processed >= self.test_dataset_size:
+                        break
+                        
+                    imgs = imgs.to(self.device, non_blocking=True)
+                    if targets is not None:
+                        targets = targets.to(self.device, non_blocking=True)
+
+                    with self.accelerator.autocast():
+                        recs = self.model(imgs, sample=True, inference_with_n_slots=self.test_num_slots, cfg=cfg_value)
+
+                    psnr_val = psnr(recs, imgs, data_range=1.0)
+                    ssim_val = ssim(recs, imgs, data_range=1.0)
+                    
+                    recs = concat_all_gather(recs).detach()
+                    psnr_val = concat_all_gather(psnr_val.view(1))
+                    ssim_val = concat_all_gather(ssim_val.view(1))
+
+                    # Remove padding after gathering from all GPUs
+                    samples_in_batch = min(
+                        recs.size(0),  # Always use the gathered size
+                        self.test_dataset_size - total_processed
+                    )
+                    recs = recs[:samples_in_batch]
+                    psnr_val = psnr_val[:samples_in_batch]
+                    ssim_val = ssim_val[:samples_in_batch]
+                    psnr_values.append(psnr_val)
+                    ssim_values.append(ssim_val)
+
+                    if self.accelerator.is_main_process:
+                        rec_paths = [os.path.join(save_dir, f"step_{self.steps}_slots{self.test_num_slots}_{batch_i}_{j}_rec_cfg_{cfg_value}_slots{self.test_num_slots}.png") 
+                                   for j in range(recs.size(0))]
+                        save_img_batch(recs.cpu(), rec_paths)
+                    
+                    total_processed += samples_in_batch
+
+                    self.accelerator.wait_for_everyone()
+                    
+                return torch.cat(psnr_values).mean(), torch.cat(ssim_values).mean()
+
+            # Helper function to calculate and log metrics
+            def calculate_and_log_metrics(real_dir, rec_dir, cfg_value, psnr_val, ssim_val):
+                if self.accelerator.is_main_process:
+                    metrics_dict = get_fid_stats(real_dir, rec_dir, self.fid_stats)
+                    fid = metrics_dict["frechet_inception_distance"]
+                    inception_score = metrics_dict["inception_score_mean"]
+                    
+                    metric_prefix = "fid"
+                    isc_prefix = "isc"
+                    self.accelerator.log({
+                        metric_prefix: fid,
+                        isc_prefix: inception_score,
+                        f"psnr": psnr_val,
+                        f"ssim": ssim_val,
+                        "cfg": cfg_value
+                    }, step=self.steps)
+                    
+                    print(f"{'CFG: {cfg_value}'} "
+                          f"FID: {fid:.2f}, ISC: {inception_score:.2f}, "
+                          f"PSNR: {psnr_val:.2f}, SSIM: {ssim_val:.4f}")
+
+            # Process without CFG
+            if self.cfg == 1.0 or not self.test_only:
+                psnr_val, ssim_val = process_batch(1.0, rec_dir, 'Testing: w/o CFG')
+                calculate_and_log_metrics(real_dir, rec_dir, 1.0, psnr_val, ssim_val)
+
+            # Process with CFG if needed
+            if self.cfg != 1.0:
+                psnr_val, ssim_val = process_batch(self.cfg, rec_cfg_dir, 'Testing: w/ CFG')
+                calculate_and_log_metrics(real_dir, rec_cfg_dir, self.cfg, psnr_val, ssim_val)
+
+            # Cleanup
+            if self.accelerator.is_main_process:
+                shutil.rmtree(rec_dir)
+                if self.cfg != 1.0:
+                    shutil.rmtree(rec_cfg_dir)
+        self.model.train()

semanticist/engine/gpt_trainer.py

@@ -0,0 +1,694 @@
+import os, torch
+import os.path as osp
+import shutil
+import numpy as np
+import copy
+import torch.nn as nn
+from tqdm.auto import tqdm
+from accelerate import Accelerator
+from torchvision.utils import make_grid, save_image
+from torch.utils.data import DataLoader, DistributedSampler
+from semanticist.utils.logger import SmoothedValue, MetricLogger, empty_cache
+from accelerate.utils import DistributedDataParallelKwargs
+from semanticist.stage2.gpt import GPT_models
+from semanticist.stage2.generate import generate
+from pathlib import Path
+import time
+
+from semanticist.engine.trainer_utils import (
+    instantiate_from_config, concat_all_gather,
+    save_img_batch, get_fid_stats,
+    EMAModel, create_scheduler, load_state_dict, load_safetensors,
+    setup_result_folders, create_optimizer,
+    CacheDataLoader
+)
+
+class GPTTrainer(nn.Module):
+    def __init__(
+        self,
+        ae_model,
+        gpt_model,
+        dataset,
+        test_only=False,
+        num_test_images=50000,
+        num_epoch=400,
+        eval_classes=[1, 7, 282, 604, 724, 207, 250, 751, 404, 850], # goldfish, cock, tiger cat, hourglass, ship, golden retriever, husky, race car, airliner, teddy bear
+        blr=1e-4,
+        cosine_lr=False,
+        lr_min=0,
+        warmup_epochs=100,
+        warmup_steps=None,
+        warmup_lr_init=0,
+        decay_steps=None,
+        batch_size=32,
+        cache_bs=8,
+        test_bs=100,
+        num_workers=8,
+        pin_memory=False,
+        max_grad_norm=None,
+        grad_accum_steps=1,
+        precision='bf16',
+        save_every=10000,
+        sample_every=1000,
+        fid_every=50000,
+        result_folder=None,
+        log_dir="./log",
+        ae_cfg=1.0,
+        cfg=6.0,
+        cfg_schedule="linear",
+        temperature=1.0,
+        train_num_slots=None,
+        test_num_slots=None,
+        eval_fid=False,
+        fid_stats=None,
+        enable_ema=False,
+        compile=False,
+        enable_cache_latents=True,
+        cache_dir='/dev/shm/slot_cache'
+    ):
+        super().__init__()
+        kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+        self.accelerator = Accelerator(
+            kwargs_handlers=[kwargs],
+            mixed_precision=precision,
+            gradient_accumulation_steps=grad_accum_steps,
+            log_with="tensorboard",
+            project_dir=log_dir,
+        )
+
+        self.ae_model = instantiate_from_config(ae_model)
+        ae_model_path = ae_model.params.ckpt_path
+        assert ae_model_path.endswith(".safetensors") or ae_model_path.endswith(".pt") or ae_model_path.endswith(".pth") or ae_model_path.endswith(".pkl")
+        assert osp.exists(ae_model_path), f"AE model checkpoint {ae_model_path} does not exist"
+        self._load_checkpoint(ae_model_path, self.ae_model)
+
+        self.ae_model.to(self.device)
+        for param in self.ae_model.parameters():
+            param.requires_grad = False
+        self.ae_model.eval()
+
+        self.model_name = gpt_model.target
+        if 'GPT' in gpt_model.target:
+            self.gpt_model = GPT_models[gpt_model.target](**gpt_model.params)
+        else:
+            raise ValueError(f"Unknown model type: {gpt_model.target}")
+        self.num_slots = ae_model.params.num_slots
+        self.slot_dim = ae_model.params.slot_dim
+
+        self.test_only = test_only
+        self.test_bs = test_bs
+        self.num_test_images = num_test_images
+        self.num_classes = gpt_model.params.num_classes
+        self.batch_size = batch_size
+        if not test_only:
+            self.train_ds = instantiate_from_config(dataset)
+            train_size = len(self.train_ds)
+            if self.accelerator.is_main_process:
+                print(f"train dataset size: {train_size}")
+
+            sampler = DistributedSampler(
+                self.train_ds,
+                num_replicas=self.accelerator.num_processes,
+                rank=self.accelerator.process_index,
+                shuffle=True,
+            )
+            self.train_dl = DataLoader(
+                self.train_ds,
+                batch_size=batch_size if not enable_cache_latents else cache_bs,
+                sampler=sampler,
+                num_workers=num_workers,
+                pin_memory=pin_memory,
+                drop_last=True,
+            )
+
+            effective_bs = batch_size * grad_accum_steps * self.accelerator.num_processes
+            lr = blr * effective_bs / 256
+            if self.accelerator.is_main_process:
+                print(f"Effective batch size is {effective_bs}")
+
+            self.g_optim = create_optimizer(self.gpt_model, weight_decay=0.05, learning_rate=lr)
+            
+            if warmup_epochs is not None:
+                warmup_steps = warmup_epochs * len(self.train_dl)
+                
+            self.g_sched = create_scheduler(
+                self.g_optim,
+                num_epoch,
+                len(self.train_dl),
+                lr_min,
+                warmup_steps,
+                warmup_lr_init,
+                decay_steps,
+                cosine_lr
+            )
+            self.accelerator.register_for_checkpointing(self.g_sched)
+            self.gpt_model, self.g_optim, self.g_sched = self.accelerator.prepare(self.gpt_model, self.g_optim, self.g_sched)
+        else:
+            self.gpt_model = self.accelerator.prepare(self.gpt_model)
+
+        self.steps = 0
+        self.loaded_steps = -1
+
+        if compile:
+            self.ae_model = torch.compile(self.ae_model, mode="reduce-overhead")
+            _model = self.accelerator.unwrap_model(self.gpt_model)
+            _model = torch.compile(_model, mode="reduce-overhead")
+
+        self.enable_ema = enable_ema
+        if self.enable_ema and not self.test_only: # when testing, we directly load the ema dict and skip here
+            self.ema_model = EMAModel(self.accelerator.unwrap_model(self.gpt_model), self.device)
+            self.accelerator.register_for_checkpointing(self.ema_model)
+
+        self._load_checkpoint(gpt_model.params.ckpt_path)
+        if self.test_only:
+            self.steps = self.loaded_steps
+
+        self.num_epoch = num_epoch
+        self.save_every = save_every
+        self.sample_every = sample_every
+        self.fid_every = fid_every
+        self.max_grad_norm = max_grad_norm
+        self.eval_classes = eval_classes
+        self.cfg = cfg
+        self.ae_cfg = ae_cfg
+        self.cfg_schedule = cfg_schedule
+        self.temperature = temperature
+        self.train_num_slots = train_num_slots
+        self.test_num_slots = test_num_slots
+        if self.train_num_slots is not None:
+            self.train_num_slots = min(self.train_num_slots, self.num_slots)
+        else:
+            self.train_num_slots = self.num_slots
+        if self.test_num_slots is not None:
+            self.num_slots_to_gen = min(self.test_num_slots, self.train_num_slots)
+        else:
+            self.num_slots_to_gen = self.train_num_slots
+        self.eval_fid = eval_fid
+        if eval_fid:
+            assert fid_stats is not None
+        self.fid_stats = fid_stats
+
+        # Setup result folders
+        self.result_folder = result_folder
+        self.model_saved_dir, self.image_saved_dir = setup_result_folders(result_folder)
+
+        # Setup cache
+        self.cache_dir = Path(cache_dir)
+        self.enable_cache_latents = enable_cache_latents
+        self.cache_loader = None
+
+    @property
+    def device(self):
+        return self.accelerator.device
+
+    def _load_checkpoint(self, ckpt_path=None, model=None):
+        if ckpt_path is None or not osp.exists(ckpt_path):
+            return
+
+        if model is None:
+            model = self.accelerator.unwrap_model(self.gpt_model)
+
+        if osp.isdir(ckpt_path):
+            self.loaded_steps = int(
+                ckpt_path.split("step")[-1].split("/")[0]
+            )
+            if not self.test_only:
+                self.accelerator.load_state(ckpt_path)
+            else:
+                if self.enable_ema:
+                    model_path = osp.join(ckpt_path, "custom_checkpoint_1.pkl")
+                    if osp.exists(model_path):
+                        state_dict = torch.load(model_path, map_location="cpu")
+                        load_state_dict(state_dict, model)
+                        if self.accelerator.is_main_process:
+                            print(f"Loaded ema model from {model_path}")
+                else:
+                    model_path = osp.join(ckpt_path, "model.safetensors")
+                    if osp.exists(model_path):
+                        load_safetensors(model_path, model)
+        else:
+            if ckpt_path.endswith(".safetensors"):
+                load_safetensors(ckpt_path, model)
+            else:
+                state_dict = torch.load(ckpt_path, map_location="cpu")
+                load_state_dict(state_dict, model)
+        if self.accelerator.is_main_process:
+            print(f"Loaded checkpoint from {ckpt_path}")
+
+    def _build_cache(self):
+        """Build cache for slots and targets."""
+        rank = self.accelerator.process_index
+        world_size = self.accelerator.num_processes
+        
+        # Clean up any existing cache files first
+        slots_file = self.cache_dir / f"slots_rank{rank}_of_{world_size}.mmap"
+        targets_file = self.cache_dir / f"targets_rank{rank}_of_{world_size}.mmap"
+        
+        if slots_file.exists():
+            os.remove(slots_file)
+        if targets_file.exists():
+            os.remove(targets_file)
+        
+        dataset_size = len(self.train_dl.dataset)
+        shard_size = dataset_size // world_size
+        
+        # Detect number of augmentations from first batch
+        with torch.no_grad():
+            sample_batch = next(iter(self.train_dl))
+            img, _ = sample_batch
+            num_augs = img.shape[1] if len(img.shape) == 5 else 1
+        
+        print(f"Rank {rank}: Creating new cache with {num_augs} augmentations per image...")
+        os.makedirs(self.cache_dir, exist_ok=True)
+        slots_file = self.cache_dir / f"slots_rank{rank}_of_{world_size}.mmap"
+        targets_file = self.cache_dir / f"targets_rank{rank}_of_{world_size}.mmap"
+        
+        # Create memory-mapped files
+        slots_mmap = np.memmap(
+            slots_file,
+            dtype='float32',
+            mode='w+',
+            shape=(shard_size * num_augs, self.train_num_slots, self.slot_dim)
+        )
+        
+        targets_mmap = np.memmap(
+            targets_file,
+            dtype='int64',
+            mode='w+',
+            shape=(shard_size * num_augs,)
+        )
+        
+        # Cache data
+        with torch.no_grad():
+            for i, batch in enumerate(tqdm(
+                self.train_dl, 
+                desc=f"Rank {rank}: Caching data",
+                disable=not self.accelerator.is_local_main_process
+            )):
+                imgs, targets = batch
+                if len(imgs.shape) == 5:  # [B, num_augs, C, H, W]
+                    B, A, C, H, W = imgs.shape
+                    imgs = imgs.view(-1, C, H, W)  # [B*num_augs, C, H, W]
+                    targets = targets.unsqueeze(1).expand(-1, A).reshape(-1)  # [B*num_augs]
+                
+                # Split imgs into n chunks
+                num_splits = num_augs
+                split_size = imgs.shape[0] // num_splits
+                imgs_splits = torch.split(imgs, split_size)
+                targets_splits = torch.split(targets, split_size)
+                
+                start_idx = i * self.train_dl.batch_size * num_augs
+                
+                for split_idx, (img_split, targets_split) in enumerate(zip(imgs_splits, targets_splits)):
+                    img_split = img_split.to(self.device, non_blocking=True)
+                    slots_split = self.ae_model.encode_slots(img_split)[:, :self.train_num_slots, :]
+                    
+                    split_start = start_idx + (split_idx * split_size)
+                    split_end = split_start + img_split.shape[0]
+                    
+                    # Write directly to mmap files
+                    slots_mmap[split_start:split_end] = slots_split.cpu().numpy()
+                    targets_mmap[split_start:split_end] = targets_split.numpy()
+        
+        # Close the mmap files
+        del slots_mmap
+        del targets_mmap
+        
+        # Reopen in read mode
+        self.cached_latents = np.memmap(
+            slots_file,
+            dtype='float32',
+            mode='r',
+            shape=(shard_size * num_augs, self.train_num_slots, self.slot_dim)
+        )
+        
+        self.cached_targets = np.memmap(
+            targets_file,
+            dtype='int64',
+            mode='r',
+            shape=(shard_size * num_augs,)
+        )
+        
+        # Store the number of augmentations for the cache loader
+        self.num_augs = num_augs
+
+    def _setup_cache(self):
+        """Setup cache if enabled."""
+        self._build_cache()
+        self.accelerator.wait_for_everyone()
+
+        # Initialize cache loader if cache exists
+        if self.cached_latents is not None:
+            self.cache_loader = CacheDataLoader(
+                slots=self.cached_latents,
+                targets=self.cached_targets,
+                batch_size=self.batch_size,
+                num_augs=self.num_augs,
+                seed=42 + self.accelerator.process_index
+            )
+
+    def __del__(self):
+        """Cleanup cache files."""
+        if self.enable_cache_latents:
+            rank = self.accelerator.process_index
+            world_size = self.accelerator.num_processes
+            
+            # Clean up slots cache
+            slots_file = self.cache_dir / f"slots_rank{rank}_of_{world_size}.mmap"
+            if slots_file.exists():
+                os.remove(slots_file)
+            
+            # Clean up targets cache
+            targets_file = self.cache_dir / f"targets_rank{rank}_of_{world_size}.mmap"
+            if targets_file.exists():
+                os.remove(targets_file)
+
+    def _train_step(self, slots, targets=None):
+        """Execute single training step."""
+        
+        with self.accelerator.accumulate(self.gpt_model):
+            with self.accelerator.autocast():
+                loss = self.gpt_model(slots, targets)
+            
+            self.accelerator.backward(loss)
+            if self.accelerator.sync_gradients and self.max_grad_norm is not None:
+                self.accelerator.clip_grad_norm_(self.gpt_model.parameters(), self.max_grad_norm)
+            self.g_optim.step()
+            if self.g_sched is not None:
+                self.g_sched.step_update(self.steps)
+            self.g_optim.zero_grad()
+
+        # Update EMA model if enabled
+        if self.enable_ema:
+            self.ema_model.update(self.accelerator.unwrap_model(self.gpt_model))
+        
+        return loss
+
+    def _train_epoch_cached(self, epoch, logger):
+        """Train one epoch using cached data."""
+        self.cache_loader.set_epoch(epoch)
+        header = f'Epoch: [{epoch}/{self.num_epoch}]'
+        
+        for batch in logger.log_every(self.cache_loader, 20, header):
+            slots, targets = (b.to(self.device, non_blocking=True) for b in batch)
+            
+            self.steps += 1
+
+            if self.steps == 1:
+                print(f"Training batch size: {len(slots)}")
+                print(f"Hello from index {self.accelerator.local_process_index}")
+            
+            loss = self._train_step(slots, targets)
+            self._handle_periodic_ops(loss, logger)
+
+    def _train_epoch_uncached(self, epoch, logger):
+        """Train one epoch using raw data."""
+        header = f'Epoch: [{epoch}/{self.num_epoch}]'
+        
+        for batch in logger.log_every(self.train_dl, 20, header):
+            img, targets = (b.to(self.device, non_blocking=True) for b in batch)
+            
+            self.steps += 1
+            
+            if self.steps == 1:
+                print(f"Training batch size: {img.size(0)}")
+                print(f"Hello from index {self.accelerator.local_process_index}")
+
+            slots = self.ae_model.encode_slots(img)[:, :self.train_num_slots, :]
+            loss = self._train_step(slots, targets)
+            self._handle_periodic_ops(loss, logger)
+
+    def _handle_periodic_ops(self, loss, logger):
+        """Handle periodic operations and logging."""
+        logger.update(loss=loss.item())
+        logger.update(lr=self.g_optim.param_groups[0]["lr"])
+        
+        if self.steps % self.save_every == 0:
+            self.save()
+        
+        if (self.steps % self.sample_every == 0) or (self.eval_fid and self.steps % self.fid_every == 0):
+            empty_cache()
+            self.evaluate()
+            self.accelerator.wait_for_everyone()
+            empty_cache()
+
+    def _save_config(self, config):
+        """Save configuration file."""
+        if config is not None and self.accelerator.is_main_process:
+            import shutil
+            from omegaconf import OmegaConf
+
+            if isinstance(config, str) and osp.exists(config):
+                shutil.copy(config, osp.join(self.result_folder, "config.yaml"))
+            else:
+                config_save_path = osp.join(self.result_folder, "config.yaml")
+                OmegaConf.save(config, config_save_path)
+
+    def _should_skip_epoch(self, epoch):
+        """Check if epoch should be skipped due to loaded checkpoint."""
+        loader = self.train_dl if not self.enable_cache_latents else self.cache_loader
+        if ((epoch + 1) * len(loader)) <= self.loaded_steps:
+            if self.accelerator.is_main_process:
+                print(f"Epoch {epoch} is skipped because it is loaded from ckpt")
+            self.steps += len(loader)
+            return True
+        
+        if self.steps < self.loaded_steps:
+            for _ in loader:
+                self.steps += 1
+                if self.steps >= self.loaded_steps:
+                    break
+        return False
+
+    def train(self, config=None):
+        """Main training loop."""
+        # Initial setup
+        n_parameters = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        if self.accelerator.is_main_process:
+            print(f"number of learnable parameters: {n_parameters//1e6}M")
+        
+        self._save_config(config)
+        self.accelerator.init_trackers("gpt")
+
+        # Handle test-only mode
+        if self.test_only:
+            empty_cache()
+            self.evaluate()
+            self.accelerator.wait_for_everyone()
+            empty_cache()
+            return
+
+        # Setup cache if enabled
+        if self.enable_cache_latents:
+            self._setup_cache()
+
+        # Training loop
+        for epoch in range(self.num_epoch):
+            if self._should_skip_epoch(epoch):
+                continue
+
+            self.gpt_model.train()
+            logger = MetricLogger(delimiter="  ")
+            logger.add_meter('lr', SmoothedValue(window_size=1, fmt='{value:.6f}'))
+
+            # Choose training path based on cache availability
+            if self.enable_cache_latents:
+                self._train_epoch_cached(epoch, logger)
+            else:
+                self._train_epoch_uncached(epoch, logger)
+
+            # Synchronize and log epoch stats
+            logger.synchronize_between_processes()
+            if self.accelerator.is_main_process:
+                print("Averaged stats:", logger)
+
+        # Finish training
+        self.accelerator.end_training()
+        self.save()
+        if self.accelerator.is_main_process:
+            print("Train finished!")
+
+    def save(self):
+        self.accelerator.wait_for_everyone()
+        self.accelerator.save_state(
+            os.path.join(self.model_saved_dir, f"step{self.steps}")
+        )
+
+    @torch.no_grad()
+    def evaluate(self, use_ema=True):
+        self.gpt_model.eval()
+        unwraped_gpt_model = self.accelerator.unwrap_model(self.gpt_model)
+        # switch to ema params, only when eval_fid is True
+        # if test_only, we directly load the ema dict and skip here
+        use_ema = use_ema and self.enable_ema and self.eval_fid and not self.test_only
+        if use_ema:
+            if hasattr(self, "ema_model"):
+                model_without_ddp = self.accelerator.unwrap_model(self.gpt_model)
+                model_state_dict = copy.deepcopy(model_without_ddp.state_dict())
+                ema_state_dict = copy.deepcopy(model_without_ddp.state_dict())
+                for i, (name, _value) in enumerate(model_without_ddp.named_parameters()):
+                    if "nested_sampler" in name:
+                        continue
+                    ema_state_dict[name] = self.ema_model.state_dict()[name]
+                if self.accelerator.is_main_process:
+                    print("Switch to ema")
+                model_without_ddp.load_state_dict(ema_state_dict)
+            else:
+                print("EMA model not found, using original model")
+                use_ema = False
+
+        if not self.test_only:
+            classes = torch.tensor(self.eval_classes, device=self.device)
+            with self.accelerator.autocast():
+                slots = generate(unwraped_gpt_model, classes, self.num_slots_to_gen, cfg_scale=self.cfg, cfg_schedule=self.cfg_schedule, temperature=self.temperature)
+                if self.num_slots_to_gen < self.num_slots:
+                    null_slots = self.ae_model.dit.null_cond.expand(slots.shape[0], -1, -1)
+                    null_slots = null_slots[:, self.num_slots_to_gen:, :]
+                    slots = torch.cat([slots, null_slots], dim=1)
+                imgs = self.ae_model.sample(slots, targets=classes, cfg=self.ae_cfg) # targets are not used for now
+
+            imgs = concat_all_gather(imgs)
+            if self.accelerator.num_processes > 16:
+                imgs = imgs[:16*len(self.eval_classes)]
+            imgs = imgs.detach().cpu()
+            grid = make_grid(
+                imgs, nrow=len(self.eval_classes), normalize=True, value_range=(0, 1)
+            )
+            if self.accelerator.is_main_process:
+                save_image(
+                    grid,
+                    os.path.join(
+                        self.image_saved_dir, f"step{self.steps}_aecfg-{self.ae_cfg}_cfg-{self.cfg_schedule}-{self.cfg}_slots{self.num_slots_to_gen}_temp{self.temperature}.jpg"
+                    ),
+                )
+        if self.eval_fid and (self.test_only or (self.steps % self.fid_every == 0)):
+            # Create output directory (only on main process)
+            save_folder = os.path.join(self.image_saved_dir, f"gen_step{self.steps}_aecfg-{self.ae_cfg}_cfg-{self.cfg_schedule}-{self.cfg}_slots{self.num_slots_to_gen}_temp{self.temperature}")
+            if self.accelerator.is_main_process:
+                os.makedirs(save_folder, exist_ok=True)
+
+            # Setup for distributed generation
+            world_size = self.accelerator.num_processes
+            local_rank = self.accelerator.process_index
+            batch_size = self.test_bs
+            
+            # Create balanced class distribution
+            num_classes = self.num_classes
+            images_per_class = self.num_test_images // num_classes
+            class_labels = np.repeat(np.arange(num_classes), images_per_class)
+            
+            # Shuffle the class labels to ensure random ordering
+            np.random.shuffle(class_labels)
+            
+            total_images = len(class_labels)
+
+            padding_size = world_size * batch_size - (total_images % (world_size * batch_size))
+            class_labels = np.pad(class_labels, (0, padding_size), 'constant')
+            padded_total_images = len(class_labels)
+
+            # Distribute workload across GPUs
+            images_per_gpu = padded_total_images // world_size
+            start_idx = local_rank * images_per_gpu
+            end_idx = min(start_idx + images_per_gpu, padded_total_images)
+            local_class_labels = class_labels[start_idx:end_idx]
+            local_num_steps = len(local_class_labels) // batch_size
+            
+            if self.accelerator.is_main_process:
+                print(f"Generating {total_images} images ({images_per_class} per class) across {world_size} GPUs")
+            
+            used_time = 0
+            gen_img_cnt = 0
+            
+            for i in range(local_num_steps):
+                if self.accelerator.is_main_process and i % 10 == 0:
+                    print(f"Generation step {i}/{local_num_steps}")
+
+                # Get and pad labels for current batch
+                batch_start = i * batch_size
+                batch_end = batch_start + batch_size
+                labels = local_class_labels[batch_start:batch_end]
+                
+                # Convert to tensors and track real vs padding
+                labels = torch.tensor(labels, device=self.device)
+                
+                # Generate images
+                self.accelerator.wait_for_everyone()
+                start_time = time.time()
+                with torch.no_grad():
+                    with self.accelerator.autocast():
+                        slots = generate(unwraped_gpt_model, labels, self.num_slots_to_gen, 
+                                            cfg_scale=self.cfg, 
+                                            cfg_schedule=self.cfg_schedule,
+                                            temperature=self.temperature)
+                        if self.num_slots_to_gen < self.num_slots:
+                            null_slots = self.ae_model.dit.null_cond.expand(slots.shape[0], -1, -1)
+                            null_slots = null_slots[:, self.num_slots_to_gen:, :]
+                            slots = torch.cat([slots, null_slots], dim=1)
+                        imgs = self.ae_model.sample(slots, targets=labels, cfg=self.ae_cfg)
+
+                samples_in_batch = min(batch_size * world_size, total_images - gen_img_cnt)
+
+                # Update timing stats
+                used_time += time.time() - start_time
+                gen_img_cnt += samples_in_batch
+                if self.accelerator.is_main_process and i % 10 == 0:
+                    print(f"Avg generation time: {used_time/gen_img_cnt:.5f} sec/image")
+
+                gathered_imgs = concat_all_gather(imgs)
+                gathered_imgs = gathered_imgs[:samples_in_batch]
+                
+                # Save images (only on main process)
+                if self.accelerator.is_main_process:
+                    real_imgs = gathered_imgs.detach().cpu()
+                    
+                    save_paths = [
+                        os.path.join(save_folder, f"{str(idx).zfill(5)}.png")
+                        for idx in range(gen_img_cnt - samples_in_batch, gen_img_cnt)
+                    ]
+                    save_img_batch(real_imgs, save_paths)
+
+            # Calculate metrics (only on main process)
+            self.accelerator.wait_for_everyone()
+            if self.accelerator.is_main_process:
+                generated_files = len(os.listdir(save_folder))
+                print(f"Generated {generated_files} images out of {total_images} expected")
+                
+                metrics_dict = get_fid_stats(save_folder, None, self.fid_stats)
+                fid = metrics_dict["frechet_inception_distance"]
+                inception_score = metrics_dict["inception_score_mean"]
+                
+                metric_prefix = "fid_ema" if use_ema else "fid"
+                isc_prefix = "isc_ema" if use_ema else "isc"
+                
+                self.accelerator.log({
+                    metric_prefix: fid,
+                    isc_prefix: inception_score,
+                    "gpt_cfg": self.cfg,
+                    "ae_cfg": self.ae_cfg,
+                    "cfg_schedule": self.cfg_schedule,
+                    "temperature": self.temperature,
+                    "num_slots": self.test_num_slots if self.test_num_slots is not None else self.train_num_slots
+                }, step=self.steps)
+                
+                # Print comprehensive CFG information
+                cfg_info = (
+                    f"{'EMA ' if use_ema else ''}CFG params: "
+                    f"gpt_cfg={self.cfg}, ae_cfg={self.ae_cfg}, "
+                    f"cfg_schedule={self.cfg_schedule}, "
+                    f"num_slots={self.test_num_slots if self.test_num_slots is not None else self.train_num_slots}, "
+                    f"temperature={self.temperature}"
+                )
+                print(cfg_info)
+                print(f"FID: {fid:.2f}, ISC: {inception_score:.2f}")
+
+                # Cleanup
+                shutil.rmtree(save_folder)
+
+        # back to no ema
+        if use_ema:
+            if self.accelerator.is_main_process:
+                print("Switch back from ema")
+            model_without_ddp.load_state_dict(model_state_dict)
+
+        self.gpt_model.train()
+

semanticist/engine/trainer_utils.py

@@ -0,0 +1,251 @@
+import os, torch
+import cv2
+import numpy as np
+import torch_fidelity
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor
+import importlib
+from torch.optim import AdamW
+from semanticist.utils.lr_scheduler import build_scheduler
+
+
+def get_obj_from_str(string, reload=False):
+    """Get object from string path."""
+    module, cls = string.rsplit(".", 1)
+    if reload:
+        module_imp = importlib.import_module(module)
+        importlib.reload(module_imp)
+    return getattr(importlib.import_module(module, package=None), cls)
+
+
+def instantiate_from_config(config):
+    """Instantiate an object from a config dictionary."""
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def is_dist_avail_and_initialized():
+    """Check if distributed training is available and initialized."""
+    if not torch.distributed.is_initialized():
+        return False
+    return True
+
+
+def is_main_process():
+    """Check if the current process is the main process."""
+    return not is_dist_avail_and_initialized() or torch.distributed.get_rank() == 0
+
+
+def concat_all_gather(tensor):
+    """
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+    tensors_gather = [torch.ones_like(tensor)
+        for _ in range(torch.distributed.get_world_size())]
+    torch.distributed.all_gather(tensors_gather, tensor, async_op=False)
+
+    output = torch.cat(tensors_gather, dim=0)
+    return output
+
+
+def requires_grad(model, flag=True):
+    """Set requires_grad flag for all model parameters."""
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def save_img(img, save_path):
+    """Save a single image to disk."""
+    img = np.clip(img.float().numpy().transpose([1, 2, 0]) * 255, 0, 255)
+    img = img.astype(np.uint8)[:, :, ::-1]
+    cv2.imwrite(save_path, img)
+
+
+def save_img_batch(imgs, save_paths):
+    """Process and save multiple images at once using a thread pool."""
+    # Convert to numpy and prepare all images in one go
+    imgs = np.clip(imgs.float().numpy().transpose(0, 2, 3, 1) * 255, 0, 255).astype(np.uint8)
+    imgs = imgs[:, :, :, ::-1]  # RGB to BGR for all images at once
+    
+    with ThreadPoolExecutor(max_workers=32) as pool:
+        # Submit all tasks at once
+        futures = [pool.submit(cv2.imwrite, path, img) 
+                  for path, img in zip(save_paths, imgs)]
+        # Wait for all tasks to complete
+        for future in futures:
+            future.result()  # This will raise any exceptions that occurred
+
+
+def get_fid_stats(real_dir, rec_dir, fid_stats):
+    """Calculate FID statistics between real and reconstructed images."""
+    stats = torch_fidelity.calculate_metrics(
+        input1=rec_dir,
+        input2=real_dir,
+        fid_statistics_file=fid_stats,
+        cuda=True,
+        isc=True,
+        fid=True,
+        kid=False,
+        prc=False,
+        verbose=False,
+    )
+    return stats
+
+
+def create_scheduler(optimizer, num_epoch, steps_per_epoch, lr_min, warmup_steps, 
+                    warmup_lr_init, decay_steps, cosine_lr):
+    """Create a learning rate scheduler."""
+    scheduler = build_scheduler(
+        optimizer,
+        num_epoch,
+        steps_per_epoch,
+        lr_min,
+        warmup_steps,
+        warmup_lr_init,
+        decay_steps,
+        cosine_lr,
+    )
+    return scheduler
+
+
+def load_state_dict(state_dict, model):
+    """Helper to load a state dict with proper prefix handling."""
+    if 'state_dict' in state_dict:
+        state_dict = state_dict['state_dict']
+    # Remove '_orig_mod' prefix if present
+    state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+    missing, unexpected = model.load_state_dict(
+        state_dict, strict=False
+    )
+    if is_main_process():
+        print(f"Loaded model. Missing: {missing}, Unexpected: {unexpected}")
+
+
+def load_safetensors(path, model):
+    """Helper to load a safetensors checkpoint."""
+    from safetensors.torch import safe_open
+    with safe_open(path, framework="pt", device="cpu") as f:
+        state_dict = {k: f.get_tensor(k) for k in f.keys()}
+    load_state_dict(state_dict, model)
+
+
+def setup_result_folders(result_folder):
+    """Setup result folders for saving models and images."""
+    model_saved_dir = os.path.join(result_folder, "models")
+    os.makedirs(model_saved_dir, exist_ok=True)
+
+    image_saved_dir = os.path.join(result_folder, "images")
+    os.makedirs(image_saved_dir, exist_ok=True)
+    
+    return model_saved_dir, image_saved_dir
+
+
+def create_optimizer(model, weight_decay, learning_rate, betas=(0.9, 0.95)):
+    """Create an AdamW optimizer with weight decay for 2D parameters only."""
+    # start with all of the candidate parameters
+    param_dict = {pn: p for pn, p in model.named_parameters()}
+    # filter out those that do not require grad
+    param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
+    # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
+    # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
+    decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
+    nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
+    optim_groups = [
+        {'params': decay_params, 'weight_decay': weight_decay},
+        {'params': nodecay_params, 'weight_decay': 0.0}
+    ]
+    num_decay_params = sum(p.numel() for p in decay_params)
+    num_nodecay_params = sum(p.numel() for p in nodecay_params)
+    if is_main_process():
+        print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
+        print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
+    optimizer = AdamW(optim_groups, lr=learning_rate, betas=betas)
+    return optimizer
+
+
+class EMAModel:
+    """Model Exponential Moving Average."""
+    def __init__(self, model, device, decay=0.999):
+        self.device = device
+        self.decay = decay
+        self.ema_params = OrderedDict(
+            (name, param.clone().detach().to(device))
+            for name, param in model.named_parameters()
+            if param.requires_grad
+        )
+
+    @torch.no_grad()
+    def update(self, model):
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                if name in self.ema_params:
+                    self.ema_params[name].lerp_(param.data, 1 - self.decay)
+                else:
+                    self.ema_params[name] = param.data.clone().detach()
+
+    def state_dict(self):
+        return self.ema_params
+
+    def load_state_dict(self, params):
+        self.ema_params = OrderedDict(
+            (name, param.clone().detach().to(self.device))
+            for name, param in params.items()
+        )
+
+
+class PaddedDataset(torch.utils.data.Dataset):
+    """Dataset wrapper that pads a dataset to ensure even distribution across processes."""
+    def __init__(self, dataset, padding_size):
+        self.dataset = dataset
+        self.padding_size = padding_size
+        
+    def __len__(self):
+        return len(self.dataset) + self.padding_size
+        
+    def __getitem__(self, idx):
+        if idx < len(self.dataset):
+            return self.dataset[idx]
+        return self.dataset[0]
+
+class CacheDataLoader:
+    """DataLoader-like interface for cached data with epoch-based shuffling."""
+    def __init__(self, slots, targets=None, batch_size=32, num_augs=1, seed=None):
+        self.slots = slots
+        self.targets = targets
+        self.batch_size = batch_size
+        self.num_augs = num_augs
+        self.seed = seed
+        self.epoch = 0
+        # Original dataset size (before augmentations)
+        self.num_samples = len(slots) // num_augs
+    
+    def set_epoch(self, epoch):
+        """Set epoch for deterministic shuffling."""
+        self.epoch = epoch
+    
+    def __len__(self):
+        """Return number of batches based on original dataset size."""
+        return self.num_samples // self.batch_size
+    
+    def __iter__(self):
+        """Return random indices for current epoch."""
+        g = torch.Generator()
+        g.manual_seed(self.seed + self.epoch if self.seed is not None else self.epoch)
+        
+        # Randomly sample indices from the entire augmented dataset
+        indices = torch.randint(
+            0, len(self.slots), 
+            (self.num_samples,), 
+            generator=g
+        ).numpy()
+        
+        # Yield batches of indices
+        for start in range(0, self.num_samples, self.batch_size):
+            end = min(start + self.batch_size, self.num_samples)
+            batch_indices = indices[start:end]
+            yield (
+                torch.from_numpy(self.slots[batch_indices]),
+                torch.from_numpy(self.targets[batch_indices])
+            )

semanticist/stage1/diffuse_slot.py

@@ -0,0 +1,452 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers import AutoencoderKL
+from semanticist.stage1 import vision_transformer
+from semanticist.stage1.diffusion import create_diffusion
+from semanticist.stage1.diffusion_transfomer import DiT
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+
+class DiT_with_autoenc_cond(DiT):
+    def __init__(
+        self,
+        *args,
+        num_autoenc=32,
+        autoenc_dim=4,
+        use_repa=False,
+        z_dim=768,
+        encoder_depth=8,
+        projector_dim=2048,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.autoenc_dim = autoenc_dim
+        self.hidden_size = kwargs["hidden_size"]
+        self.null_cond = nn.Parameter(torch.zeros(1, num_autoenc, autoenc_dim))
+        torch.nn.init.normal_(self.null_cond, std=.02)
+        self.autoenc_cond_embedder = nn.Linear(autoenc_dim, self.hidden_size)
+        self.y_embedder = nn.Identity()
+        self.cond_drop_prob = 0.1
+        
+        self.use_repa = use_repa
+        self._repa_hook = None
+        self.encoder_depth = encoder_depth
+        if use_repa:
+            self.projector = build_mlp(self.hidden_size, projector_dim, z_dim)
+
+    def embed_cond(self, autoenc_cond, drop_mask=None):
+        # autoenc_cond: (N, K, D)
+        # drop_ids: (N)
+        # self.null_cond: (1, K, D)
+        batch_size = autoenc_cond.shape[0]
+        if drop_mask is None:
+            # randomly drop all conditions, for classifier-free guidance
+            if self.training:
+                drop_ids = (
+                    torch.rand(batch_size, 1, 1, device=autoenc_cond.device)
+                    < self.cond_drop_prob
+                )
+                autoenc_cond_drop = torch.where(drop_ids, self.null_cond, autoenc_cond)
+            else:
+                autoenc_cond_drop = autoenc_cond
+        else:
+            # randomly drop some conditions according to the drop_mask (N, K)
+            # True means keep
+            autoenc_cond_drop = torch.where(drop_mask[:, :, None], autoenc_cond, self.null_cond)
+        return self.autoenc_cond_embedder(autoenc_cond_drop)
+
+    def forward(self, x, t, autoenc_cond, drop_mask=None):
+        """
+        Forward pass of DiT.
+        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        autoenc_cond: (N, K, D) tensor of autoencoder conditions (slots)
+        """
+        x = (
+            self.x_embedder(x) + self.pos_embed
+        )  # (N, T, D), where T = H * W / patch_size ** 2
+        c = self.t_embedder(t)  # (N, D)
+        autoenc = self.embed_cond(autoenc_cond, drop_mask)
+        num_tokens = x.shape[1]
+        x = torch.cat((x, autoenc), dim=1)
+
+        for i, block in enumerate(self.blocks):
+            x = block(x, c)  # (N, T, D)
+            if (i + 1) == self.encoder_depth and self.use_repa:
+                projected = self.projector(x)
+                self._repa_hook = projected[:, :num_tokens]
+
+        x = x[:, :num_tokens]
+        x = self.final_layer(x, c)  # (N, T, patch_size ** 2 * out_channels)
+        x = self.unpatchify(x)  # (N, out_channels, H, W)
+        return x
+
+    def forward_with_cfg(self, x, t, autoenc_cond, drop_mask, y=None, cfg_scale=1.0):
+        """
+        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, autoenc_cond, drop_mask)
+        eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+
+#################################################################################
+#                                   DiT Configs                                  #
+#################################################################################
+
+
+def DiT_with_autoenc_cond_XL_2(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_XL_4(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_XL_8(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_L_2(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_L_4(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_L_8(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_B_2(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_B_4(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_B_8(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_S_2(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_S_4(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs
+    )
+
+
+def DiT_with_autoenc_cond_S_8(**kwargs):
+    return DiT_with_autoenc_cond(
+        depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs
+    )
+
+
+DiT_with_autoenc_cond_models = {
+    "DiT-XL-2": DiT_with_autoenc_cond_XL_2,
+    "DiT-XL-4": DiT_with_autoenc_cond_XL_4,
+    "DiT-XL-8": DiT_with_autoenc_cond_XL_8,
+    "DiT-L-2": DiT_with_autoenc_cond_L_2,
+    "DiT-L-4": DiT_with_autoenc_cond_L_4,
+    "DiT-L-8": DiT_with_autoenc_cond_L_8,
+    "DiT-B-2": DiT_with_autoenc_cond_B_2,
+    "DiT-B-4": DiT_with_autoenc_cond_B_4,
+    "DiT-B-8": DiT_with_autoenc_cond_B_8,
+    "DiT-S-2": DiT_with_autoenc_cond_S_2,
+    "DiT-S-4": DiT_with_autoenc_cond_S_4,
+    "DiT-S-8": DiT_with_autoenc_cond_S_8,
+}
+
+class NestedSampler(nn.Module):
+    def __init__(
+        self,
+        num_slots,
+    ):
+        super().__init__()
+        self.num_slots = num_slots
+        self.register_buffer("arange", torch.arange(num_slots))
+
+    def uniform_sample(self, num):
+        return torch.randint(1, self.num_slots + 1, (num,))
+
+    def sample(self, num):
+        samples = self.uniform_sample(num)
+        return samples
+
+    def forward(self, batch_size, device, inference_with_n_slots=-1):
+        if self.training:
+            b = self.sample(batch_size).to(device)
+        else:
+            if inference_with_n_slots != -1:
+                b = torch.full((batch_size,), inference_with_n_slots, device=device)
+            else:
+                b = torch.full((batch_size,), self.num_slots, device=device)
+        b = torch.clamp(b, max=self.num_slots)
+
+        slot_mask = self.arange[None, :] < b[:, None]  # (batch_size, num_slots)
+        return slot_mask
+
+class DiffuseSlot(nn.Module):
+    def __init__(
+        self,
+        encoder="vit_base_patch16",
+        drop_path_rate=0.1,
+        enc_img_size=256,
+        enc_causal=True,
+        num_slots=16,
+        slot_dim=256,
+        norm_slots=False,
+        enable_nest=False,
+        enable_nest_after=-1,
+        vae="stabilityai/sd-vae-ft-ema",
+        dit_model="DiT-B-4",
+        num_sampling_steps="ddim25",
+        use_repa=False,
+        repa_encoder_depth=8,
+        repa_loss_weight=1.0,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.use_repa = use_repa
+        self.repa_loss_weight = repa_loss_weight
+        if use_repa:
+            self.repa_encoder = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
+            self.repa_encoder.image_size = 224
+            for param in self.repa_encoder.parameters():
+                param.requires_grad = False
+            self.repa_encoder.eval()
+
+        self.diffusion = create_diffusion(timestep_respacing="")
+        self.gen_diffusion = create_diffusion(timestep_respacing=num_sampling_steps)
+        self.dit_input_size = enc_img_size // 8 if not "mar" in vae else enc_img_size // 16
+        self.dit_in_channels = 4 if not "mar" in vae else 16
+        self.dit = DiT_with_autoenc_cond_models[dit_model](
+            input_size=self.dit_input_size,
+            in_channels=self.dit_in_channels,
+            num_autoenc=num_slots,
+            autoenc_dim=slot_dim,
+            use_repa=use_repa,
+            encoder_depth=repa_encoder_depth,
+            z_dim=768,
+        )
+        self.vae = AutoencoderKL.from_pretrained(vae)
+        self.scaling_factor = self.vae.config.scaling_factor
+        self.vae.eval().requires_grad_(False)
+
+        self.enc_img_size = enc_img_size
+        self.enc_causal = enc_causal
+        encoder_fn = vision_transformer.__dict__[encoder]
+
+        self.encoder = encoder_fn(
+            img_size=[enc_img_size],
+            num_slots=num_slots,
+            drop_path_rate=drop_path_rate,
+        )
+        self.num_slots = num_slots
+        self.norm_slots = norm_slots
+        self.num_channels = self.encoder.num_features
+        
+        self.encoder2slot = nn.Linear(self.num_channels, slot_dim)
+        self.nested_sampler = NestedSampler(num_slots)
+        self.enable_nest = enable_nest
+        self.enable_nest_after = enable_nest_after
+
+    @torch.no_grad()
+    def vae_encode(self, x):
+        x = x * 2 - 1
+        x = self.vae.encode(x)
+        if hasattr(x, 'latent_dist'):
+            x = x.latent_dist
+        return x.sample().mul_(self.scaling_factor)
+
+    @torch.no_grad()
+    def vae_decode(self, z):
+        z = self.vae.decode(z / self.scaling_factor)
+        if hasattr(z, 'sample'):
+            z = z.sample
+        return (z + 1) / 2
+
+    @torch.no_grad()
+    def repa_encode(self, x):
+        mean = torch.Tensor(IMAGENET_DEFAULT_MEAN).to(x.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+        std = torch.Tensor(IMAGENET_DEFAULT_STD).to(x.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+        x = (x - mean) / std
+        if self.repa_encoder.image_size != self.enc_img_size:
+            x = torch.nn.functional.interpolate(x, self.repa_encoder.image_size, mode='bicubic')
+        x = self.repa_encoder.forward_features(x)['x_norm_patchtokens']
+        return x
+
+    def encode_slots(self, x):
+        slots = self.encoder(x, is_causal=self.enc_causal)
+        slots = self.encoder2slot(slots)
+        if self.norm_slots:
+            slots_std = torch.std(slots, dim=-1, keepdim=True)
+            slots_mean = torch.mean(slots, dim=-1, keepdim=True)
+            slots = (slots - slots_mean) / slots_std
+        return slots
+    
+    def forward_with_latents(self,
+                             x_vae,
+                             slots,
+                             z,
+                             sample=False,
+                             epoch=None,
+                             inference_with_n_slots=-1, 
+                             cfg=1.0):
+        losses = {}
+        batch_size = x_vae.shape[0]
+        device = x_vae.device
+        
+        if (
+            epoch is not None
+            and epoch >= self.enable_nest_after
+            and self.enable_nest_after != -1
+        ):
+            self.enable_nest = True
+
+        t = torch.randint(0, 1000, (x_vae.shape[0],), device=device)
+
+        if self.enable_nest or inference_with_n_slots != -1:
+            drop_mask = self.nested_sampler(
+                batch_size, device, 
+                inference_with_n_slots=inference_with_n_slots, 
+            )
+        else:
+            drop_mask = None
+            
+        if sample:
+            return self.sample(slots, drop_mask=drop_mask, cfg=cfg)
+
+        model_kwargs = dict(autoenc_cond=slots, drop_mask=drop_mask)
+        loss_dict = self.diffusion.training_losses(self.dit, x_vae, t, model_kwargs)
+        diff_loss = loss_dict["loss"].mean()
+        losses["diff_loss"] = diff_loss
+        
+        if self.use_repa:
+            assert self.dit._repa_hook is not None and z is not None
+            z_tilde = self.dit._repa_hook
+            
+            if z_tilde.shape[1] != z.shape[1]:
+                z_tilde = interpolate_features(z_tilde, z.shape[1])
+            
+            z_tilde = F.normalize(z_tilde, dim=-1)
+            z = F.normalize(z, dim=-1)
+            repa_loss = -torch.sum(z_tilde * z, dim=-1)
+            losses["repa_loss"] = repa_loss.mean() * self.repa_loss_weight
+        
+        return losses
+        
+
+    def forward(self, 
+                x,
+                sample=False,
+                epoch=None,
+                inference_with_n_slots=-1,
+                cfg=1.0):
+
+        x_vae = self.vae_encode(x)
+        z = self.repa_encode(x) if self.use_repa else None
+        slots = self.encode_slots(x)
+        return self.forward_with_latents(x_vae, slots, z, sample, epoch, inference_with_n_slots, cfg)
+
+
+    @torch.no_grad()
+    def sample(self, slots, drop_mask=None, cfg=1.0):
+        batch_size = slots.shape[0]
+        device = slots.device
+        z = torch.randn(batch_size, self.dit_in_channels, self.dit_input_size, self.dit_input_size, device=device)
+        if cfg != 1.0:
+            z = torch.cat([z, z], 0)
+            null_slots = self.dit.null_cond.expand(batch_size, -1, -1)
+            slots = torch.cat([slots, null_slots], 0)
+            if drop_mask is not None:
+                null_cond_mask = torch.ones_like(drop_mask)
+                drop_mask = torch.cat([drop_mask, null_cond_mask], 0)
+            model_kwargs = dict(autoenc_cond=slots, drop_mask=drop_mask, cfg_scale=cfg)
+            sample_fn = self.dit.forward_with_cfg
+        else:
+            model_kwargs = dict(autoenc_cond=slots, drop_mask=drop_mask)
+            sample_fn = self.dit.forward
+        samples = self.gen_diffusion.p_sample_loop(
+            sample_fn,
+            z.shape,
+            z,
+            clip_denoised=False,
+            model_kwargs=model_kwargs,
+            progress=False,
+            device=device,
+        )
+        if cfg != 1.0:
+            samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+        samples = self.vae_decode(samples)
+        return samples
+
+    def train(self, mode=True):
+        """Override train() to keep certain components in eval mode"""
+        super().train(mode)
+        self.vae.eval()
+        return self
+
+
+def build_mlp(hidden_size, projector_dim, z_dim):
+    return nn.Sequential(
+                nn.Linear(hidden_size, projector_dim),
+                nn.SiLU(),
+                nn.Linear(projector_dim, projector_dim),
+                nn.SiLU(),
+                nn.Linear(projector_dim, z_dim),
+            )
+
+def interpolate_features(x, target_len):
+    """Interpolate features to match target sequence length.
+    Args:
+        x: tensor of shape (B, T1, D)
+        target_len: desired sequence length T2
+    Returns:
+        tensor of shape (B, T2, D)
+    """
+    B, T1, D = x.shape
+    H1 = W1 = int(math.sqrt(T1))
+    H2 = W2 = int(math.sqrt(target_len))
+    
+    # Reshape to 2D spatial dimensions and move channels to second dimension
+    x = x.reshape(B, H1, W1, D).permute(0, 3, 1, 2)
+    
+    # Interpolate
+    x = F.interpolate(x, size=(H2, W2), mode='bicubic', align_corners=False)
+    
+    # Reshape back to sequence
+    return x.permute(0, 2, 3, 1).reshape(B, target_len, D)

semanticist/stage1/diffusion/__init__.py

@@ -0,0 +1,46 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+from . import gaussian_diffusion as gd
+from .respace import SpacedDiffusion, space_timesteps
+
+
+def create_diffusion(
+    timestep_respacing,
+    noise_schedule="linear", 
+    use_kl=False,
+    sigma_small=False,
+    predict_xstart=False,
+    learn_sigma=True,
+    rescale_learned_sigmas=False,
+    diffusion_steps=1000
+):
+    betas = gd.get_named_beta_schedule(noise_schedule, diffusion_steps)
+    if use_kl:
+        loss_type = gd.LossType.RESCALED_KL
+    elif rescale_learned_sigmas:
+        loss_type = gd.LossType.RESCALED_MSE
+    else:
+        loss_type = gd.LossType.MSE
+    if timestep_respacing is None or timestep_respacing == "":
+        timestep_respacing = [diffusion_steps]
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(diffusion_steps, timestep_respacing),
+        betas=betas,
+        model_mean_type=(
+            gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X
+        ),
+        model_var_type=(
+            (
+                gd.ModelVarType.FIXED_LARGE
+                if not sigma_small
+                else gd.ModelVarType.FIXED_SMALL
+            )
+            if not learn_sigma
+            else gd.ModelVarType.LEARNED_RANGE
+        ),
+        loss_type=loss_type
+        # rescale_timesteps=rescale_timesteps,
+    )

semanticist/stage1/diffusion/diffusion_utils.py

@@ -0,0 +1,88 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+import torch as th
+import numpy as np
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def continuous_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a continuous Gaussian distribution.
+    :param x: the targets
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    normalized_x = centered_x * inv_stdv
+    log_probs = th.distributions.Normal(th.zeros_like(x), th.ones_like(x)).log_prob(normalized_x)
+    return log_probs
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs

semanticist/stage1/diffusion/gaussian_diffusion.py

@@ -0,0 +1,886 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+
+import math
+
+import numpy as np
+import torch as th
+import enum
+
+from .diffusion_utils import discretized_gaussian_log_likelihood, normal_kl
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
+    START_X = enum.auto()  # the model predicts x_0
+    EPSILON = enum.auto()  # the model predicts epsilon
+
+
+class ModelVarType(enum.Enum):
+    """
+    What is used as the model's output variance.
+    The LEARNED_RANGE option has been added to allow the model to predict
+    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
+    """
+
+    LEARNED = enum.auto()
+    FIXED_SMALL = enum.auto()
+    FIXED_LARGE = enum.auto()
+    LEARNED_RANGE = enum.auto()
+
+
+class LossType(enum.Enum):
+    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        enum.auto()
+    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    KL = enum.auto()  # use the variational lower-bound
+    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
+    betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    warmup_time = int(num_diffusion_timesteps * warmup_frac)
+    betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64)
+    return betas
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    """
+    This is the deprecated API for creating beta schedules.
+    See get_named_beta_schedule() for the new library of schedules.
+    """
+    if beta_schedule == "quad":
+        betas = (
+            np.linspace(
+                beta_start ** 0.5,
+                beta_end ** 0.5,
+                num_diffusion_timesteps,
+                dtype=np.float64,
+            )
+            ** 2
+        )
+    elif beta_schedule == "linear":
+        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "warmup10":
+        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1)
+    elif beta_schedule == "warmup50":
+        betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5)
+    elif beta_schedule == "const":
+        betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
+        betas = 1.0 / np.linspace(
+            num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64
+        )
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        return get_beta_schedule(
+            "linear",
+            beta_start=scale * 0.0001,
+            beta_end=scale * 0.02,
+            num_diffusion_timesteps=num_diffusion_timesteps,
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+    Original ported from this codebase:
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type
+    ):
+
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        ) if len(self.posterior_variance) > 1 else np.array([])
+
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+        In other words, sample from q(x_t | x_0).
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+            q(x_{t-1} | x_t, x_0)
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, t, **model_kwargs)
+        if isinstance(model_output, tuple):
+            model_output, extra = model_output
+        else:
+            extra = None
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape)
+            max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+            # The model_var_values is [-1, 1] for [min_var, max_var].
+            frac = (model_var_values + 1) / 2
+            model_log_variance = frac * max_log + (1 - frac) * min_log
+            model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                ModelVarType.FIXED_LARGE: (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                ModelVarType.FIXED_SMALL: (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if self.model_mean_type == ModelMeanType.START_X:
+            pred_xstart = process_xstart(model_output)
+        else:
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+            )
+        model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
+
+        assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+            "extra": extra,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, t, **model_kwargs)
+        new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+        See condition_mean() for details on cond_fn.
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t)
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        temperature=1.0
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param temperature: temperature scaling during Diff Loss sampling.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        # scale the noise by temperature
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise * temperature
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        temperature=1.0,
+    ):
+        """
+        Generate samples from the model.
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :param temperature: temperature scaling during Diff Loss sampling.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            temperature=temperature,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        temperature=1.0,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    temperature=temperature,
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+            self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
+    ):
+        """
+        Get a term for the variational lower-bound.
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(
+            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
+        )
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {"output": output, "pred_xstart": out["pred_xstart"]}
+
+    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
+        """
+        Compute training losses for a single timestep.
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            terms["loss"] = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )["output"]
+            if self.loss_type == LossType.RESCALED_KL:
+                terms["loss"] *= self.num_timesteps
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            model_output = model(x_t, t, **model_kwargs)
+
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                if len(model_output.shape) == len(x_t.shape) + 1:
+                    x_t = x_t.unsqueeze(-1).expand(*([-1] * (len(x_t.shape))), model_output.shape[-1])
+                    x_start = x_start.unsqueeze(-1).expand(*([-1] * (len(x_start.shape))), model_output.shape[-1])
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
+                    x_start=x_start, x_t=x_t, t=t
+                )[0],
+                ModelMeanType.START_X: x_start,
+                ModelMeanType.EPSILON: noise,
+            }[self.model_mean_type]
+            if len(model_output.shape) == len(target.shape) + 1:
+                target = target.unsqueeze(-1).expand(*([-1] * (len(target.shape))), model_output.shape[-1])
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+        This term can't be optimized, as it only depends on the encoder.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(
+            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+        )
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + th.zeros(broadcast_shape, device=timesteps.device)

semanticist/stage1/diffusion/respace.py

@@ -0,0 +1,130 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+import numpy as np
+import torch as th
+
+from .gaussian_diffusion import GaussianDiffusion
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+    :param use_timesteps: a collection (sequence or set) of timesteps from the
+                          original diffusion process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model):
+        if isinstance(model, _WrappedModel):
+            return model
+        return _WrappedModel(
+            model, self.timestep_map, self.original_num_steps
+        )
+
+    def _scale_timesteps(self, t):
+        # Scaling is done by the wrapped model.
+        return t
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        # self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        # if self.rescale_timesteps:
+        #     new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, new_ts, **kwargs)
+    

semanticist/stage1/diffusion/timestep_sampler.py

@@ -0,0 +1,150 @@
+# Modified from OpenAI's diffusion repos
+#     GLIDE: https://github.com/openai/glide-text2im/blob/main/glide_text2im/gaussian_diffusion.py
+#     ADM:   https://github.com/openai/guided-diffusion/blob/main/guided_diffusion
+#     IDDPM: https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch as th
+import torch.distributed as dist
+
+
+def create_named_schedule_sampler(name, diffusion):
+    """
+    Create a ScheduleSampler from a library of pre-defined samplers.
+    :param name: the name of the sampler.
+    :param diffusion: the diffusion object to sample for.
+    """
+    if name == "uniform":
+        return UniformSampler(diffusion)
+    elif name == "loss-second-moment":
+        return LossSecondMomentResampler(diffusion)
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {name}")
+
+
+class ScheduleSampler(ABC):
+    """
+    A distribution over timesteps in the diffusion process, intended to reduce
+    variance of the objective.
+    By default, samplers perform unbiased importance sampling, in which the
+    objective's mean is unchanged.
+    However, subclasses may override sample() to change how the resampled
+    terms are reweighted, allowing for actual changes in the objective.
+    """
+
+    @abstractmethod
+    def weights(self):
+        """
+        Get a numpy array of weights, one per diffusion step.
+        The weights needn't be normalized, but must be positive.
+        """
+
+    def sample(self, batch_size, device):
+        """
+        Importance-sample timesteps for a batch.
+        :param batch_size: the number of timesteps.
+        :param device: the torch device to save to.
+        :return: a tuple (timesteps, weights):
+                 - timesteps: a tensor of timestep indices.
+                 - weights: a tensor of weights to scale the resulting losses.
+        """
+        w = self.weights()
+        p = w / np.sum(w)
+        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)
+        indices = th.from_numpy(indices_np).long().to(device)
+        weights_np = 1 / (len(p) * p[indices_np])
+        weights = th.from_numpy(weights_np).float().to(device)
+        return indices, weights
+
+
+class UniformSampler(ScheduleSampler):
+    def __init__(self, diffusion):
+        self.diffusion = diffusion
+        self._weights = np.ones([diffusion.num_timesteps])
+
+    def weights(self):
+        return self._weights
+
+
+class LossAwareSampler(ScheduleSampler):
+    def update_with_local_losses(self, local_ts, local_losses):
+        """
+        Update the reweighting using losses from a model.
+        Call this method from each rank with a batch of timesteps and the
+        corresponding losses for each of those timesteps.
+        This method will perform synchronization to make sure all of the ranks
+        maintain the exact same reweighting.
+        :param local_ts: an integer Tensor of timesteps.
+        :param local_losses: a 1D Tensor of losses.
+        """
+        batch_sizes = [
+            th.tensor([0], dtype=th.int32, device=local_ts.device)
+            for _ in range(dist.get_world_size())
+        ]
+        dist.all_gather(
+            batch_sizes,
+            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
+        )
+
+        # Pad all_gather batches to be the maximum batch size.
+        batch_sizes = [x.item() for x in batch_sizes]
+        max_bs = max(batch_sizes)
+
+        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
+        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
+        dist.all_gather(timestep_batches, local_ts)
+        dist.all_gather(loss_batches, local_losses)
+        timesteps = [
+            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
+        ]
+        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
+        self.update_with_all_losses(timesteps, losses)
+
+    @abstractmethod
+    def update_with_all_losses(self, ts, losses):
+        """
+        Update the reweighting using losses from a model.
+        Sub-classes should override this method to update the reweighting
+        using losses from the model.
+        This method directly updates the reweighting without synchronizing
+        between workers. It is called by update_with_local_losses from all
+        ranks with identical arguments. Thus, it should have deterministic
+        behavior to maintain state across workers.
+        :param ts: a list of int timesteps.
+        :param losses: a list of float losses, one per timestep.
+        """
+
+
+class LossSecondMomentResampler(LossAwareSampler):
+    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
+        self.diffusion = diffusion
+        self.history_per_term = history_per_term
+        self.uniform_prob = uniform_prob
+        self._loss_history = np.zeros(
+            [diffusion.num_timesteps, history_per_term], dtype=np.float64
+        )
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+
+    def weights(self):
+        if not self._warmed_up():
+            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
+        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights /= np.sum(weights)
+        weights *= 1 - self.uniform_prob
+        weights += self.uniform_prob / len(weights)
+        return weights
+
+    def update_with_all_losses(self, ts, losses):
+        for t, loss in zip(ts, losses):
+            if self._loss_counts[t] == self.history_per_term:
+                # Shift out the oldest loss term.
+                self._loss_history[t, :-1] = self._loss_history[t, 1:]
+                self._loss_history[t, -1] = loss
+            else:
+                self._loss_history[t, self._loss_counts[t]] = loss
+                self._loss_counts[t] += 1
+
+    def _warmed_up(self):
+        return (self._loss_counts == self.history_per_term).all()

semanticist/stage1/diffusion_transfomer.py

@@ -0,0 +1,372 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# References:
+# GLIDE: https://github.com/openai/glide-text2im
+# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from timm.models.vision_transformer import PatchEmbed, Mlp
+from semanticist.stage1.fused_attention import Attention
+
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+#               Embedding Layers for Timesteps and Class Labels                 #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
+        return t_emb
+
+
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+#################################################################################
+#                                 Core DiT Model                                #
+#################################################################################
+
+class DiTBlock(nn.Module):
+    """
+    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        approx_gelu = lambda: nn.GELU(approximate="tanh")
+        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c, mask=None):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), mask)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    The final layer of DiT.
+    """
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class DiT(nn.Module):
+    """
+    Diffusion model with a Transformer backbone.
+    """
+    def __init__(
+        self,
+        input_size=32,
+        patch_size=2,
+        in_channels=4,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        class_dropout_prob=0.1,
+        num_classes=1000,
+        learn_sigma=True,
+    ):
+        super().__init__()
+        self.learn_sigma = learn_sigma
+        self.in_channels = in_channels
+        self.out_channels = in_channels * 2 if learn_sigma else in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+
+        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
+        num_patches = self.x_embedder.num_patches
+        # Will use fixed sin-cos embedding:
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+
+        self.blocks = nn.ModuleList([
+            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)
+        ])
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize (and freeze) pos_embed by sin-cos embedding:
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches ** 0.5))
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
+        w = self.x_embedder.proj.weight.data
+        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj.bias, 0)
+
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum('nhwpqc->nchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, h * p))
+        return imgs
+
+    def forward(self, x, t, y):
+        """
+        Forward pass of DiT.
+        x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+        t: (N,) tensor of diffusion timesteps
+        y: (N,) tensor of class labels
+        """
+        x = self.x_embedder(x) + self.pos_embed  # (N, T, D), where T = H * W / patch_size ** 2
+        t = self.t_embedder(t)                   # (N, D)
+        y = self.y_embedder(y, self.training)    # (N, D)
+        c = t + y                                # (N, D)
+        for block in self.blocks:
+            x = block(x, c)                      # (N, T, D)
+        x = self.final_layer(x, c)                # (N, T, patch_size ** 2 * out_channels)
+        x = self.unpatchify(x)                   # (N, out_channels, H, W)
+        return x
+
+    def forward_with_cfg(self, x, t, y, cfg_scale):
+        """
+        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/notebooks/text2im.ipynb
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, y)
+        # For exact reproducibility reasons, we apply classifier-free guidance on only
+        # three channels by default. The standard approach to cfg applies it to all channels.
+        # This can be done by uncommenting the following line and commenting-out the line following that.
+        eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        # eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+
+
+#################################################################################
+#                   Sine/Cosine Positional Embedding Functions                  #
+#################################################################################
+# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+#################################################################################
+#                                   DiT Configs                                  #
+#################################################################################
+
+def DiT_XL_2(**kwargs):
+    return DiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+def DiT_XL_4(**kwargs):
+    return DiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs)
+
+def DiT_XL_8(**kwargs):
+    return DiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs)
+
+def DiT_L_2(**kwargs):
+    return DiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs)
+
+def DiT_L_4(**kwargs):
+    return DiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs)
+
+def DiT_L_8(**kwargs):
+    return DiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs)
+
+def DiT_B_2(**kwargs):
+    return DiT(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs)
+
+def DiT_B_4(**kwargs):
+    return DiT(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs)
+
+def DiT_B_8(**kwargs):
+    return DiT(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs)
+
+def DiT_S_2(**kwargs):
+    return DiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs)
+
+def DiT_S_4(**kwargs):
+    return DiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs)
+
+def DiT_S_8(**kwargs):
+    return DiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs)
+
+
+DiT_models = {
+    'DiT-XL/2': DiT_XL_2,  'DiT-XL/4': DiT_XL_4,  'DiT-XL/8': DiT_XL_8,
+    'DiT-L/2':  DiT_L_2,   'DiT-L/4':  DiT_L_4,   'DiT-L/8':  DiT_L_8,
+    'DiT-B/2':  DiT_B_2,   'DiT-B/4':  DiT_B_4,   'DiT-B/8':  DiT_B_8,
+    'DiT-S/2':  DiT_S_2,   'DiT-S/4':  DiT_S_4,   'DiT-S/8':  DiT_S_8,
+}

semanticist/stage1/fused_attention.py

@@ -0,0 +1,45 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Type
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_bias: bool = True,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: Type[nn.Module] = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attn_mask=None):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_drop.p if self.training else 0.,
+        )
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

semanticist/stage1/pos_embed.py

@@ -0,0 +1,102 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+import numpy as np
+
+import torch
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_1d_sincos_pos_embed(embed_dim, grid_size):
+    grid = np.arange(grid_size, dtype=np.float32)
+    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=float)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# --------------------------------------------------------
+# Interpolate position embeddings for high-resolution
+# References:
+# DeiT: https://github.com/facebookresearch/deit
+# --------------------------------------------------------
+def interpolate_pos_embed(model, checkpoint_model):
+    if 'pos_embed' in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.patch_embed.num_patches
+        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            checkpoint_model['pos_embed'] = new_pos_embed

semanticist/stage1/transport/__init__.py

@@ -0,0 +1,63 @@
+from .transport import Transport, ModelType, WeightType, PathType, Sampler
+
+def create_transport(
+    path_type='Linear',
+    prediction="velocity",
+    loss_weight=None,
+    train_eps=None,
+    sample_eps=None,
+):
+    """function for creating Transport object
+    **Note**: model prediction defaults to velocity
+    Args:
+    - path_type: type of path to use; default to linear
+    - learn_score: set model prediction to score
+    - learn_noise: set model prediction to noise
+    - velocity_weighted: weight loss by velocity weight
+    - likelihood_weighted: weight loss by likelihood weight
+    - train_eps: small epsilon for avoiding instability during training
+    - sample_eps: small epsilon for avoiding instability during sampling
+    """
+
+    if prediction == "noise":
+        model_type = ModelType.NOISE
+    elif prediction == "score":
+        model_type = ModelType.SCORE
+    else:
+        model_type = ModelType.VELOCITY
+
+    if loss_weight == "velocity":
+        loss_type = WeightType.VELOCITY
+    elif loss_weight == "likelihood":
+        loss_type = WeightType.LIKELIHOOD
+    else:
+        loss_type = WeightType.NONE
+
+    path_choice = {
+        "Linear": PathType.LINEAR,
+        "GVP": PathType.GVP,
+        "VP": PathType.VP,
+    }
+
+    path_type = path_choice[path_type]
+
+    if (path_type in [PathType.VP]):
+        train_eps = 1e-5 if train_eps is None else train_eps
+        sample_eps = 1e-3 if train_eps is None else sample_eps
+    elif (path_type in [PathType.GVP, PathType.LINEAR] and model_type != ModelType.VELOCITY):
+        train_eps = 1e-3 if train_eps is None else train_eps
+        sample_eps = 1e-3 if train_eps is None else sample_eps
+    else: # velocity & [GVP, LINEAR] is stable everywhere
+        train_eps = 0
+        sample_eps = 0
+    
+    # create flow state
+    state = Transport(
+        model_type=model_type,
+        path_type=path_type,
+        loss_type=loss_type,
+        train_eps=train_eps,
+        sample_eps=sample_eps,
+    )
+    
+    return state

semanticist/stage1/transport/integrators.py

@@ -0,0 +1,130 @@
+import numpy as np
+import torch as th
+import torch.nn as nn
+from torchdiffeq import odeint
+from functools import partial
+from tqdm import tqdm
+
+class sde:
+    """SDE solver class"""
+    def __init__(
+        self, 
+        drift,
+        diffusion,
+        *,
+        t0,
+        t1,
+        num_steps,
+        sampler_type,
+        temperature=1.0,
+    ):
+        assert t0 < t1, "SDE sampler has to be in forward time"
+
+        self.num_timesteps = num_steps
+        self.t = th.linspace(t0, t1, num_steps)
+        self.dt = self.t[1] - self.t[0]
+        self.drift = drift
+        self.diffusion = diffusion
+        self.sampler_type = sampler_type
+        self.temperature = temperature
+
+    def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs):
+        w_cur = th.randn(x.size()).to(x)
+        t = th.ones(x.size(0)).to(x) * t
+        dw = w_cur * th.sqrt(self.dt)
+        drift = self.drift(x, t, model, **model_kwargs)
+        diffusion = self.diffusion(x, t)
+        mean_x = x + drift * self.dt
+        x = mean_x + th.sqrt(2 * diffusion) * dw * self.temperature
+        return x, mean_x
+    
+    def __Heun_step(self, x, _, t, model, **model_kwargs):
+        w_cur = th.randn(x.size()).to(x)
+        dw = w_cur * th.sqrt(self.dt) * self.temperature
+        t_cur = th.ones(x.size(0)).to(x) * t
+        diffusion = self.diffusion(x, t_cur)
+        xhat = x + th.sqrt(2 * diffusion) * dw
+        K1 = self.drift(xhat, t_cur, model, **model_kwargs)
+        xp = xhat + self.dt * K1
+        K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs)
+        return xhat + 0.5 * self.dt * (K1 + K2), xhat # at last time point we do not perform the heun step
+
+    def __forward_fn(self):
+        """TODO: generalize here by adding all private functions ending with steps to it"""
+        sampler_dict = {
+            "Euler": self.__Euler_Maruyama_step,
+            "Heun": self.__Heun_step,
+        }
+
+        try:
+            sampler = sampler_dict[self.sampler_type]
+        except:
+            raise NotImplementedError("Smapler type not implemented.")
+    
+        return sampler
+
+    def sample(self, init, model, **model_kwargs):
+        """forward loop of sde"""
+        x = init
+        mean_x = init 
+        samples = []
+        sampler = self.__forward_fn()
+        for ti in self.t[:-1]:
+            with th.no_grad():
+                x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs)
+                samples.append(x)
+
+        return samples
+
+class ode:
+    """ODE solver class"""
+    def __init__(
+        self,
+        drift,
+        *,
+        t0,
+        t1,
+        sampler_type,
+        num_steps,
+        atol,
+        rtol,
+        temperature=1.0,
+    ):
+        assert t0 < t1, "ODE sampler has to be in forward time"
+
+        self.drift = drift
+        self.t = th.linspace(t0, t1, num_steps)
+        self.atol = atol
+        self.rtol = rtol
+        self.sampler_type = sampler_type
+        self.temperature = temperature
+
+    def sample(self, x, model, **model_kwargs):
+        
+        device = x[0].device if isinstance(x, tuple) else x.device
+        def _fn(t, x):
+            t = th.ones(x[0].size(0)).to(device) * t if isinstance(x, tuple) else th.ones(x.size(0)).to(device) * t
+            # For ODE, we scale the drift by the temperature
+            # This is equivalent to scaling time by 1/temperature
+            model_output = self.drift(x, t, model, **model_kwargs)
+            if self.temperature != 1.0:
+                # If it's a tuple (for likelihood calculation), only scale the first element
+                if isinstance(model_output, tuple):
+                    scaled_output = (model_output[0] / self.temperature, model_output[1])
+                    return scaled_output
+                else:
+                    return model_output / self.temperature
+            return model_output
+
+        t = self.t.to(device)
+        atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol]
+        rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol]
+        samples = odeint(
+            _fn,
+            x,
+            t,
+            method=self.sampler_type,
+            atol=atol,
+            rtol=rtol
+        )
+        return samples

semanticist/stage1/transport/path.py

@@ -0,0 +1,192 @@
+import torch as th
+import numpy as np
+from functools import partial
+
+def expand_t_like_x(t, x):
+    """Function to reshape time t to broadcastable dimension of x
+    Args:
+      t: [batch_dim,], time vector
+      x: [batch_dim,...], data point
+    """
+    dims = [1] * (len(x.size()) - 1)
+    t = t.view(t.size(0), *dims)
+    return t
+
+
+#################### Coupling Plans ####################
+
+class ICPlan:
+    """Linear Coupling Plan"""
+    def __init__(self, sigma=0.0):
+        self.sigma = sigma
+
+    def compute_alpha_t(self, t):
+        """Compute the data coefficient along the path"""
+        return t, 1
+    
+    def compute_sigma_t(self, t):
+        """Compute the noise coefficient along the path"""
+        return 1 - t, -1
+    
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Compute the ratio between d_alpha and alpha"""
+        return 1 / t
+
+    def compute_drift(self, x, t):
+        """We always output sde according to score parametrization; """
+        t = expand_t_like_x(t, x)
+        alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        drift = alpha_ratio * x
+        diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
+
+        return -drift, diffusion
+
+    def compute_diffusion(self, x, t, form="constant", norm=1.0):
+        """Compute the diffusion term of the SDE
+        Args:
+          x: [batch_dim, ...], data point
+          t: [batch_dim,], time vector
+          form: str, form of the diffusion term
+          norm: float, norm of the diffusion term
+        """
+        t = expand_t_like_x(t, x)
+        choices = {
+            "constant": norm,
+            "SBDM": norm * self.compute_drift(x, t)[1],
+            "sigma": norm * self.compute_sigma_t(t)[0],
+            "linear": norm * (1 - t),
+            "decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2,
+            "inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2,
+        }
+
+        try:
+            diffusion = choices[form]
+        except KeyError:
+            raise NotImplementedError(f"Diffusion form {form} not implemented")
+        
+        return diffusion
+
+    def get_score_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to score
+        Args:
+            velocity: [batch_dim, ...] shaped tensor; velocity model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
+        score = (reverse_alpha_ratio * velocity - mean) / var
+        return score
+    
+    def get_noise_from_velocity(self, velocity, x, t):
+        """Wrapper function: transfrom velocity prediction model to denoiser
+        Args:
+            velocity: [batch_dim, ...] shaped tensor; velocity model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        alpha_t, d_alpha_t = self.compute_alpha_t(t)
+        sigma_t, d_sigma_t = self.compute_sigma_t(t)
+        mean = x
+        reverse_alpha_ratio = alpha_t / d_alpha_t
+        var = reverse_alpha_ratio * d_sigma_t - sigma_t
+        noise = (reverse_alpha_ratio * velocity - mean) / var
+        return noise
+
+    def get_velocity_from_score(self, score, x, t):
+        """Wrapper function: transfrom score prediction model to velocity
+        Args:
+            score: [batch_dim, ...] shaped tensor; score model output
+            x: [batch_dim, ...] shaped tensor; x_t data point
+            t: [batch_dim,] time tensor
+        """
+        t = expand_t_like_x(t, x)
+        drift, var = self.compute_drift(x, t)
+        velocity = var * score - drift
+        return velocity
+
+    def compute_mu_t(self, t, x0, x1):
+        """Compute the mean of time-dependent density p_t"""
+        t = expand_t_like_x(t, x1)
+        alpha_t, _ = self.compute_alpha_t(t)
+        sigma_t, _ = self.compute_sigma_t(t)
+        return alpha_t * x1 + sigma_t * x0
+    
+    def compute_xt(self, t, x0, x1):
+        """Sample xt from time-dependent density p_t; rng is required"""
+        xt = self.compute_mu_t(t, x0, x1)
+        return xt
+    
+    def compute_ut(self, t, x0, x1, xt):
+        """Compute the vector field corresponding to p_t"""
+        t = expand_t_like_x(t, x1)
+        _, d_alpha_t = self.compute_alpha_t(t)
+        _, d_sigma_t = self.compute_sigma_t(t)
+        return d_alpha_t * x1 + d_sigma_t * x0
+    
+    def plan(self, t, x0, x1):
+        xt = self.compute_xt(t, x0, x1)
+        ut = self.compute_ut(t, x0, x1, xt)
+        return t, xt, ut
+    
+
+class VPCPlan(ICPlan):
+    """class for VP path flow matching"""
+
+    def __init__(self, sigma_min=0.1, sigma_max=20.0):
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.log_mean_coeff = lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min) - 0.5 * (1 - t) * self.sigma_min 
+        self.d_log_mean_coeff = lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min) + 0.5 * self.sigma_min
+
+
+    def compute_alpha_t(self, t):
+        """Compute coefficient of x1"""
+        alpha_t = self.log_mean_coeff(t)
+        alpha_t = th.exp(alpha_t)
+        d_alpha_t = alpha_t * self.d_log_mean_coeff(t)
+        return alpha_t, d_alpha_t
+    
+    def compute_sigma_t(self, t):
+        """Compute coefficient of x0"""
+        p_sigma_t = 2 * self.log_mean_coeff(t)
+        sigma_t = th.sqrt(1 - th.exp(p_sigma_t))
+        d_sigma_t = th.exp(p_sigma_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * sigma_t)
+        return sigma_t, d_sigma_t
+    
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
+        return self.d_log_mean_coeff(t)
+
+    def compute_drift(self, x, t):
+        """Compute the drift term of the SDE"""
+        t = expand_t_like_x(t, x)
+        beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min)
+        return -0.5 * beta_t * x, beta_t / 2
+    
+
+class GVPCPlan(ICPlan):
+    def __init__(self, sigma=0.0):
+        super().__init__(sigma)
+    
+    def compute_alpha_t(self, t):
+        """Compute coefficient of x1"""
+        alpha_t = th.sin(t * np.pi / 2)
+        d_alpha_t = np.pi / 2 * th.cos(t * np.pi / 2)
+        return alpha_t, d_alpha_t
+    
+    def compute_sigma_t(self, t):
+        """Compute coefficient of x0"""
+        sigma_t = th.cos(t * np.pi / 2)
+        d_sigma_t = -np.pi / 2 * th.sin(t * np.pi / 2)
+        return sigma_t, d_sigma_t
+    
+    def compute_d_alpha_alpha_ratio_t(self, t):
+        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
+        return np.pi / (2 * th.tan(t * np.pi / 2))

semanticist/stage1/transport/transport.py

@@ -0,0 +1,456 @@
+import torch as th
+import numpy as np
+import logging
+
+import enum
+
+from . import path
+from .utils import EasyDict, log_state, mean_flat
+from .integrators import ode, sde
+
+class ModelType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    NOISE = enum.auto()  # the model predicts epsilon
+    SCORE = enum.auto()  # the model predicts \nabla \log p(x)
+    VELOCITY = enum.auto()  # the model predicts v(x)
+
+class PathType(enum.Enum):
+    """
+    Which type of path to use.
+    """
+
+    LINEAR = enum.auto()
+    GVP = enum.auto()
+    VP = enum.auto()
+
+class WeightType(enum.Enum):
+    """
+    Which type of weighting to use.
+    """
+
+    NONE = enum.auto()
+    VELOCITY = enum.auto()
+    LIKELIHOOD = enum.auto()
+
+
+class Transport:
+
+    def __init__(
+        self,
+        *,
+        model_type,
+        path_type,
+        loss_type,
+        train_eps,
+        sample_eps,
+    ):
+        path_options = {
+            PathType.LINEAR: path.ICPlan,
+            PathType.GVP: path.GVPCPlan,
+            PathType.VP: path.VPCPlan,
+        }
+
+        self.loss_type = loss_type
+        self.model_type = model_type
+        self.path_sampler = path_options[path_type]()
+        self.train_eps = train_eps
+        self.sample_eps = sample_eps
+
+    def prior_logp(self, z):
+        '''
+            Standard multivariate normal prior
+            Assume z is batched
+        '''
+        shape = th.tensor(z.size())
+        N = th.prod(shape[1:])
+        _fn = lambda x: -N / 2. * np.log(2 * np.pi) - th.sum(x ** 2) / 2.
+        return th.vmap(_fn)(z)
+    
+
+    def check_interval(
+        self, 
+        train_eps, 
+        sample_eps, 
+        *, 
+        diffusion_form="SBDM",
+        sde=False, 
+        reverse=False, 
+        eval=False,
+        last_step_size=0.0,
+    ):
+        t0 = 0
+        t1 = 1
+        eps = train_eps if not eval else sample_eps
+        if (type(self.path_sampler) in [path.VPCPlan]):
+
+            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
+
+        elif (type(self.path_sampler) in [path.ICPlan, path.GVPCPlan]) \
+            and (self.model_type != ModelType.VELOCITY or sde): # avoid numerical issue by taking a first semi-implicit step
+
+            t0 = eps if (diffusion_form == "SBDM" and sde) or self.model_type != ModelType.VELOCITY else 0
+            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
+        
+        if reverse:
+            t0, t1 = 1 - t0, 1 - t1
+
+        return t0, t1
+
+
+    def sample(self, x1):
+        """Sampling x0 & t based on shape of x1 (if needed)
+          Args:
+            x1 - data point; [batch, *dim]
+        """
+        
+        x0 = th.randn_like(x1)
+        t0, t1 = self.check_interval(self.train_eps, self.sample_eps)
+        t = th.rand((x1.shape[0],)) * (t1 - t0) + t0
+        t = t.to(x1)
+        return t, x0, x1
+    
+
+    def training_losses(
+        self, 
+        model,  
+        x1, 
+        model_kwargs=None
+    ):
+        """Loss for training the score model
+        Args:
+        - model: backbone model; could be score, noise, or velocity
+        - x1: datapoint
+        - model_kwargs: additional arguments for the model
+        """
+        if model_kwargs == None:
+            model_kwargs = {}
+        
+        t, x0, x1 = self.sample(x1)
+        t, xt, ut = self.path_sampler.plan(t, x0, x1)
+        model_output = model(xt, t, **model_kwargs)
+        if len(model_output.shape) == len(xt.shape) + 1:
+            x0 = x0.unsqueeze(-1).expand(*([-1] * (len(x0.shape))), model_output.shape[-1])
+            xt = xt.unsqueeze(-1).expand(*([-1] * (len(xt.shape))), model_output.shape[-1])
+            ut = ut.unsqueeze(-1).expand(*([-1] * (len(ut.shape))), model_output.shape[-1])
+        B, C = xt.shape[:2]
+        assert model_output.shape == (B, C, *xt.shape[2:])
+
+        terms = {}
+        terms['pred'] = model_output
+        if self.model_type == ModelType.VELOCITY:
+            terms['loss'] = mean_flat(((model_output - ut) ** 2))
+        else: 
+            _, drift_var = self.path_sampler.compute_drift(xt, t)
+            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, xt))
+            if self.loss_type in [WeightType.VELOCITY]:
+                weight = (drift_var / sigma_t) ** 2
+            elif self.loss_type in [WeightType.LIKELIHOOD]:
+                weight = drift_var / (sigma_t ** 2)
+            elif self.loss_type in [WeightType.NONE]:
+                weight = 1
+            else:
+                raise NotImplementedError()
+            
+            if self.model_type == ModelType.NOISE:
+                terms['loss'] = mean_flat(weight * ((model_output - x0) ** 2))
+            else:
+                terms['loss'] = mean_flat(weight * ((model_output * sigma_t + x0) ** 2))
+                
+        return terms
+    
+
+    def get_drift(
+        self
+    ):
+        """member function for obtaining the drift of the probability flow ODE"""
+        def score_ode(x, t, model, **model_kwargs):
+            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
+            model_output = model(x, t, **model_kwargs)
+            return (-drift_mean + drift_var * model_output) # by change of variable
+        
+        def noise_ode(x, t, model, **model_kwargs):
+            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
+            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))
+            model_output = model(x, t, **model_kwargs)
+            score = model_output / -sigma_t
+            return (-drift_mean + drift_var * score)
+        
+        def velocity_ode(x, t, model, **model_kwargs):
+            model_output = model(x, t, **model_kwargs)
+            return model_output
+
+        if self.model_type == ModelType.NOISE:
+            drift_fn = noise_ode
+        elif self.model_type == ModelType.SCORE:
+            drift_fn = score_ode
+        else:
+            drift_fn = velocity_ode
+        
+        def body_fn(x, t, model, **model_kwargs):
+            model_output = drift_fn(x, t, model, **model_kwargs)
+            assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape"
+            return model_output
+
+        return body_fn
+    
+
+    def get_score(
+        self,
+    ):
+        """member function for obtaining score of 
+            x_t = alpha_t * x + sigma_t * eps"""
+        if self.model_type == ModelType.NOISE:
+            score_fn = lambda x, t, model, **kwargs: model(x, t, **kwargs) / -self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))[0]
+        elif self.model_type == ModelType.SCORE:
+            score_fn = lambda x, t, model, **kwagrs: model(x, t, **kwagrs)
+        elif self.model_type == ModelType.VELOCITY:
+            score_fn = lambda x, t, model, **kwargs: self.path_sampler.get_score_from_velocity(model(x, t, **kwargs), x, t)
+        else:
+            raise NotImplementedError()
+        
+        return score_fn
+
+
+class Sampler:
+    """Sampler class for the transport model"""
+    def __init__(
+        self,
+        transport,
+    ):
+        """Constructor for a general sampler; supporting different sampling methods
+        Args:
+        - transport: an tranport object specify model prediction & interpolant type
+        """
+        
+        self.transport = transport
+        self.drift = self.transport.get_drift()
+        self.score = self.transport.get_score()
+    
+    def __get_sde_diffusion_and_drift(
+        self,
+        *,
+        diffusion_form="SBDM",
+        diffusion_norm=1.0,
+    ):
+
+        def diffusion_fn(x, t):
+            diffusion = self.transport.path_sampler.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm)
+            return diffusion
+        
+        sde_drift = \
+            lambda x, t, model, **kwargs: \
+                self.drift(x, t, model, **kwargs) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs)
+    
+        sde_diffusion = diffusion_fn
+
+        return sde_drift, sde_diffusion
+    
+    def __get_last_step(
+        self,
+        sde_drift,
+        *,
+        last_step,
+        last_step_size,
+    ):
+        """Get the last step function of the SDE solver"""
+    
+        if last_step is None:
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x
+        elif last_step == "Mean":
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x + sde_drift(x, t, model, **model_kwargs) * last_step_size
+        elif last_step == "Tweedie":
+            alpha = self.transport.path_sampler.compute_alpha_t # simple aliasing; the original name was too long
+            sigma = self.transport.path_sampler.compute_sigma_t
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x / alpha(t)[0][0] + (sigma(t)[0][0] ** 2) / alpha(t)[0][0] * self.score(x, t, model, **model_kwargs)
+        elif last_step == "Euler":
+            last_step_fn = \
+                lambda x, t, model, **model_kwargs: \
+                    x + self.drift(x, t, model, **model_kwargs) * last_step_size
+        else:
+            raise NotImplementedError()
+
+        return last_step_fn
+
+    def sample_sde(
+        self,
+        *,
+        sampling_method="Euler",
+        diffusion_form="SBDM",
+        diffusion_norm=1.0,
+        last_step="Mean",
+        last_step_size=0.04,
+        num_steps=250,
+        temperature=1.0,
+    ):
+        """returns a sampling function with given SDE settings
+        Args:
+        - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama
+        - diffusion_form: function form of diffusion coefficient; default to be matching SBDM
+        - diffusion_norm: function magnitude of diffusion coefficient; default to 1
+        - last_step: type of the last step; default to identity
+        - last_step_size: size of the last step; default to match the stride of 250 steps over [0,1]
+        - num_steps: total integration step of SDE
+        - temperature: temperature scaling for the noise during sampling; default to 1.0
+        """
+
+        if last_step is None:
+            last_step_size = 0.0
+
+        sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift(
+            diffusion_form=diffusion_form,
+            diffusion_norm=diffusion_norm,
+        )
+
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            diffusion_form=diffusion_form,
+            sde=True,
+            eval=True,
+            reverse=False,
+            last_step_size=last_step_size,
+        )
+
+        _sde = sde(
+            sde_drift,
+            sde_diffusion,
+            t0=t0,
+            t1=t1,
+            num_steps=num_steps,
+            sampler_type=sampling_method,
+            temperature=temperature
+        )
+
+        last_step_fn = self.__get_last_step(sde_drift, last_step=last_step, last_step_size=last_step_size)
+            
+
+        def _sample(init, model, **model_kwargs):
+            xs = _sde.sample(init, model, **model_kwargs)
+            ts = th.ones(init.size(0), device=init.device) * t1
+            x = last_step_fn(xs[-1], ts, model, **model_kwargs)
+            xs.append(x)
+
+            assert len(xs) == num_steps, "Samples does not match the number of steps"
+
+            return xs
+
+        return _sample
+    
+    def sample_ode(
+        self,
+        *,
+        sampling_method="dopri5",
+        num_steps=50,
+        atol=1e-6,
+        rtol=1e-3,
+        reverse=False,
+        temperature=1.0,
+    ):
+        """returns a sampling function with given ODE settings
+        Args:
+        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
+        - num_steps: 
+            - fixed solver (Euler, Heun): the actual number of integration steps performed
+            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
+        - atol: absolute error tolerance for the solver
+        - rtol: relative error tolerance for the solver
+        - reverse: whether solving the ODE in reverse (data to noise); default to False
+        - temperature: temperature scaling for the drift during sampling; default to 1.0
+        """
+        if reverse:
+            drift = lambda x, t, model, **kwargs: self.drift(x, th.ones_like(t) * (1 - t), model, **kwargs)
+        else:
+            drift = self.drift
+
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            sde=False,
+            eval=True,
+            reverse=reverse,
+            last_step_size=0.0,
+        )
+
+        _ode = ode(
+            drift=drift,
+            t0=t0,
+            t1=t1,
+            sampler_type=sampling_method,
+            num_steps=num_steps,
+            atol=atol,
+            rtol=rtol,
+            temperature=temperature,
+        )
+        
+        return _ode.sample
+
+    def sample_ode_likelihood(
+        self,
+        *,
+        sampling_method="dopri5",
+        num_steps=50,
+        atol=1e-6,
+        rtol=1e-3,
+        temperature=1.0,
+    ):
+        
+        """returns a sampling function for calculating likelihood with given ODE settings
+        Args:
+        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
+        - num_steps: 
+            - fixed solver (Euler, Heun): the actual number of integration steps performed
+            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
+        - atol: absolute error tolerance for the solver
+        - rtol: relative error tolerance for the solver
+        - temperature: temperature scaling for the drift during sampling; default to 1.0
+        """
+        def _likelihood_drift(x, t, model, **model_kwargs):
+            x, _ = x
+            eps = th.randint(2, x.size(), dtype=th.float, device=x.device) * 2 - 1
+            t = th.ones_like(t) * (1 - t)
+            with th.enable_grad():
+                x.requires_grad = True
+                grad = th.autograd.grad(th.sum(self.drift(x, t, model, **model_kwargs) * eps), x)[0]
+                logp_grad = th.sum(grad * eps, dim=tuple(range(1, len(x.size()))))
+                drift = self.drift(x, t, model, **model_kwargs)
+            return (-drift, logp_grad)
+        
+        t0, t1 = self.transport.check_interval(
+            self.transport.train_eps,
+            self.transport.sample_eps,
+            sde=False,
+            eval=True,
+            reverse=False,
+            last_step_size=0.0,
+        )
+
+        _ode = ode(
+            drift=_likelihood_drift,
+            t0=t0,
+            t1=t1,
+            sampler_type=sampling_method,
+            num_steps=num_steps,
+            atol=atol,
+            rtol=rtol,
+            temperature=temperature,
+        )
+
+        def _sample_fn(x, model, **model_kwargs):
+            init_logp = th.zeros(x.size(0)).to(x)
+            input = (x, init_logp)
+            drift, delta_logp = _ode.sample(input, model, **model_kwargs)
+            drift, delta_logp = drift[-1], delta_logp[-1]
+            prior_logp = self.transport.prior_logp(drift)
+            logp = prior_logp - delta_logp
+            return logp, drift
+
+        return _sample_fn

semanticist/stage1/transport/utils.py

@@ -0,0 +1,29 @@
+import torch as th
+
+class EasyDict:
+
+    def __init__(self, sub_dict):
+        for k, v in sub_dict.items():
+            setattr(self, k, v)
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+def mean_flat(x):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return th.mean(x, dim=list(range(1, len(x.size()))))
+
+def log_state(state):
+    result = []
+    
+    sorted_state = dict(sorted(state.items()))
+    for key, value in sorted_state.items():
+        # Check if the value is an instance of a class
+        if "<object" in str(value) or "object at" in str(value):
+            result.append(f"{key}: [{value.__class__.__name__}]")
+        else:
+            result.append(f"{key}: {value}")
+    
+    return '\n'.join(result)

semanticist/stage1/vision_transformer.py

@@ -0,0 +1,259 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Mostly copy-paste from timm library.
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+import math
+import torch
+import torch.nn as nn
+
+from functools import partial
+from semanticist.stage1.fused_attention import Attention
+
+__all__ = ['VisionTransformer', 'vit_tiny_patch16', 'vit_small_patch16',
+           'vit_base_patch16', 'vit_large_patch16', 'vit_huge_patch14']
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0.,
+                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, init_values=0):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, attn_mask=None):
+        y = self.attn(self.norm1(x), attn_mask=attn_mask)
+        if self.gamma_1 is None:
+            x = x + self.drop_path(y)
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * y)
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        return self.proj(x)
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+
+    def __init__(self, img_size=[224], patch_size=16, in_chans=3, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=partial(nn.LayerNorm, eps=1e-6),
+                 init_values=0, num_slots=16):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.num_slots = num_slots
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1 + self.num_slots, embed_dim))
+        self.slot_embed = nn.Parameter(torch.zeros(1, num_slots, embed_dim))
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values)
+            for i in range(depth)])
+        
+        self.norm = norm_layer(embed_dim)
+
+        nn.init.trunc_normal_(self.pos_embed, std=.02)
+        nn.init.trunc_normal_(self.cls_token, std=.02)
+        nn.init.trunc_normal_(self.slot_embed, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1 - self.num_slots
+        N = self.pos_embed.shape[1] - 1 - self.num_slots
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:1+npatch]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size[0]
+        h0 = h // self.patch_embed.patch_size[1]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+        slots_pos_embed = self.pos_embed[:, 1+npatch:]
+        slots_pos_embed = slots_pos_embed.view(1, -1, dim) # (1, num_slots, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed, slots_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)
+        x = x.flatten(2).transpose(1, 2)
+        x = torch.cat((self.cls_token.expand(B, -1, -1), x, self.slot_embed.expand(B, -1, -1)), dim=1)
+        x = x + self.interpolate_pos_encoding(x, w, h)
+        return self.pos_drop(x)
+
+    def forward(self, x, is_causal=True):
+        x = self.prepare_tokens(x)
+        if is_causal:
+            attn_mask = torch.ones(x.shape[1], x.shape[1], device=x.device, dtype=torch.bool)
+            # slots are causal to each other
+            causal_mask = torch.ones(self.num_slots, self.num_slots, device=x.device, dtype=torch.bool).tril(diagonal=0)
+            attn_mask[-self.num_slots:, -self.num_slots:] = causal_mask
+            # cls token and patches should not see slots
+            attn_mask[:-self.num_slots, -self.num_slots:] = False
+        else:
+            attn_mask = None
+
+        for blk in self.blocks:
+            x = blk(x, attn_mask=attn_mask)
+
+        x = self.norm(x)
+        outcome = x[:, -self.num_slots:] # return the slots
+        return outcome
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+
+def vit_tiny_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_small_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_base_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_large_patch16(**kwargs):
+    model = VisionTransformer(
+        patch_size=16, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_huge_patch14(**kwargs):
+    model = VisionTransformer(
+        patch_size=14, embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model

semanticist/stage2/diffloss.py

@@ -0,0 +1,267 @@
+import math
+import torch
+import torch.nn as nn
+
+from semanticist.stage1.diffusion import create_diffusion
+from semanticist.stage1.transport import create_transport, Sampler
+
+
+class DiffLoss(nn.Module):
+    """Diffusion Loss"""
+    def __init__(self, target_channels, z_channels, depth, width, num_sampling_steps, predict_xstart=False, use_si=False, cond_method="adaln"):
+        super(DiffLoss, self).__init__()
+        self.in_channels = target_channels
+        self.net = SimpleMLPAdaLN(
+            in_channels=target_channels,
+            model_channels=width,
+            out_channels=target_channels * 2 if not use_si else target_channels,  # for vlb loss
+            z_channels=z_channels,
+            num_res_blocks=depth,
+            cond_method=cond_method,
+        )
+        self.use_si = use_si
+        if not use_si:
+            self.train_diffusion = create_diffusion(timestep_respacing="", noise_schedule="cosine", predict_xstart=predict_xstart)
+            self.gen_diffusion = create_diffusion(timestep_respacing=num_sampling_steps, noise_schedule="cosine", predict_xstart=predict_xstart)
+        else:
+            self.transport = create_transport()
+            self.sampler = Sampler(self.transport)
+
+    def forward(self, target, z, mask=None):
+        model_kwargs = dict(c=z)
+        if not self.use_si:
+            t = torch.randint(0, self.train_diffusion.num_timesteps, (target.shape[0],), device=target.device)
+            loss_dict = self.train_diffusion.training_losses(self.net, target, t, model_kwargs)
+        else:
+            loss_dict = self.transport.training_losses(self.net, target, model_kwargs)
+        loss = loss_dict["loss"]
+        if mask is not None:
+            loss = (loss * mask).sum() / mask.sum()
+        return loss.mean()
+
+    def sample(self, z, temperature=1.0, cfg=1.0):
+        # diffusion loss sampling
+        device = z.device
+        if not cfg == 1.0:
+            noise = torch.randn(z.shape[0] // 2, self.in_channels, device=device)
+            noise = torch.cat([noise, noise], dim=0)
+            model_kwargs = dict(c=z, cfg_scale=cfg)
+            sample_fn = self.net.forward_with_cfg
+        else:
+            noise = torch.randn(z.shape[0], self.in_channels, device=device)
+            model_kwargs = dict(c=z)
+            sample_fn = self.net.forward
+
+        if not self.use_si:
+            sampled_token_latent = self.gen_diffusion.p_sample_loop(
+                sample_fn, noise.shape, noise, clip_denoised=False, model_kwargs=model_kwargs, progress=False,
+                temperature=temperature, device=device
+            )
+        else:
+            sde_sample_fn = self.sampler.sample_sde(diffusion_form="sigma", temperature=temperature)
+            sampled_token_latent = sde_sample_fn(noise, sample_fn, **model_kwargs)[-1]
+        if cfg != 1.0:
+            sampled_token_latent, _ = sampled_token_latent.chunk(2, dim=0)
+        return sampled_token_latent
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale) + shift
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class ResBlock(nn.Module):
+    """
+    A residual block with AdaLN for timestep and optional concatenation for condition.
+    """
+    def __init__(
+        self,
+        channels,
+        cond_method="adaln",
+    ):
+        super().__init__()
+        self.channels = channels
+        self.cond_method = cond_method
+
+        self.in_ln = nn.LayerNorm(channels, eps=1e-6)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(channels, 3 * channels, bias=True)
+        )
+        
+        # Input dimension depends on conditioning method
+        mlp_in_dim = channels * 2 if cond_method == "concat" else channels
+        self.mlp = nn.Sequential(
+            nn.Linear(mlp_in_dim, channels, bias=True),
+            nn.SiLU(),
+            nn.Linear(channels, channels, bias=True),
+        )
+
+    def forward(self, x, t, c=None):
+        # Apply timestep embedding via AdaLN
+        shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(3, dim=-1)
+        h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
+        
+        # Concatenate condition if using concat method
+        if self.cond_method == "concat" and c is not None:
+            h = torch.cat([h, c], dim=-1)
+        
+        h = self.mlp(h)
+        x = x + gate_mlp * h
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    Final layer with AdaLN for timestep and optional concatenation for condition.
+    """
+    def __init__(self, model_channels, out_channels, cond_method="adaln"):
+        super().__init__()
+        self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
+        self.cond_method = cond_method
+        
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(model_channels, 2 * model_channels, bias=True)
+        )
+        
+        # Output dimension depends on conditioning method
+        linear_in_dim = model_channels * 2 if cond_method == "concat" else model_channels
+        self.linear = nn.Linear(linear_in_dim, out_channels, bias=True)
+
+    def forward(self, x, t, c=None):
+        # Apply timestep embedding via AdaLN
+        shift, scale = self.adaLN_modulation(t).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        
+        # Concatenate condition if using concat method
+        if self.cond_method == "concat" and c is not None:
+            x = torch.cat([x, c], dim=-1)
+            
+        return self.linear(x)
+
+
+class SimpleMLPAdaLN(nn.Module):
+    """
+    MLP for Diffusion Loss with AdaLN for timestep and optional concatenation for condition.
+    """
+    def __init__(
+        self,
+        in_channels,
+        model_channels,
+        out_channels,
+        z_channels,
+        num_res_blocks,
+        cond_method="adaln"
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.cond_method = cond_method
+
+        self.time_embed = TimestepEmbedder(model_channels)
+        self.cond_embed = nn.Linear(z_channels, model_channels)
+        self.input_proj = nn.Linear(in_channels, model_channels)
+
+        # Create residual blocks
+        res_blocks = [ResBlock(model_channels, cond_method) for _ in range(num_res_blocks)]
+        self.res_blocks = nn.ModuleList(res_blocks)
+        
+        self.final_layer = FinalLayer(model_channels, out_channels, cond_method=cond_method)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Basic initialization for all linear layers
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP
+        nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers (always used for timestep)
+        for i, block in enumerate(self.res_blocks):
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def forward(self, x, t, c):
+        """
+        Apply the model to an input batch.
+        :param x: an [N x C] Tensor of inputs.
+        :param t: a 1-D batch of timesteps.
+        :param c: conditioning from AR transformer.
+        :return: an [N x C] Tensor of outputs.
+        """
+        x = self.input_proj(x)
+        t_emb = self.time_embed(t)
+        c_emb = self.cond_embed(c)
+
+        # Prepare conditioning based on method
+        if self.cond_method == "adaln":
+            t_combined, c_for_concat = t_emb + c_emb, None
+        else:  # concat
+            t_combined, c_for_concat = t_emb, c_emb
+
+        for block in self.res_blocks:
+            x = block(x, t_combined, c_for_concat)
+        return self.final_layer(x, t_combined, c_for_concat)
+
+    def forward_with_cfg(self, x, t, c, cfg_scale):
+        half = x[: len(x) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self.forward(combined, t, c)
+        eps, rest = model_out[:, :self.in_channels], model_out[:, self.in_channels:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)

semanticist/stage2/generate.py

@@ -0,0 +1,88 @@
+# Modified from:
+#   gpt-fast: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
+#   DiT:      https://github.com/facebookresearch/DiT/blob/main/models.py
+import torch
+
+def prefill(model, cond_idx: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, temperature: float = 1.0):
+    tokens = model(None, cond_idx, input_pos, cfg=cfg_scale, temperature=temperature)
+    return tokens.unsqueeze(1)
+
+
+def decode_one_token(model, x: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, temperature: float = 1.0):
+    assert input_pos.shape[-1] == 1
+    if cfg_scale > 1.0:
+        x = torch.cat([x, x])
+    tokens = model(x, cond_idx=None, input_pos=input_pos, cfg=cfg_scale, temperature=temperature)
+    return tokens
+
+
+def decode_n_tokens(
+    model, cur_token: torch.Tensor, input_pos: torch.Tensor, num_new_tokens: int, 
+    cfg_scale: float, cfg_schedule = "constant", temperature: float = 1.0):
+    new_tokens = []
+    for i in range(num_new_tokens):
+        cfg_iter = get_cfg(cfg_scale, i + 1, num_new_tokens + 1, cfg_schedule)
+        next_token = decode_one_token(model, cur_token, input_pos, cfg_iter, temperature=temperature).unsqueeze(1)
+        input_pos += 1
+        new_tokens.append(next_token.clone())
+        cur_token = next_token
+    
+    return new_tokens
+
+
+@torch.no_grad()
+def generate(model, cond, max_new_tokens, emb_masks=None, cfg_scale=1.0, cfg_schedule = "constant", temperature: float = 1.0):
+    if cfg_scale > 1.0:
+        cond_null = torch.ones_like(cond) * model.num_classes
+        cond_combined = torch.cat([cond, cond_null])
+    else:
+        cond_combined = cond
+    T = model.cls_token_num
+
+    T_new = T + max_new_tokens
+    max_seq_length = T_new
+    max_batch_size = cond.shape[0]
+
+    device = cond.device
+    dtype = model.z_proj.weight.dtype
+    if torch.is_autocast_enabled():
+        dtype = torch.get_autocast_dtype(device_type=device.type)
+    with torch.device(device):
+        max_batch_size_cfg = max_batch_size * 2 if cfg_scale > 1.0 else max_batch_size
+        model.setup_caches(max_batch_size=max_batch_size_cfg, max_seq_length=max_seq_length, dtype=dtype)
+    
+    if emb_masks is not None:
+        assert emb_masks.shape[0] == max_batch_size
+        assert emb_masks.shape[-1] == T
+        if cfg_scale > 1.0:
+            model.causal_mask[:, :, :T] = model.causal_mask[:, :, :T] * torch.cat([emb_masks, emb_masks]).unsqueeze(1)
+        else:
+            model.causal_mask[:, :, :T] = model.causal_mask[:, :, :T] * emb_masks.unsqueeze(1)
+
+        eye_matrix = torch.eye(model.causal_mask.size(1), model.causal_mask.size(2), device=device)
+        model.causal_mask[:] = model.causal_mask * (1 - eye_matrix) + eye_matrix
+
+    # create an empty tensor of the expected final shape and fill in the current tokens
+    seq = torch.empty((max_batch_size, T_new, model.slot_dim), dtype=dtype, device=device)
+
+    input_pos = torch.arange(0, T, device=device)
+    cfg_iter = get_cfg(cfg_scale, 0, max_new_tokens, cfg_schedule)
+    next_token = prefill(model, cond_combined, input_pos, cfg_iter, temperature=temperature)
+    seq[:, T:T+1] = next_token
+    
+    if max_new_tokens > 1:
+        input_pos = torch.tensor([T], device=device, dtype=torch.int)
+        generated_tokens = decode_n_tokens(model, next_token, input_pos, max_new_tokens - 1, cfg_scale, cfg_schedule=cfg_schedule, temperature=temperature)
+        seq[:, T+1:] = torch.cat(generated_tokens, dim=1)
+
+    model.reset_caches()
+    return seq[:, T:]
+
+
+def get_cfg(cfg, cur_step, total_step, cfg_schedule="constant"):
+    if cfg_schedule == "linear":
+        return 1 + (cfg - 1) * (cur_step + 1) / total_step
+    elif cfg_schedule == "constant":
+        return cfg
+    else:
+        raise NotImplementedError

semanticist/stage2/gpt.py

@@ -0,0 +1,431 @@
+# Modified from:
+#   VQGAN:    https://github.com/CompVis/taming-transformers/blob/master/taming/modules/transformer/mingpt.py
+#   DiT:      https://github.com/facebookresearch/DiT/blob/main/models.py  
+#   nanoGPT:  https://github.com/karpathy/nanoGPT/blob/master/model.py
+#   llama:    https://github.com/facebookresearch/llama/blob/main/llama/model.py
+#   gpt-fast: https://github.com/pytorch-labs/gpt-fast/blob/main/model.py
+#   PixArt:   https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
+from typing import Optional, List, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from semanticist.stage1.vision_transformer import DropPath
+from semanticist.stage2.diffloss import DiffLoss
+
+def find_multiple(n: int, k: int):
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+
+#################################################################################
+#                      Embedding Layers for Class Labels                        #
+#################################################################################
+class LabelEmbedder(nn.Module):
+    """
+    Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+    """
+    def __init__(self, num_classes, hidden_size, dropout_prob):
+        super().__init__()
+        use_cfg_embedding = dropout_prob > 0
+        self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+        self.num_classes = num_classes
+        self.dropout_prob = dropout_prob
+
+    def token_drop(self, labels, force_drop_ids=None):
+        """
+        Drops labels to enable classifier-free guidance.
+        """
+        if force_drop_ids is None:
+            drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+        else:
+            drop_ids = force_drop_ids == 1
+        labels = torch.where(drop_ids, self.num_classes, labels)
+        return labels
+
+    def forward(self, labels, train, force_drop_ids=None):
+        use_dropout = self.dropout_prob > 0
+        if (train and use_dropout) or (force_drop_ids is not None):
+            labels = self.token_drop(labels, force_drop_ids)
+        embeddings = self.embedding_table(labels).unsqueeze(1)
+        return embeddings
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features, out_features):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=False)
+        self.act = nn.GELU(approximate='tanh')
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=False)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+
+#################################################################################
+#                                  GPT Model                                    #
+#################################################################################
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        multiple_of: int = 256,
+        ffn_dropout_p: float = 0.0,
+    ):
+        super().__init__()
+        hidden_dim = 4 * dim
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = find_multiple(hidden_dim, multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.ffn_dropout = nn.Dropout(ffn_dropout_p)
+
+    def forward(self, x):
+        return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+
+class KVCache(nn.Module):
+    def __init__(self, max_batch_size, max_seq_length, n_head, head_dim, dtype):
+        super().__init__()
+        cache_shape = (max_batch_size, n_head, max_seq_length, head_dim)
+        self.register_buffer('k_cache', torch.zeros(cache_shape, dtype=dtype))
+        self.register_buffer('v_cache', torch.zeros(cache_shape, dtype=dtype))
+
+    def update(self, input_pos, k_val, v_val):
+        # input_pos: [S], k_val: [B, H, S, D]
+        assert input_pos.shape[0] == k_val.shape[2]
+        k_out = self.k_cache
+        v_out = self.v_cache
+        k_out[:, :, input_pos] = k_val
+        v_out[:, :, input_pos] = v_val
+
+        return k_out, v_out
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_head: int,
+        attn_dropout_p: float = 0.0,
+        resid_dropout_p: float = 0.1,
+    ):
+        super().__init__()
+        assert dim % n_head == 0
+        self.dim = dim
+        self.head_dim = dim // n_head
+        self.n_head = n_head
+
+        # key, query, value projections for all heads, but in a batch
+        self.wqkv = nn.Linear(dim, dim * 3, bias=False)
+        self.wo = nn.Linear(dim, dim, bias=False)
+        self.kv_cache = None
+
+        # regularization
+        self.attn_dropout_p = attn_dropout_p
+        self.resid_dropout = nn.Dropout(resid_dropout_p)
+
+    def forward(
+        self, x: torch.Tensor, 
+        input_pos: Optional[torch.Tensor] = None, 
+        mask: Optional[torch.Tensor] = None
+    ):
+        bsz, seqlen, _ = x.shape
+        xq, xk, xv = self.wqkv(x).split([self.dim, self.dim, self.dim], dim=-1)
+
+        xq = xq.view(bsz, seqlen, self.n_head, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_head, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_head, self.head_dim)
+
+        xq, xk, xv = map(lambda x: x.transpose(1, 2), (xq, xk, xv))
+
+        if self.kv_cache is not None:
+            keys, values = self.kv_cache.update(input_pos, xk, xv)
+        else:
+            keys, values = xk, xv
+
+        output = F.scaled_dot_product_attention(
+            xq, keys, values, 
+            attn_mask=mask, 
+            is_causal=True if mask is None else False, # is_causal=False is for KV cache
+            dropout_p=self.attn_dropout_p if self.training else 0)            
+        
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, self.dim)
+
+        output = self.resid_dropout(self.wo(output))
+        return output
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_head: int,
+        multiple_of: int = 256,
+        norm_eps: float = 1e-5,
+        attn_dropout_p: float = 0.0,
+        ffn_dropout_p: float = 0.1,
+        resid_dropout_p: float = 0.1,
+        drop_path: float = 0.0,
+    ):
+        super().__init__()
+        self.attention = Attention(
+            dim=dim,
+            n_head=n_head,
+            attn_dropout_p=attn_dropout_p,
+            resid_dropout_p=resid_dropout_p,
+        )
+        self.feed_forward = FeedForward(
+            dim=dim,
+            multiple_of=multiple_of,
+            ffn_dropout_p=ffn_dropout_p,
+        )
+        self.attention_norm = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm = RMSNorm(dim, eps=norm_eps)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor, start_pos: int, mask: Optional[torch.Tensor] = None):
+        h = x + self.drop_path(self.attention(self.attention_norm(x), start_pos, mask))
+        out = h + self.drop_path(self.feed_forward(self.ffn_norm(h)))
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        dim: int = 4096,
+        n_layer: int = 32,
+        n_head: int = 32,
+        attn_dropout_p: float = 0.0,
+        resid_dropout_p: float = 0.1,
+        ffn_dropout_p: float = 0.1,
+        drop_path_rate: float = 0.0,
+        num_classes: Union[int, List[int]] = 1000,
+        class_dropout_prob: float = 0.1,
+
+        cls_token_num: int = 1,
+        num_slots: int = 16,
+        slot_dim: int = 256,
+
+        diffloss_d: int = 3,
+        diffloss_w: int = 1024,
+        num_sampling_steps: str = '100',
+        diffusion_batch_mul: int = 4,
+        predict_xstart: bool = False,
+        use_si: bool = False,
+        cond_method: str = "adaln",
+        **kwargs,
+    ):
+        super().__init__()
+        
+        # Store configuration
+        self.dim = dim
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.num_slots = num_slots
+        self.slot_dim = slot_dim
+        self.num_classes = num_classes
+        self.cls_token_num = cls_token_num
+        
+        # Initialize embeddings
+        self.cls_embedding = LabelEmbedder(num_classes, dim, class_dropout_prob)
+        self.z_proj = nn.Linear(slot_dim, dim, bias=True)
+        self.z_proj_ln = RMSNorm(dim)
+        self.pos_embed_learned = nn.Parameter(torch.zeros(1, num_slots + cls_token_num, dim))
+
+        # transformer blocks
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, n_layer)]
+        self.layers = torch.nn.ModuleList()
+        for layer_id in range(n_layer):
+            self.layers.append(TransformerBlock(
+                dim=dim,
+                n_head=n_head,
+                ffn_dropout_p=ffn_dropout_p,
+                attn_dropout_p=attn_dropout_p,
+                resid_dropout_p=resid_dropout_p,
+                drop_path=dpr[layer_id],
+            ))
+
+        # output layer
+        self.norm = RMSNorm(dim)
+
+        self.diffusion_pos_embed_learned = nn.Parameter(torch.zeros(1, num_slots, dim))
+
+        # KVCache
+        self.max_batch_size = -1
+        self.max_seq_length = -1
+
+        self.initialize_weights()
+
+        # Diffusion Loss
+        self.diffloss = DiffLoss(
+            target_channels=slot_dim,
+            z_channels=dim,
+            width=diffloss_w,
+            depth=diffloss_d,
+            num_sampling_steps=num_sampling_steps,
+            predict_xstart=predict_xstart,
+            use_si=use_si,
+            cond_method=cond_method,
+        )
+        self.diffusion_batch_mul = diffusion_batch_mul
+
+    def initialize_weights(self):
+        nn.init.normal_(self.pos_embed_learned, std=0.02)
+        nn.init.normal_(self.diffusion_pos_embed_learned, std=0.02)
+        # Initialize nn.Linear and nn.Embedding
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(std=0.02)
+
+    def setup_caches(self, max_batch_size, max_seq_length, dtype):
+        # if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size:
+        #     return
+        head_dim = self.dim // self.n_head
+        max_seq_length = find_multiple(max_seq_length, 8)
+        self.max_seq_length = max_seq_length
+        self.max_batch_size = max_batch_size
+        for b in self.layers:
+            b.attention.kv_cache = KVCache(max_batch_size, max_seq_length, self.n_head, head_dim, dtype)
+
+        causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool))
+        self.causal_mask = causal_mask.unsqueeze(0).repeat(self.max_batch_size, 1, 1)
+
+    def reset_caches(self):
+        self.max_seq_length = -1
+        self.max_batch_size = -1
+        for b in self.layers:
+            b.attention.kv_cache = None
+
+    def forward_loss(self, z, target):
+        bsz, seq_len, _ = target.shape
+        target = target.reshape(bsz * seq_len, -1).repeat(self.diffusion_batch_mul, 1)
+        z = z.reshape(bsz*seq_len, -1).repeat(self.diffusion_batch_mul, 1)
+        loss = self.diffloss(z=z, target=target)
+        return loss
+    
+    def forward_cfg(self, h, cfg):
+        if cfg > 1.0:
+            h_cond, h_uncond = h.chunk(2, dim=0)
+            h = h_uncond + cfg * (h_cond - h_uncond)
+        return h
+
+    def forward(
+        self,
+        slots: torch.Tensor,
+        cond_idx: torch.Tensor,
+        input_pos:  Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        cfg: float = 1.0,
+        temperature: float = 1.0
+    ):
+        if slots is not None and cond_idx is not None: # training or naive inference
+            cond_embeddings = self.cls_embedding(cond_idx, train=self.training)
+            cond_embeddings = cond_embeddings.expand(-1, self.cls_token_num, -1)
+            token_embeddings = self.z_proj(slots)
+            token_embeddings = torch.cat((cond_embeddings, token_embeddings), dim=1)
+        else:
+            if cond_idx is not None: # prefill in inference
+                token_embeddings = self.cls_embedding(cond_idx, train=self.training)
+                token_embeddings = token_embeddings.expand(-1, self.cls_token_num, -1)
+            else: # decode_n_tokens(kv cache) in inference
+                token_embeddings = self.z_proj(slots)
+            
+            bs = token_embeddings.shape[0]
+            mask = self.causal_mask[:bs, None, input_pos]
+        
+        h = token_embeddings
+        if self.training:
+            h = h + self.pos_embed_learned
+        else:
+            h = h + self.pos_embed_learned[:, input_pos].view(1, -1, self.dim)
+        
+        h = self.z_proj_ln(h) # not sure if this is needed
+
+        # transformer blocks
+        for layer in self.layers:
+            h = layer(h, input_pos, mask)
+        
+        h = self.norm(h)
+        
+        if self.training:
+            h = h[:, self.cls_token_num - 1 : -1].contiguous()
+            h = h + self.diffusion_pos_embed_learned
+            loss = self.forward_loss(h, slots.detach())
+            return loss
+        else:
+            h = h[:, -1]
+            h = h + self.diffusion_pos_embed_learned[:, input_pos[-1] - self.cls_token_num + 1]
+            next_tokens = self.diffloss.sample(h, temperature=temperature, cfg=cfg)
+            return next_tokens
+
+
+    def get_fsdp_wrap_module_list(self) -> List[nn.Module]:
+        return list(self.layers)
+
+
+
+#################################################################################
+#                                GPT Configs                                    #
+#################################################################################
+### text-conditional
+def GPT_7B(**kwargs):
+    return Transformer(n_layer=32, n_head=32, dim=4096, **kwargs) # 6.6B
+
+def GPT_3B(**kwargs):
+    return Transformer(n_layer=24, n_head=32, dim=3200, **kwargs) # 3.1B
+
+def GPT_1B(**kwargs):
+    return Transformer(n_layer=22, n_head=32, dim=2048, **kwargs) # 1.2B
+
+### class-conditional
+def GPT_XXXL(**kwargs):
+    return Transformer(n_layer=48, n_head=40, dim=2560, **kwargs) # 3.9B
+
+def GPT_XXL(**kwargs):
+    return Transformer(n_layer=48, n_head=24, dim=1536, **kwargs) # 1.4B
+
+def GPT_XL(**kwargs):
+    return Transformer(n_layer=36, n_head=20, dim=1280, **kwargs) # 775M
+
+def GPT_L(**kwargs):
+    return Transformer(n_layer=24, n_head=16, dim=1024, **kwargs) # 343M
+
+def GPT_B(**kwargs):
+    return Transformer(n_layer=12, n_head=12, dim=768, **kwargs) # 111M
+        
+
+GPT_models = {
+    'GPT-B': GPT_B, 'GPT-L': GPT_L, 'GPT-XL': GPT_XL, 'GPT-XXL': GPT_XXL, 'GPT-XXXL': GPT_XXXL,
+    'GPT-1B': GPT_1B, 'GPT-3B': GPT_3B, 'GPT-7B': GPT_7B, 
+}

semanticist/utils/datasets.py

@@ -0,0 +1,72 @@
+import torch
+import torchvision
+import numpy as np
+import os.path as osp
+from PIL import Image
+import torchvision
+import torchvision.transforms as TF
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def center_crop_arr(pil_image, image_size):
+    """
+    Center cropping implementation from ADM.
+    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
+    """
+    while min(*pil_image.size) >= 2 * image_size:
+        pil_image = pil_image.resize(
+            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
+        )
+
+    scale = image_size / min(*pil_image.size)
+    pil_image = pil_image.resize(
+        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
+    )
+
+    arr = np.array(pil_image)
+    crop_y = (arr.shape[0] - image_size) // 2
+    crop_x = (arr.shape[1] - image_size) // 2
+    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+
+def vae_transforms(split, aug='randcrop', img_size=256):
+    t = []
+    if split == 'train':
+        if aug == 'randcrop':
+            t.append(TF.Resize(img_size, interpolation=TF.InterpolationMode.BICUBIC, antialias=True))
+            t.append(TF.RandomCrop(img_size))
+        elif aug == 'centercrop':
+            t.append(TF.Lambda(lambda x: center_crop_arr(x, img_size)))
+        else:
+            raise ValueError(f"Invalid augmentation: {aug}")
+        t.append(TF.RandomHorizontalFlip(p=0.5))
+    else:
+        t.append(TF.Lambda(lambda x: center_crop_arr(x, img_size)))
+        
+    t.append(TF.ToTensor())
+
+    return TF.Compose(t)
+
+
+def cached_transforms(aug='tencrop', img_size=256, crop_ranges=[1.05, 1.10]):
+    t = []
+    if 'centercrop' in aug:
+        t.append(TF.Lambda(lambda x: center_crop_arr(x, img_size)))
+        t.append(TF.Lambda(lambda x: torch.stack([TF.ToTensor()(x), TF.ToTensor()(TF.functional.hflip(x))])))
+    elif 'tencrop' in aug:
+        crop_sizes = [int(img_size * crop_range) for crop_range in crop_ranges]
+        t.append(TF.Lambda(lambda x: [center_crop_arr(x, crop_size) for crop_size in crop_sizes]))
+        t.append(TF.Lambda(lambda crops: [crop for crop_tuple in [TF.TenCrop(img_size)(crop) for crop in crops] for crop in crop_tuple]))
+        t.append(TF.Lambda(lambda crops: torch.stack([TF.ToTensor()(crop) for crop in crops])))
+    else:
+        raise ValueError(f"Invalid augmentation: {aug}")
+
+    return TF.Compose(t)
+
+class ImageNet(torchvision.datasets.ImageFolder):
+    def __init__(self, root, split='train', aug='randcrop', img_size=256):
+        super().__init__(osp.join(root, split))
+        if not 'cache' in aug:
+            self.transform = vae_transforms(split, aug=aug, img_size=img_size)
+        else:
+            self.transform = cached_transforms(aug=aug, img_size=img_size)

semanticist/utils/device_utils.py

@@ -0,0 +1,18 @@
+import torch
+
+def configure_compute_backend():
+    """Configure PyTorch compute backend settings for CUDA."""
+    if torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True 
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+    else:
+        raise ValueError("No CUDA available")
+
+def get_device():
+    """Get the device to use for training."""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    else:
+        raise ValueError("No CUDA available")

semanticist/utils/logger.py

@@ -0,0 +1,170 @@
+from collections import defaultdict, deque
+import datetime
+import time
+import torch
+import torch.distributed as dist
+from semanticist.engine.trainer_utils import is_dist_avail_and_initialized, is_main_process
+from semanticist.utils.device_utils import get_device
+
+def synchronize_processes():
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    else: # do nothing
+        pass
+
+def empty_cache():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    else: # do nothing
+        pass
+
+class SmoothedValue(object):
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float32, device=get_device())
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=self.max,
+            value=self.value)
+
+
+class MetricLogger(object):
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, attr))
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(
+                "{}: {}".format(name, str(meter))
+            )
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        if torch.cuda.is_available():
+            log_msg.append('mem: {memory:.0f}')
+            log_msg.append("util: {util:.1f}%")
+        log_msg = self.delimiter.join(log_msg)
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    if is_main_process():
+                        memory = torch.cuda.max_memory_allocated()
+                        util = torch.cuda.utilization()
+                        print(log_msg.format(
+                            i, len(iterable), eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time), data=str(data_time),
+                            memory=memory / MB, util=util))
+                else:
+                    if is_main_process():
+                        print(log_msg.format(
+                            i, len(iterable), eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        if is_main_process():
+            print('{} Total time: {} ({:.4f} s / it)'.format(
+                header, total_time_str, total_time / len(iterable)))

semanticist/utils/lr_scheduler.py

@@ -0,0 +1,15 @@
+from timm.scheduler.cosine_lr import CosineLRScheduler
+from timm.scheduler.step_lr import StepLRScheduler
+
+def build_scheduler(optimizer, n_epoch, n_iter_per_epoch, lr_min=0, warmup_steps=0, warmup_lr_init=0, decay_steps=None, cosine_lr=True):
+    if decay_steps is None:
+        decay_steps = n_epoch * n_iter_per_epoch
+    
+    if cosine_lr:
+        scheduler = CosineLRScheduler(optimizer, t_initial=decay_steps, lr_min=lr_min, warmup_t=warmup_steps, warmup_lr_init=warmup_lr_init, 
+                                      cycle_limit=1, t_in_epochs=False, warmup_prefix=True)
+    else:
+        scheduler = StepLRScheduler(optimizer, decay_t=decay_steps, warmup_t=warmup_steps, warmup_lr_init=warmup_lr_init, 
+                                    t_in_epochs=False, warmup_prefix=True)
+    
+    return scheduler

semanticist/utils/transform.py

@@ -0,0 +1,35 @@
+import PIL
+import torchvision.transforms as T
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+def stage1_transform(img_size=256, is_train=True, scale=0.8):
+    
+    resize = pair(int(img_size/scale))
+    t = []
+    t.append(T.Resize(resize, interpolation=PIL.Image.BICUBIC))
+    if is_train:
+        t.append(T.RandomCrop(img_size))
+        t.append(T.RandomHorizontalFlip(p=0.5))
+    else:
+        t.append(T.CenterCrop(img_size))
+        
+    t.append(T.ToTensor())
+    t.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))),
+    
+    return T.Compose(t)
+        
+def stage2_transform(img_size=256, is_train=True, scale=0.8):
+    resize = pair(int(img_size/scale))
+    t = []
+    t.append(T.Resize(resize, interpolation=PIL.Image.BICUBIC))
+    if is_train:
+        t.append(T.RandomCrop(img_size))
+    else:
+        t.append(T.CenterCrop(img_size))
+        
+    t.append(T.ToTensor())
+    t.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))),
+    
+    return T.Compose(t)