upload

Dockerfile

@@ -0,0 +1,58 @@
+FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04
+LABEL maintainer="Bingchen Zhao"
+LABEL repository="Semanticist"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get -y update \
+    && apt-get install -y software-properties-common \
+    && add-apt-repository ppa:deadsnakes/ppa
+
+RUN apt install -y bash \
+    build-essential \
+    git \
+    git-lfs \
+    curl \
+    ca-certificates \
+    libsndfile1-dev \
+    libgl1 \
+    python3.10 \
+    python3.10-dev \
+    python3-pip \
+    python3.10-venv rsync sudo tmux && \
+    rm -rf /var/lib/apt/lists
+
+# make sure to use venv
+RUN python3.10 -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# pre-install the heavy dependencies (these can later be overridden by the deps from setup.py)
+RUN python3.10 -m pip install --no-cache-dir --upgrade pip uv==0.1.11 && \
+    python3.10 -m uv pip install --no-cache-dir \
+    torch \
+    torchvision \
+    torchaudio \
+    invisible_watermark && \
+    python3.10 -m pip install --no-cache-dir \
+    accelerate \
+    datasets \
+    hf-doc-builder \
+    huggingface-hub \
+    hf_transfer \
+    Jinja2 \
+    librosa \
+    numpy==1.26.4 \
+    scipy \
+    tensorboard \
+    transformers \
+    pytorch-lightning matplotlib \
+    hf_transfer
+
+# start Semanticist part
+COPY . /work/Semanticist
+WORKDIR /work/Semanticist
+RUN ls && python3.10 -m pip install -r req_min.txt && \
+    python3.10 -m pip install git+https://github.com/cocodataset/panopticapi.git
+
+CMD ["/bin/bash"]
+# docker run -it --rm --runtime=nvidia --gpus all xx/xx:xx

README.md

@@ -5,7 +5,7 @@ colorFrom: gray
 colorTo: red
 sdk: gradio
 sdk_version: 5.20.1
-app_file: app.py
+app_file: demo.py
 pinned: false
 license: mit
 ---

configs/autoregressive_config.yaml

@@ -0,0 +1,77 @@
+trainer:
+  target: paintmind.engine.gpt_trainer.GPTTrainer
+  params:
+    num_epoch: 400
+    blr: 1e-4
+    cosine_lr: False
+    warmup_epochs: 100
+    batch_size: 32
+    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.5
+    cfg: 1.5
+    cfg_schedule: "constant"
+    train_num_slots: 32
+    test_num_slots: 32
+    compile: True
+    enable_cache_latents: True
+    ae_model:
+      target: paintmind.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_mask_type: 'replace'
+        cond_method: 'token'
+        dit_model: 'DiT-XL-2'
+        vae: 'xwen99/mar-vae-kl16'
+        num_sampling_steps: '250'
+        ckpt_path: ./output/tokenizer/models/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
+        token_drop_prob: 0
+        use_si: True
+        cond_method: 'concat'
+        decoupled_cfg: False
+        ckpt_path: None
+
+    dataset:
+      target: paintmind.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: train
+        aug:  tencrop_cached
+        img_size: 256
+
+    test_dataset:
+      target: paintmind.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: val
+        img_size: 256

configs/onenode_config.yaml

@@ -0,0 +1,12 @@
+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
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+use_cpu: false

configs/tokenizer_config.yaml

@@ -0,0 +1,64 @@
+trainer:
+  target: paintmind.engine.trainer.DiffusionTrainer
+  params:
+    num_epoch: 400
+    valid_size: 64
+    blr: 2.5e-5
+    cosine_lr: True
+    warmup_epochs: 100
+    batch_size: 32
+    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"
+    log_dit: "./output/tokenizer/logs"
+    cfg: 3.0
+    compile: True
+    model:
+      target: paintmind.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_mask_type: 'replace'
+        cond_method: 'token'
+        dit_model: 'DiT-XL-2'
+        vae: 'xwen99/mar-vae-kl16'
+        enable_nest: False
+        enable_nest_after: 50
+        nest_rho: 0.03
+        nest_dist: uniform
+        nest_null_prob: 0
+        nest_allow_zero: False
+        use_repa: True
+        repa_encoder: dinov2_vitb
+        repa_encoder_depth: 8
+        repa_loss_weight: 1.0
+        eval_fid: True
+        fid_stats: 'fid_stats/adm_in256_stats.npz'
+        num_sampling_steps: '250'
+        ckpt_path: None
+
+    dataset:
+      target: paintmind.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: train
+        img_size: 256
+
+    test_dataset:
+      target: paintmind.utils.datasets.ImageNet
+      params:
+        root: ./dataset/imagenet/
+        split: val
+        img_size: 256

demo.py

@@ -0,0 +1,170 @@
+import gradio as gr
+import numpy as np
+from PIL import Image
+import os
+import torch
+import matplotlib.pyplot as plt
+from omegaconf import OmegaConf
+from huggingface_hub import hf_hub_download
+
+from paintmind.engine.util import instantiate_from_config
+from paintmind.stage1.diffuse_slot import DiffuseSlot
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+ckpt_path = hf_hub_download(repo_id='tennant/semanticist', filename='semanticist_tok_XL.pkl')
+config_path = 'configs/tokenizer_config.yaml'
+cfg = OmegaConf.load(config_path)
+ckpt = torch.load(ckpt_path, map_location='cpu')
+from paintmind.utils.datasets import vae_transforms
+from PIL import Image
+
+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
+
+ckpt = {k.replace('._orig_mod', ''): v for k, v in ckpt.items()}
+
+model = DiffuseSlot(**cfg['trainer']['params']['model']['params'])
+msg = model.load_state_dict(ckpt, strict=False)
+model = model.to(device)
+model = model.eval()
+model.enable_nest = True
+
+def viz_diff_slots(model, img, nums, cfg=1.0, return_img=False):
+    n_slots_inf = []
+    for num_slots_to_inference in nums:
+        recon_n = model(
+            img, None, sample=True, cfg=cfg,
+            inference_with_n_slots=num_slots_to_inference,
+        )
+        n_slots_inf.append(recon_n)
+    return [convert_np(n_slots_inf[i][0]) for i in range(len(n_slots_inf))]
+
+# Removed process_image function as its functionality is now in the update_outputs function
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        # First column - Input and configs
+        with gr.Column(scale=1):
+            gr.Markdown("## Input")
+            input_image = gr.Image(label="Upload an image", type="numpy")
+            
+            with gr.Group():
+                gr.Markdown("### Configuration")
+                show_gallery = gr.Checkbox(label="Show Gallery", value=False)
+                # You can add more config options here
+                # slider = gr.Slider(minimum=0, maximum=10, value=5, label="Processing Intensity")
+                slider = gr.Slider(minimum=0.1, maximum=20.0, value=4.0, label="CFG value")
+                labels_input = gr.Textbox(
+                    label="Gallery Labels (comma-separated)", 
+                    value="1, 4, 16, 64, 256", 
+                    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=False) 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="Processed Image", type="numpy")
+    
+    # Handle form submission
+    submit_btn = gr.Button("Process")
+    
+    # Define the processing logic
+    def update_outputs(image, show_gallery_value, slider_value, labels_text):
+        # 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)
+        
+        # Process the image based on configurations
+        if image is None:
+            # Return placeholder if no image is uploaded
+            placeholder = np.zeros((300, 300, 3), dtype=np.uint8)
+            return gallery_container, [], placeholder
+        image = Image.fromarray(image)
+        img = transform(image)
+        img = img.unsqueeze(0).to(device)
+        recon = viz_diff_slots(model, img, [256], cfg=slider_value)[0]
+
+            
+        if not show_gallery_value:
+            # If only the image should be shown, return just the processed image
+            return gallery_container, [], recon
+        else:
+            model_decompose = viz_diff_slots(model, img, labels, cfg=slider_value)
+            # Create image variations and pair them with labels
+            gallery_images = [
+                (image, 'GT'),
+                # (np.array(Image.fromarray(image).convert("L").convert("RGB")), labels[1]),
+                # (np.array(Image.fromarray(image).rotate(180)), labels[2])
+            ] + [(img, 'Recon. with ' + str(label) + ' tokens') for img, label in zip(model_decompose, labels)]
+            return gallery_container, gallery_images, image
+    
+    # Connect the inputs and outputs
+    submit_btn.click(
+        fn=update_outputs,
+        inputs=[input_image, 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 = [
+        ["examples/city.jpg", False, 4.0, "1,4,16,64,256"],
+        ["examples/food.jpg", True, 4.0, "1,4,16,64,256"],
+        ["examples/highland.webp", True, 4.0, "1,4,16,64,256"],
+    ]
+    
+    gr.Examples(
+        examples=examples,
+        inputs=[input_image, show_gallery, slider, labels_input],
+        outputs=[gallery_container, gallery, output_image],
+        fn=update_outputs,
+        cache_examples=True
+    )
+
+# Launch the demo
+if __name__ == "__main__":
+    demo.launch()

fid_stats/adm_in256_stats.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e842c756177806d210e891893914620bee2cf5b779a5613ba5af5145d7c85289
+size 33563124

gen_demo.py

@@ -0,0 +1,137 @@
+import gradio as gr
+import numpy as np
+import cv2
+import time
+from PIL import Image
+import io
+
+def apply_transformations(image, numbers_str):
+    """
+    Apply a series of transformations to an image based on a list of numbers.
+    Shows the progressive changes as each transformation is applied.
+    Returns both the current image and the full gallery of transformations.
+    """
+    try:
+        # Parse the input numbers
+        numbers = [float(n.strip()) for n in numbers_str.split(',') if n.strip()]
+        if not numbers:
+            return image, [(image, "Original Image")]
+        
+        # Convert PIL Image to numpy array for OpenCV operations
+        img = np.array(image)
+        
+        # Initialize the result list with the original image
+        results = [(image, "Original Image")]
+        current_image = image
+        
+        # Apply transformations based on each number
+        for i, value in enumerate(numbers):
+            # Make a copy of the current numpy image
+            if i == 0:
+                current_img = img.copy()
+            else:
+                current_img = np.array(current_image)
+            
+            # Apply different transformations based on the value
+            transformation_type = ""
+            if i % 5 == 0:  # Brightness adjustment
+                # Scale value to reasonable brightness adjustment
+                brightness = max(min(value, 100), -100)  # Limit between -100 and 100
+                current_img = cv2.addWeighted(current_img, 1, np.zeros_like(current_img), 0, brightness)
+                transformation_type = f"Brightness: {brightness:.1f}"
+            
+            elif i % 5 == 1:  # Rotation
+                # Scale value to reasonable rotation angle
+                angle = value % 360
+                h, w = current_img.shape[:2]
+                center = (w // 2, h // 2)
+                rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
+                current_img = cv2.warpAffine(current_img, rotation_matrix, (w, h))
+                transformation_type = f"Rotation: {angle:.1f}°"
+            
+            elif i % 5 == 2:  # Contrast adjustment
+                # Scale value to reasonable contrast adjustment
+                contrast = max(min(value / 10, 3), 0.5)  # Limit between 0.5 and 3
+                current_img = cv2.convertScaleAbs(current_img, alpha=contrast, beta=0)
+                transformation_type = f"Contrast: {contrast:.1f}x"
+            
+            elif i % 5 == 3:  # Blur
+                # Scale value to reasonable blur kernel size
+                blur_amount = max(int(abs(value) % 20), 1)
+                if blur_amount % 2 == 0:  # Ensure kernel size is odd
+                    blur_amount += 1
+                current_img = cv2.GaussianBlur(current_img, (blur_amount, blur_amount), 0)
+                transformation_type = f"Blur: {blur_amount}px"
+            
+            elif i % 5 == 4:  # Hue shift (for color images)
+                if current_img.shape[-1] == 3:  # Only for color images
+                    # Convert to HSV
+                    hsv_img = cv2.cvtColor(current_img, cv2.COLOR_RGB2HSV)
+                    # Shift hue
+                    hue_shift = int(value) % 180
+                    hsv_img[:, :, 0] = (hsv_img[:, :, 0] + hue_shift) % 180
+                    # Convert back to RGB
+                    current_img = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2RGB)
+                    transformation_type = f"Hue Shift: {hue_shift}°"
+            
+            # Convert back to PIL Image and add to results
+            current_image = Image.fromarray(current_img)
+            
+            # Add to results with a label for the gallery
+            results.append((current_image, f"Step {i+1}: {transformation_type}"))
+            
+            # (Progress updates removed)
+            
+            # Add a small delay to make the progressive changes visible
+            time.sleep(4)
+            
+            # Yield intermediate results for real-time updates
+            if i < len(numbers) - 1:
+                yield current_image, results
+        
+        return current_image, results
+    
+    except Exception as e:
+        error_msg = f"Error: {str(e)}"
+        return image, [(image, "Error")]
+
+# Create Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Image Transformation Demo")
+    gr.Markdown("Upload an image and provide a comma-separated list of numbers. The demo will apply a series of transformations to the image based on these numbers.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(label="Input Image", type="pil")
+            numbers_input = gr.Textbox(label="Transformation Values (comma-separated numbers)", 
+                                      placeholder="e.g., 50, -30, 1.5, 5, 90, 20",
+                                      value="30, 45, 1.5, 3, 60, -20, 90, 1.8, 7, 120")
+            transform_btn = gr.Button("Apply Transformations")
+            
+            explanation = gr.Markdown("""
+            ## How the transformations work:
+            
+            The numbers you input will be used to apply these transformations in sequence:
+            1. First number: Brightness adjustment (-100 to 100)
+            2. Second number: Rotation (degrees)
+            3. Third number: Contrast adjustment (0.5 to 3)
+            4. Fourth number: Blur (kernel size)
+            5. Fifth number: Hue shift (color images only)
+            
+            And the pattern repeats for longer lists of numbers.
+            """)
+            
+        with gr.Column(scale=2):
+            with gr.Row():
+                current_image = gr.Image(label="Current Transformation", type="pil")
+            with gr.Row():
+                gallery = gr.Gallery(label="Transformation History", show_label=True, columns=4, rows=2, height="auto")
+    
+    transform_btn.click(
+        fn=apply_transformations,
+        inputs=[input_image, numbers_input],
+        outputs=[current_image, gallery]
+    )
+
+# Launch the app
+demo.launch()

making_cache.py

@@ -0,0 +1,46 @@
+import numpy as np
+import os, pdb, time
+import torch_fidelity
+import tqdm
+import torch
+import os.path as osp
+import argparse
+from omegaconf import OmegaConf
+from paintmind.engine.util import instantiate_from_config
+
+
+@torch.no_grad()
+def caching():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cfg', type=str, default='configs/vit_vqgan.yaml')
+    args = parser.parse_args()
+
+    cfg_file = args.cfg
+    assert osp.exists(cfg_file)
+    config = OmegaConf.load(cfg_file)
+    dataset = instantiate_from_config(config.trainer.params.dataset)
+    model = instantiate_from_config(config.trainer.params.model)
+    dataloader = torch.utils.data.DataLoader(
+        dataset, 
+        batch_size=config.trainer.params.batch_size, 
+        shuffle=False, 
+        num_workers=config.trainer.params.num_workers,
+    )
+    # Each batch will give us a (N, C, H, W) tensor of images
+    # We need to cache them and save them to a pth file
+    cache_save_file = config.trainer.params.latent_cache_file
+    cache = []
+    # import ipdb; ipdb.set_trace()
+    model.cuda()
+    model.eval()
+    for idx, batch in enumerate(tqdm.tqdm(dataloader)):
+        batch = batch[0].cuda()
+        latent = model.vae_encode(batch)
+        cache.append(latent.cpu())
+    cache = torch.cat(cache, dim=0)
+    torch.save(cache, cache_save_file)
+
+if __name__ == '__main__':
+
+    caching()
+

paintmind/__init__.py

@@ -0,0 +1,2 @@
+from .config import Config
+from .version import __version__

paintmind/config.py

@@ -0,0 +1,110 @@
+import json, pdb
+from copy import deepcopy
+
+class Config:
+    def __init__(self, config=None):
+        if config is not None:
+            self.from_dict(config)
+    
+    def __repr__(self):
+        return str(self.to_json_string())
+    
+    def to_dict(self):
+        return deepcopy(self.__dict__)
+    
+    def to_json(self, path):
+        with open(path, 'w') as f:
+            json.dump(self.to_dict(), f, indent=2)
+            
+    def to_json_string(self):
+        return json.dumps(self.to_dict(), indent=2)
+            
+    def from_dict(self, dct):
+        self.clear()
+        for key, value in dct.items():
+            self.__dict__[key] = value
+            
+        return self.to_dict()
+    
+    def from_json(self, json_path):
+        with open(json_path, 'r') as f:
+            config = json.load(f)
+            self.from_dict(config)
+            
+        return self.to_dict()
+    
+    def clear(self):
+        del self.__dict__
+        
+
+vit_s_vqgan_config = {
+    'n_embed'     :8192,
+    'embed_dim'   :16,
+    'beta'        :0.25,
+    'enc':{
+        'image_size':320, 
+        'patch_size':8, 
+        'dim':512, 
+        'depth':8, 
+        'num_head':8, 
+        'mlp_dim':2048, 
+        'in_channels':3, 
+        'dim_head':64, 
+        'dropout':0.0,
+    }, 
+    'dec':{
+        'image_size':320, 
+        'patch_size':8, 
+        'dim':512, 
+        'depth':8, 
+        'num_head':8, 
+        'mlp_dim':2048, 
+        'out_channels':3, 
+        'dim_head':64, 
+        'dropout':0.0,
+    },     
+}
+vit_m_vqgan_config = {
+    'n_embed'     :8192,
+    'embed_dim'   :32,
+    'beta'        :0.25,
+    'enc':{
+        'image_size':256, 
+        'patch_size':8, 
+        'dim': 1024, 
+        'depth': 16, 
+        'num_head':16, 
+        'mlp_dim':2048, 
+        'in_channels':3, 
+        'dim_head':64, 
+        'dropout':0.0,
+    }, 
+    'dec':{
+        'image_size':256, 
+        'patch_size':8, 
+        'dim':1024, 
+        'depth':16, 
+        'num_head':16, 
+        'mlp_dim':2048, 
+        'out_channels':3, 
+        'dim_head':64, 
+        'dropout':0.0,
+    },     
+}
+
+pipeline_v1_config = {
+    'stage1'         :'vit-s-vqgan',
+    't5'             :'t5-l',
+    'dim'            :1024, 
+    'dim_head'       :64,
+    'mlp_dim'        :4096,
+    'num_head'       :16, 
+    'depth'          :12,
+    'dropout'        :0.1, 
+}
+
+ver2cfg = {
+    'vit-s-vqgan'  : vit_s_vqgan_config,
+    'vit-m-vqgan'  : vit_m_vqgan_config,
+    'paintmindv1'  : pipeline_v1_config,
+}

paintmind/engine/gpt_trainer.py

@@ -0,0 +1,892 @@
+import os, torch
+import os.path as osp
+import cv2
+import shutil
+import numpy as np
+import copy
+import torch_fidelity
+import torch.nn as nn
+from tqdm.auto import tqdm
+from collections import OrderedDict
+from einops import rearrange
+from accelerate import Accelerator
+from .util import instantiate_from_config
+from torchvision.utils import make_grid, save_image
+from torch.utils.data import DataLoader, random_split, DistributedSampler, Sampler
+from paintmind.utils.lr_scheduler import build_scheduler
+from paintmind.utils.logger import SmoothedValue, MetricLogger, synchronize_processes, empty_cache
+from paintmind.engine.misc import is_main_process, all_reduce_mean, concat_all_gather
+from accelerate.utils import DistributedDataParallelKwargs, AutocastKwargs
+from torch.optim import AdamW
+from concurrent.futures import ThreadPoolExecutor
+from paintmind.stage2.gpt import GPT_models
+from paintmind.stage2.causaldit import CausalDiT_models
+from paintmind.stage2.generate import generate, generate_causal_dit
+from pathlib import Path
+import time
+
+
+def requires_grad(model, flag=True):
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def save_img(img, save_path):
+    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
+    
+    # Use ProcessPoolExecutor which is generally better for CPU-bound tasks
+    # ThreadPoolExecutor is better for I/O-bound tasks like file saving
+    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):
+    stats = torch_fidelity.calculate_metrics(
+        input1=real_dir,
+        input2=rec_dir,
+        fid_statistics_file=fid_stats,
+        cuda=True,
+        isc=True,
+        fid=True,
+        kid=False,
+        prc=False,
+        verbose=False,
+    )
+    return stats
+
+
+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 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])
+            )
+
+class GPTTrainer(nn.Module):
+    def __init__(
+        self,
+        ae_model,
+        gpt_model,
+        dataset,
+        test_dataset=None,
+        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
+        lr=None,
+        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=0,
+        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",
+        steps=0,
+        cfg=1.75,
+        ae_cfg=1.5,
+        diff_cfg=2.0,
+        temperature=1.0,
+        cfg_schedule="constant",
+        diff_cfg_schedule="inv_linear",
+        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',
+        seed=42
+    ):
+        super().__init__()
+        kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+        self.accelerator = Accelerator(
+            kwargs_handlers=[kwargs],
+            mixed_precision="bf16",
+            gradient_accumulation_steps=grad_accum_steps,
+            log_with="tensorboard",
+            project_dir=log_dir,
+        )
+
+        self.ae_model = instantiate_from_config(ae_model)
+        if hasattr(ae_model.params, "ema_path") and ae_model.params.ema_path is not None:
+            ae_model_path = ae_model.params.ema_path
+        else:
+            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)
+        elif 'CausalDiT' in gpt_model.target:
+            self.gpt_model = CausalDiT_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
+
+        assert precision in ["bf16", "fp32"]
+        precision = "fp32"
+        if self.accelerator.is_main_process:
+            print("Overlooking specified precision and using autocast bf16...")
+        self.precision = precision
+
+        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
+            if lr is None:
+                lr = blr * effective_bs / 256
+            if self.accelerator.is_main_process:
+                print(f"Effective batch size is {effective_bs}")
+
+            self.g_optim = self._creat_optimizer(weight_decay=0.05, learning_rate=lr, betas=(0.9, 0.95))
+            self.g_sched = self._create_scheduler(
+                cosine_lr, warmup_epochs, warmup_steps, num_epoch, 
+                lr_min, warmup_lr_init, decay_steps
+            )
+            self.accelerator.register_for_checkpointing(self.g_sched)
+
+        self.steps = steps
+        self.loaded_steps = -1
+
+        # Prepare everything together
+        if not test_only:
+            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)
+
+        # assume _ori_model does not exist in checkpoints
+        if compile:
+            _model = self.accelerator.unwrap_model(self.gpt_model)
+            self.ae_model = torch.compile(self.ae_model, mode="reduce-overhead")
+            _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.samp_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.diff_cfg = diff_cfg
+        self.cfg_schedule = cfg_schedule
+        self.diff_cfg_schedule = diff_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
+
+        self.result_folder = result_folder
+        self.model_saved_dir = os.path.join(result_folder, "models")
+        os.makedirs(self.model_saved_dir, exist_ok=True)
+
+        self.image_saved_dir = os.path.join(result_folder, "images")
+        os.makedirs(self.image_saved_dir, exist_ok=True)
+
+        self.cache_dir = Path(cache_dir)
+        self.enable_cache_latents = enable_cache_latents
+        self.seed = seed
+        self.cache_loader = None
+
+    @property
+    def device(self):
+        return self.accelerator.device
+
+    def _creat_optimizer(self, weight_decay, learning_rate, betas):
+        # start with all of the candidate parameters
+        param_dict = {pn: p for pn, p in self.gpt_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 self.accelerator.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
+
+    def _create_scheduler(self, cosine_lr, warmup_epochs, warmup_steps, num_epoch, lr_min, warmup_lr_init, decay_steps):
+        if warmup_epochs is not None:
+            warmup_steps = warmup_epochs * len(self.train_dl)
+        else:
+            assert warmup_steps is not None
+
+        scheduler = build_scheduler(
+            self.g_optim,
+            num_epoch,
+            len(self.train_dl),
+            lr_min,
+            warmup_steps,
+            warmup_lr_init,
+            decay_steps,
+            cosine_lr, # if not cosine_lr, then use step_lr (warmup, then fix)
+        )
+        return scheduler
+
+    def _load_state_dict(self, 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 self.accelerator.is_main_process:
+            print(f"Loaded model. Missing: {missing}, Unexpected: {unexpected}")
+
+    def _load_safetensors(self, 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()}
+        self._load_state_dict(state_dict, model)
+
+    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):
+            # 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")
+                        self._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):
+                        self._load_safetensors(model_path, model)
+        else:
+            # ckpt_path is something like 'path/to/models/step10.pt'
+            if ckpt_path.endswith(".safetensors"):
+                self._load_safetensors(ckpt_path, model)
+            else:
+                state_dict = torch.load(ckpt_path, map_location="cpu")
+                self._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=self.seed + 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.samp_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
+        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
+
+        generate_fn = generate if 'GPT' in self.model_name else generate_causal_dit
+        if not self.test_only:
+            classes = torch.tensor(self.eval_classes, device=self.device)
+            with self.accelerator.autocast():
+                slots = generate_fn(unwraped_gpt_model, classes, self.num_slots_to_gen, cfg_scale=self.cfg, diff_cfg=self.diff_cfg, cfg_schedule=self.cfg_schedule, diff_cfg_schedule=self.diff_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}_diffcfg-{self.diff_cfg_schedule}-{self.diff_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}_diffcfg-{self.diff_cfg_schedule}-{self.diff_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_fn(unwraped_gpt_model, labels, self.num_slots_to_gen, 
+                                            cfg_scale=self.cfg, diff_cfg=self.diff_cfg, 
+                                            cfg_schedule=self.cfg_schedule, diff_cfg_schedule=self.diff_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,
+                    "diff_cfg": self.diff_cfg,
+                    "cfg_schedule": self.cfg_schedule,
+                    "diff_cfg_schedule": self.diff_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}, diff_cfg={self.diff_cfg}, "
+                    f"cfg_schedule={self.cfg_schedule}, diff_cfg_schedule={self.diff_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()
+

paintmind/engine/misc.py

@@ -0,0 +1,260 @@
+import socket
+import os, sys, pdb
+from torch import inf
+import os.path as osp
+from pathlib import Path
+import builtins, datetime
+import torch.distributed as dist
+import os, sys, time, torch, copy, pdb
+from collections import defaultdict, deque
+
+def print_available_port():
+
+    return _find_free_port()
+
+def _find_free_port():
+    
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    # Binding to port 0 will cause the OS to find an available port for us
+    sock.bind(("", 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    # NOTE: there is still a chance the port could be taken by other processes.
+    return port
+
+def ensure_dir(dirpath):
+    
+    if not osp.exists(dirpath):
+        os.makedirs(dirpath, exist_ok=True)
+ 
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    builtin_print = builtins.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop('force', False)
+        force = force or (get_world_size() > 8)
+        if is_master or force:
+            now = datetime.datetime.now().time()
+            builtin_print('[{}] '.format(now), end='')  # print with time stamp
+            builtin_print(*args, **kwargs)
+
+    builtins.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+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 is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+
+    if args.dist_on_itp:
+        args.rank = int(os.environ['OMPI_COMM_WORLD_RANK'])
+        args.world_size = int(os.environ['OMPI_COMM_WORLD_SIZE'])
+        args.gpu = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])
+        assert isinstance(args.port, int) & (args.port > 0) & (args.port < 1<<30)
+        port = _find_free_port()
+        # args.dist_url = "tcp://%s:%s" % (port, os.environ['MASTER_PORT'])
+        args.dist_url = f'tcp://127.0.0.1:{port}'
+        os.environ['LOCAL_RANK'] = str(args.gpu)
+        os.environ['RANK'] = str(args.rank)
+        os.environ['WORLD_SIZE'] = str(args.world_size)
+        # ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT", "LOCAL_RANK"]
+    elif 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ['WORLD_SIZE'])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print('Not using distributed mode')
+        setup_for_distributed(is_master=True)  # hack
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}, gpu {}'.format(
+        args.rank, args.dist_url, args.gpu), flush=True)
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)
+
+class NativeScalerWithGradNormCount:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=None, parameters=None, create_graph=False, update_grad=True):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    if norm_type == inf:
+        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
+    else:
+        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
+    return total_norm
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer, loss_scaler, ema_params=None):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+
+        # ema
+        if ema_params is not None:
+            ema_state_dict = copy.deepcopy(model_without_ddp.state_dict())
+            for i, (name, _value) in enumerate(model_without_ddp.named_parameters()):
+                assert name in ema_state_dict
+                ema_state_dict[name] = ema_params[i]
+        else:
+            ema_state_dict = None
+
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'model_ema': ema_state_dict,
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
+
+
+def save_model_last(args, epoch, model, model_without_ddp, optimizer, loss_scaler, ema_params=None):
+
+    output_dir = Path(args.output_dir)
+    epoch_name = 'last'
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+
+        # ema
+        if ema_params is not None:
+            ema_state_dict = copy.deepcopy(model_without_ddp.state_dict())
+            for i, (name, _value) in enumerate(model_without_ddp.named_parameters()):
+                assert name in ema_state_dict
+                ema_state_dict[name] = ema_params[i]
+        else:
+            ema_state_dict = None
+
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'model_ema': ema_state_dict,
+                'optimizer': optimizer.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
+
+
+def load_model(args, model_without_ddp, optimizer, loss_scaler):
+    
+    if osp.exists(osp.join(args.resume, "checkpoint-last.pth")):
+        resume_path = osp.join(args.resume, "checkpoint-last.pth")
+    else:
+        resume_path = args.resume
+    if args.resume:
+        checkpoint = torch.load(resume_path, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % resume_path)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'evaluate') and args.evaluate):
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+def all_reduce_mean(x):
+
+    world_size = get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x

paintmind/engine/trainer.py

@@ -0,0 +1,695 @@
+import os, torch
+import os.path as osp
+import cv2
+import shutil
+import numpy as np
+import copy
+import torch_fidelity
+import torch.nn as nn
+from tqdm.auto import tqdm
+from collections import OrderedDict
+from einops import rearrange
+from accelerate import Accelerator
+from .util import instantiate_from_config
+from torchvision.utils import make_grid, save_image
+from torch.utils.data import DataLoader, random_split, DistributedSampler
+from paintmind.utils.lr_scheduler import build_scheduler
+from paintmind.utils.logger import SmoothedValue, MetricLogger, synchronize_processes, empty_cache
+from paintmind.engine.misc import is_main_process, all_reduce_mean, concat_all_gather
+from accelerate.utils import DistributedDataParallelKwargs, AutocastKwargs
+from torch.optim import AdamW
+from concurrent.futures import ThreadPoolExecutor
+from torchmetrics.functional.image import (
+    peak_signal_noise_ratio as psnr,
+    structural_similarity_index_measure as ssim
+)
+
+def requires_grad(model, flag=True):
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def save_img(img, save_path):
+    img = np.clip(img.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.numpy().transpose(0, 2, 3, 1) * 255, 0, 255).astype(np.uint8)
+    imgs = imgs[:, :, :, ::-1]  # RGB to BGR for all images at once
+    
+    # Use ProcessPoolExecutor which is generally better for CPU-bound tasks
+    # ThreadPoolExecutor is better for I/O-bound tasks like file saving
+    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):
+    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
+
+
+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 DiffusionTrainer(nn.Module):
+    def __init__(
+        self,
+        model,
+        dataset,
+        test_dataset=None,
+        test_only=False,
+        num_epoch=400,
+        valid_size=32,
+        lr=None,
+        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=0,
+        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",
+        steps=0,
+        cfg=1.0,
+        test_num_slots=None,
+        eval_fid=False,
+        fid_stats=None,
+        enable_ema=False,
+        use_multi_epochs_dataloader=False,
+        compile=False,
+        overfit=False,
+        making_cache=False,
+        cache_mode=False,
+        latent_cache_file=None,
+    ):
+        super().__init__()
+        kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
+        self.accelerator = Accelerator(
+            kwargs_handlers=[kwargs],
+            mixed_precision="bf16",
+            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
+
+        assert precision in ["bf16", "fp32"]
+        precision = "fp32"
+        if self.accelerator.is_main_process:
+            print("Overlooking specified precision and using autocast bf16...")
+        self.precision = precision
+
+        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
+            
+            # Create a padded dataset wrapper
+            class PaddedDataset(torch.utils.data.Dataset):
+                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]
+            
+            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,
+            )
+
+            effective_bs = batch_size * grad_accum_steps * self.accelerator.num_processes
+            if lr is None:
+                lr = blr * effective_bs / 256
+            if self.accelerator.is_main_process:
+                print(f"Effective batch size is {effective_bs}")
+
+            params = filter(lambda p: p.requires_grad, self.model.parameters())
+            self.g_optim = AdamW(params, lr=lr, betas=(0.9, 0.95), weight_decay=0)
+            self.g_sched = self._create_scheduler(
+                cosine_lr, warmup_epochs, warmup_steps, num_epoch, 
+                lr_min, warmup_lr_init, decay_steps
+            )
+            if self.g_sched is not None:
+                self.accelerator.register_for_checkpointing(self.g_sched)
+
+        self.steps = steps
+        self.loaded_steps = -1
+
+        # Prepare everything together
+        if not test_only:
+            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)
+
+        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.samp_every = sample_every
+        self.fid_every = fid_every
+        self.max_grad_norm = max_grad_norm
+
+        self.cache_mode = cache_mode
+
+        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.use_vq = model.params.use_vq if hasattr(model.params, "use_vq") else False
+        self.vq_beta = model.params.code_beta if hasattr(model.params, "code_beta") else 0.25
+
+        self.result_folder = result_folder
+        self.model_saved_dir = os.path.join(result_folder, "models")
+        os.makedirs(self.model_saved_dir, exist_ok=True)
+
+        self.image_saved_dir = os.path.join(result_folder, "images")
+        os.makedirs(self.image_saved_dir, exist_ok=True)
+
+    @property
+    def device(self):
+        return self.accelerator.device
+
+    def _create_scheduler(self, cosine_lr, warmup_epochs, warmup_steps, num_epoch, lr_min, warmup_lr_init, decay_steps):
+        if warmup_epochs is not None:
+            warmup_steps = warmup_epochs * len(self.train_dl)
+        else:
+            assert warmup_steps is not None
+
+        scheduler = build_scheduler(
+            self.g_optim,
+            num_epoch,
+            len(self.train_dl),
+            lr_min,
+            warmup_steps,
+            warmup_lr_init,
+            decay_steps,
+            cosine_lr, # if not cosine_lr, then use step_lr (warmup, then fix)
+        )
+        return scheduler
+
+    def _load_state_dict(self, state_dict):
+        """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 = self.accelerator.unwrap_model(self.model).load_state_dict(
+            state_dict, strict=False
+        )
+        if self.accelerator.is_main_process:
+            print(f"Loaded model. Missing: {missing}, Unexpected: {unexpected}")
+
+    def _load_safetensors(self, path):
+        """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()}
+        self._load_state_dict(state_dict)
+
+    def _load_checkpoint(self, ckpt_path=None):
+        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/'
+            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")
+                        self._load_state_dict(state_dict)
+                        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):
+                        self._load_safetensors(model_path)
+        else:
+            # ckpt_path is something like 'path/to/models/step10.pt'
+            if ckpt_path.endswith(".safetensors"):
+                self._load_safetensors(ckpt_path)
+            else:
+                state_dict = torch.load(ckpt_path)
+                self._load_state_dict(state_dict)
+        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.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
+            import shutil
+            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("vqgan")
+
+        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)):
+                # Move batch to device once
+                if isinstance(batch, tuple) or isinstance(batch, list):
+                    batch = tuple(b.to(self.device, non_blocking=True) for b in batch)
+                    if self.cache_mode:
+                        img, latent, targets = batch[0], batch[1], batch[2]
+                        img = img.to(self.device, non_blocking=True)
+                        latent = latent.to(self.device, non_blocking=True)
+                        targets = targets.to(self.device, non_blocking=True)
+                    else:
+                        latent = None
+                        img, targets = batch[0], batch[1]
+                        img = img.to(self.device, non_blocking=True)
+                        targets = targets.to(self.device, non_blocking=True)
+                else:
+                    img = batch
+                    latent = None
+
+                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, targets, latents=latent, epoch=epoch)
+                        # combine
+                        loss = sum([v for _, v in losses.items()])
+                        diff_loss = losses["diff_loss"]
+                        if self.use_vq:
+                            vq_loss = losses["vq_loss"]
+
+                    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.accelerator.unwrap_model(self.model).cancel_gradients_encoder(epoch)
+                    self.g_optim.step()
+                    if self.g_sched is not None:
+                        self.g_sched.step_update(self.steps)
+                    self.g_optim.zero_grad()
+
+                # synchronize_processes()
+
+                # update ema with state dict
+                if self.enable_ema:
+                    self.ema_model.update(self.accelerator.unwrap_model(self.model))
+
+                logger.update(diff_loss=diff_loss.item())
+                if self.use_vq:
+                    logger.update(vq_loss=vq_loss.item() / self.vq_beta)
+                if 'kl_loss' in losses:
+                    logger.update(kl_loss=losses["kl_loss"].item())
+                if 'repa_loss' in losses:
+                    logger.update(repa_loss=losses["repa_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.samp_every == 0) or (self.steps % self.fid_every == 0):
+                    empty_cache()
+                    self.evaluate()
+                    self.accelerator.wait_for_everyone()
+                    empty_cache()
+
+                # omitted all_gather here
+                # write_dict = dict(epoch=epoch)
+                # write_dict.update(diff_loss=diff_loss.item())
+                # if "kl_loss" in losses:
+                #     write_dict.update(kl_loss=losses["kl_loss"].item())
+                # if self.use_vq:
+                #     write_dict.update(vq_loss=vq_loss.item() / self.vq_beta)
+                # 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, use_ema=True):
+        self.model.eval()
+        # switch to ema params, only when eval_fid is True
+        use_ema = use_ema and self.enable_ema and self.eval_fid and not self.test_only
+        # use_ema = False
+        if use_ema:
+            if hasattr(self, "ema_model"):
+                model_without_ddp = self.accelerator.unwrap_model(self.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
+                    if name in self.ema_model.state_dict():
+                        ema_state_dict[name] = self.ema_model.state_dict()[name]
+                if self.accelerator.is_main_process:
+                    print("Switch to ema")
+                msg = model_without_ddp.load_state_dict(ema_state_dict)
+                if self.accelerator.is_main_process:
+                    print(msg)
+            else:
+                print("EMA model not found, using original model")
+                use_ema = False
+
+        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, targets, 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, targets, 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)):
+            # Create output directories
+            if self.test_dataset_size > 10000:
+                real_dir = "./dataset/imagenet/val256"
+            else:
+                real_dir = "./dataset/coco/val2017_256"
+            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, targets, 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_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,
+                        f"psnr_{'ema' if use_ema else 'test'}": psnr_val,
+                        f"ssim_{'ema' if use_ema else 'test'}": ssim_val,
+                        "cfg": cfg_value
+                    }, step=self.steps)
+                    
+                    print(f"{'EMA ' if use_ema else ''}{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)
+
+        # 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.model.train()

paintmind/engine/util.py

@@ -0,0 +1,572 @@
+import numpy as np
+import os.path as osp
+import torch_fidelity
+from PIL import Image
+from tqdm import tqdm
+import pickle as pkl
+import os, hashlib, pdb
+from pathlib import Path
+from torch import Tensor
+import torch, torchvision
+from einops import rearrange
+from omegaconf import OmegaConf
+import torch.distributed as dist
+from typing import List, Optional
+from torchvision import transforms
+from io import BytesIO as Bytes2Data
+from smart_open import open
+from .misc import is_main_process, get_rank
+import importlib, datetime, requests, time, shutil
+from collections import defaultdict, deque, OrderedDict
+from dotwiz import DotWiz
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+    
+def customized_collate_fn(batch):
+
+    collate_fn = {}
+    if len(batch) < 2:
+        for key, value in batch[0].items():
+            collate_fn[key] = [value]
+    else:
+
+        for i, dd in enumerate(batch):
+            if i < 1:
+                for key, value in dd.items():
+                    collate_fn[key] = [value]
+            else:
+                for key, value in dd.items():
+                    collate_fn[key].append(value)
+    
+    return collate_fn
+
+
+def trivial_batch_collator(batch):
+    """
+    A batch collator that does nothing.
+    """
+    return batch
+
+class NestedTensor(object):
+    def __init__(self, tensors, mask: Optional[Tensor]):
+        self.tensors = tensors
+        self.mask = mask
+
+    def to(self, device):
+        # type: (Device) -> NestedTensor # noqa
+        cast_tensor = self.tensors.to(device)
+        mask = self.mask
+        if mask is not None:
+            assert mask is not None
+            cast_mask = mask.to(device)
+        else:
+            cast_mask = None
+        return NestedTensor(cast_tensor, cast_mask)
+
+    def decompose(self):
+        return self.tensors, self.mask
+
+    def __repr__(self):
+        return str(self.tensors)
+
+
+def _max_by_axis(the_list):
+    # type: (List[List[int]]) -> List[int]
+    maxes = the_list[0]
+    for sublist in the_list[1:]:
+        for index, item in enumerate(sublist):
+            maxes[index] = max(maxes[index], item)
+    return maxes
+
+
+def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
+    # TODO make this more general
+    if tensor_list[0].ndim == 3:
+        if torchvision._is_tracing():
+            # nested_tensor_from_tensor_list() does not export well to ONNX
+            # call _onnx_nested_tensor_from_tensor_list() instead
+            return _onnx_nested_tensor_from_tensor_list(tensor_list)
+
+        # TODO make it support different-sized images
+        max_size = _max_by_axis([list(img.shape) for img in tensor_list])
+        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
+        batch_shape = [len(tensor_list)] + max_size
+        b, c, h, w = batch_shape
+        dtype = tensor_list[0].dtype
+        device = tensor_list[0].device
+        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
+        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
+        for img, pad_img, m in zip(tensor_list, tensor, mask):
+            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
+            m[: img.shape[1], : img.shape[2]] = False
+    else:
+        raise ValueError("not supported")
+    return NestedTensor(tensor, mask)
+
+
+# _onnx_nested_tensor_from_tensor_list() is an implementation of
+# nested_tensor_from_tensor_list() that is supported by ONNX tracing.
+@torch.jit.unused
+def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor:
+    max_size = []
+    for i in range(tensor_list[0].dim()):
+        max_size_i = torch.max(
+            torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)
+        ).to(torch.int64)
+        max_size.append(max_size_i)
+    max_size = tuple(max_size)
+
+    # work around for
+    # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
+    # m[: img.shape[1], :img.shape[2]] = False
+    # which is not yet supported in onnx
+    padded_imgs = []
+    padded_masks = []
+    for img in tensor_list:
+        padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))]
+        padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0]))
+        padded_imgs.append(padded_img)
+
+        m = torch.zeros_like(img[0], dtype=torch.int, device=img.device)
+        padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1)
+        padded_masks.append(padded_mask.to(torch.bool))
+
+    tensor = torch.stack(padded_imgs)
+    mask = torch.stack(padded_masks)
+
+    return NestedTensor(tensor, mask=mask)
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+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.float64, device='cuda')
+        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('max mem: {memory:.0f}')
+        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():
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time),
+                        memory=torch.cuda.max_memory_allocated() / MB))
+                else:
+                    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)))
+        print('{} Total time: {} ({:.4f} s / it)'.format(
+            header, total_time_str, total_time / len(iterable)))
+
+
+def all_reduce_mean(x):
+    world_size = dist.get_world_size()
+    if world_size > 1:
+        x_reduce = torch.tensor(x).cuda()
+        dist.all_reduce(x_reduce)
+        x_reduce /= world_size
+        return x_reduce.item()
+    else:
+        return x
+
+
+class NativeScaler:
+    state_dict_key = "amp_scaler"
+
+    def __init__(self):
+        self._scaler = torch.cuda.amp.GradScaler()
+
+    def __call__(self, loss, optimizer, clip_grad=3., parameters=None, create_graph=False, update_grad=True):
+        self._scaler.scale(loss).backward(create_graph=create_graph)
+        if update_grad:
+            if clip_grad is not None:
+                assert parameters is not None
+                self._scaler.unscale_(optimizer)  # unscale the gradients of optimizer's assigned params in-place
+                norm = torch.nn.utils.clip_grad_norm_(parameters, clip_grad)
+            else:
+                self._scaler.unscale_(optimizer)
+                norm = get_grad_norm_(parameters)
+            self._scaler.step(optimizer)
+            self._scaler.update()
+        else:
+            norm = None
+        return norm
+
+    def state_dict(self):
+        return self._scaler.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self._scaler.load_state_dict(state_dict)
+
+
+def get_grad_norm_(parameters, norm_type: float = 2.0) -> torch.Tensor:
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = [p for p in parameters if p.grad is not None and p.requires_grad]
+    norm_type = float(norm_type)
+    if len(parameters) == 0:
+        return torch.tensor(0.)
+    device = parameters[0].grad.device
+    total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]),
+                            norm_type)
+    return total_norm
+
+
+def get_obj_from_str(string, reload=False):
+    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):
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def save_on_master(*args, **kwargs):
+    if dist.get_rank() == 0:
+        torch.save(*args, **kwargs)
+
+
+def save_model(args, epoch, model, model_without_ddp, optimizer_g, optimizer_d, loss_scaler):
+    output_dir = Path(args.output_dir)
+    epoch_name = str(epoch)
+    if loss_scaler is not None:
+        checkpoint_paths = [output_dir / ('checkpoint-%s.pth' % epoch_name)]
+        for checkpoint_path in checkpoint_paths:
+            to_save = {
+                'model': model_without_ddp.state_dict(),
+                'optimizer_g': optimizer_g.state_dict(),
+                'optimizer_d': optimizer_d.state_dict(),
+                'epoch': epoch,
+                'scaler': loss_scaler.state_dict(),
+                'args': args,
+            }
+
+            save_on_master(to_save, checkpoint_path)
+    else:
+        client_state = {'epoch': epoch}
+        model.save_checkpoint(save_dir=args.output_dir, tag="checkpoint-%s" % epoch_name, client_state=client_state)
+
+
+def load_model(args, model_without_ddp, optimizer_g, optimizer_d, loss_scaler):
+    if args.resume:
+        if args.resume.startswith('https'):
+            checkpoint = torch.hub.load_state_dict_from_url(
+                args.resume, map_location='cpu', check_hash=True)
+        else:
+            checkpoint = torch.load(args.resume, map_location='cpu')
+        model_without_ddp.load_state_dict(checkpoint['model'])
+        print("Resume checkpoint %s" % args.resume)
+        if 'optimizer_g' in checkpoint and 'epoch' in checkpoint and not (hasattr(args, 'eval') and args.eval):
+            optimizer_g.load_state_dict(checkpoint['optimizer_g'])
+            optimizer_d.load_state_dict(checkpoint['optimizer_d'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class KeyNotFoundError(Exception):
+    def __init__(self, cause, keys=None, visited=None):
+        self.cause = cause
+        self.keys = keys
+        self.visited = visited
+        messages = list()
+        if keys is not None:
+            messages.append("Key not found: {}".format(keys))
+        if visited is not None:
+            messages.append("Visited: {}".format(visited))
+        messages.append("Cause:\n{}".format(cause))
+        message = "\n".join(messages)
+        super().__init__(message)
+
+
+def retrieve(
+        list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False
+):
+    """Given a nested list or dict return the desired value at key expanding
+    callable nodes if necessary and :attr:`expand` is ``True``. The expansion
+    is done in-place.
+
+    Parameters
+    ----------
+        list_or_dict : list or dict
+            Possibly nested list or dictionary.
+        key : str
+            key/to/value, path like string describing all keys necessary to
+            consider to get to the desired value. List indices can also be
+            passed here.
+        splitval : str
+            String that defines the delimiter between keys of the
+            different depth levels in `key`.
+        default : obj
+            Value returned if :attr:`key` is not found.
+        expand : bool
+            Whether to expand callable nodes on the path or not.
+
+    Returns
+    -------
+        The desired value or if :attr:`default` is not ``None`` and the
+        :attr:`key` is not found returns ``default``.
+
+    Raises
+    ------
+        Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is
+        ``None``.
+    """
+
+    keys = key.split(splitval)
+
+    success = True
+    try:
+        visited = []
+        parent = None
+        last_key = None
+        for key in keys:
+            if callable(list_or_dict):
+                if not expand:
+                    raise KeyNotFoundError(
+                        ValueError(
+                            "Trying to get past callable node with expand=False."
+                        ),
+                        keys=keys,
+                        visited=visited,
+                    )
+                list_or_dict = list_or_dict()
+                parent[last_key] = list_or_dict
+
+            last_key = key
+            parent = list_or_dict
+
+            try:
+                if isinstance(list_or_dict, dict):
+                    list_or_dict = list_or_dict[key]
+                else:
+                    list_or_dict = list_or_dict[int(key)]
+            except (KeyError, IndexError, ValueError) as e:
+                raise KeyNotFoundError(e, keys=keys, visited=visited)
+
+            visited += [key]
+        # final expansion of retrieved value
+        if expand and callable(list_or_dict):
+            list_or_dict = list_or_dict()
+            parent[last_key] = list_or_dict
+    except KeyNotFoundError as e:
+        if default is None:
+            raise e
+        else:
+            list_or_dict = default
+            success = False
+
+    if not pass_success:
+        return list_or_dict
+    else:
+        return list_or_dict, success
+
+
+if __name__ == "__main__":
+    config = {"keya": "a",
+              "keyb": "b",
+              "keyc":
+                  {"cc1": 1,
+                   "cc2": 2,
+                   }
+              }
+    from omegaconf import OmegaConf
+
+    config = OmegaConf.create(config)
+    print(config)
+    retrieve(config, "keya")
+
+def instantiate_from_config(config):
+
+    if not "target" in config:
+        raise KeyError("Expected key `target` to instantiate.")
+    return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+def get_obj_from_str(string, reload=False):
+    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)

paintmind/stage1/diffuse_slot.py

@@ -0,0 +1,808 @@
+import torch
+import math
+import numpy as np
+import torch.nn as nn
+
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from paintmind.stage1.diffusion import create_diffusion
+from paintmind.stage1.diffusion_transfomers import DiT
+from paintmind.stage1.quantize import DiagonalGaussianDistribution
+from paintmind.stage1.transport import create_transport, Sampler
+
+from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from transformers import SiglipVisionModel, CLIPVisionModel
+
+CLIP_DEFAULT_MEAN = (0.48145466, 0.4578275, 0.40821073)
+CLIP_DEFAULT_STD = (0.26862954, 0.26130258, 0.27577711)
+SIGLIP_DEFAULT_MEAN = (0.5, 0.5, 0.5)
+SIGLIP_DEFAULT_STD = (0.5, 0.5, 0.5)
+
+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),
+            )
+
+class DiT_with_autoenc_cond(DiT):
+    def __init__(
+        self,
+        *args,
+        num_autoenc=32,
+        autoenc_dim=4,
+        cond_method="adaln",
+        mask_type="simple",
+        class_cond=False,
+        use_repa=False,
+        z_dim=768,
+        encoder_depth=8,
+        projector_dim=2048,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.autoenc_dim = autoenc_dim
+        self.class_cond = class_cond
+        self.mask_type = mask_type
+        self.hidden_size = kwargs["hidden_size"]
+        self.cond_drop_prob = self.y_embedder.dropout_prob  # 0.1 without cond guidance
+        self.null_cond = nn.Parameter(torch.zeros(1, num_autoenc, autoenc_dim))
+        torch.nn.init.normal_(self.null_cond, std=.02)
+        # NOTE: adaln is adaptive layer normalization, token fed the cond to the attention layer
+        assert cond_method in [
+            "adaln",
+            "token",
+            "token+adaln",
+        ], f"Invalid cond_method: {cond_method}"
+        self.cond_method = cond_method
+        if "token" in cond_method:
+            self.autoenc_cond_embedder = nn.Linear(autoenc_dim, self.hidden_size)
+        else:
+            self.autoenc_cond_embedder = nn.Linear(
+                num_autoenc * autoenc_dim, self.hidden_size
+            )
+
+        if cond_method == "token+adaln":
+            self.autoenc_proj_ln = nn.Linear(self.hidden_size, self.hidden_size)
+
+        if not class_cond:
+            self.y_embedder = nn.Identity()
+        
+        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)
+        # NOTE: this dropout will replace some condition from the autoencoder to null condition
+        # this is to enable classifier-free guidance.
+        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)
+        if "token" in self.cond_method:
+            return self.autoenc_cond_embedder(autoenc_cond_drop)
+        return self.autoenc_cond_embedder(autoenc_cond_drop.reshape(batch_size, -1))
+
+    # def forward(self, x, t, y, autoenc_cond):
+    def forward(self, x, t, autoenc_cond, drop_mask=None, y=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
+        y: (N,) tensor of class labels
+        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
+        N, T, D = x.shape
+
+        c = self.t_embedder(t)  # (N, D)
+        if y is not None and self.class_cond:
+            y = self.y_embedder(y, self.training)    # (N, D)
+            c = c + y                                # (N, D)
+
+        if self.mask_type == "replace":
+            autoenc = self.embed_cond(autoenc_cond, drop_mask)
+        else:
+            autoenc = self.embed_cond(autoenc_cond)
+
+        if self.cond_method == "adaln":
+            c = c + autoenc  # add the encoder condition to adaln
+        elif self.cond_method == "token":
+            num_tokens = x.shape[1]
+            # append the autoencoder condition to the token sequence
+            x = torch.cat((x, autoenc), dim=1)
+        elif self.cond_method == "token+adaln":
+            c = c + self.autoenc_proj_ln(autoenc.mean(dim=1))
+            num_tokens = x.shape[1]
+            x = torch.cat((x, autoenc), dim=1)
+        else:
+            raise ValueError(f"Invalid cond_method: {self.cond_method}")
+
+        for i, block in enumerate(self.blocks):
+            if self.mask_type == "replace":
+                x = block(x, c)  # (N, T, D)
+            else:
+                x = block(x, c, drop_mask)  # (N, T, D)
+            if (i + 1) == self.encoder_depth and self.use_repa:
+                projected = self.projector(x)
+                self._repa_hook = projected[:, :num_tokens]
+
+        if "token" in self.cond_method:
+            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, 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)
+
+#################################################################################
+#                                   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,
+}
+
+from torch.distributions import Geometric
+
+class NestedSampler(nn.Module):
+    def __init__(
+        self,
+        num_slots,
+        rho=0.03,
+        nest_dist="geometric",
+        mask_type="simple",
+        null_prob=0.1,
+        allow_zero=False,
+        one_slot_before=0,  # Use only one slot before this epoch
+    ):
+        super().__init__()
+        self.num_slots = num_slots
+        self.mask_type = mask_type
+        self.rho = rho
+        self.geometric = Geometric(rho)
+        self.nest_dist = nest_dist
+        self.null_prob = null_prob
+        self.allow_zero = allow_zero
+        self.one_slot_before = one_slot_before
+        self.register_buffer("arange", torch.arange(num_slots))
+
+    def _apply_epoch_constraint(self, samples, epoch=None):
+        if epoch is not None and epoch < self.one_slot_before:
+            return torch.ones_like(samples)
+        return samples
+
+    def _apply_null_prob(self, samples, num):
+        # First determine which samples will be 0 based on null_prob
+        null_mask = torch.rand(num, device=samples.device) < self.null_prob
+        # Replace with 0 where null_mask is True
+        return torch.where(null_mask, torch.zeros_like(samples), samples)
+
+    def geometric_sample(self, num):
+        return self.geometric.sample([num]) + int(not self.allow_zero)
+
+    def uniform_sample(self, num):
+        return torch.randint(int(not self.allow_zero), self.num_slots + 1, (num,))
+
+    def power2_uniform_sample(self, num):
+        # Get powers of 2 up to num_slots and add num_slots
+        choices = [2**i for i in range(int(math.log2(self.num_slots)) + 1)]
+        if self.num_slots not in choices:
+            choices.append(self.num_slots)
+        return torch.tensor(choices)[torch.randint(0, len(choices), (num,))]
+
+    def sample(self, num, epoch=None):
+        if self.nest_dist == "geometric":
+            samples = self.geometric_sample(num)
+        elif self.nest_dist == "uniform":
+            samples = self.uniform_sample(num)
+        elif self.nest_dist == "power2_uniform":
+            samples = self.power2_uniform_sample(num)
+        else:
+            raise ValueError(f"Invalid nest_dist: {self.nest_dist}")
+        samples = self._apply_epoch_constraint(samples, epoch)
+        return self._apply_null_prob(samples, num)
+
+    def forward(self, batch_size, num_patches, device, inference_with_n_slots=-1, coupled_value=None, epoch=None):
+        if self.training:
+            if coupled_value is None:
+                b = self.sample(batch_size, epoch).to(device)
+            else:
+                b = coupled_value.long().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)
+        if self.mask_type == "replace":
+            return slot_mask
+        else:
+            return self.get_cond_attn_mask(slot_mask.unsqueeze(1), num_patches, self.num_slots, device)
+
+    def get_cond_attn_mask(self, slot_mask, num_patches, num_slots, device):
+        num_tokens = num_patches + num_slots
+        batch_size = slot_mask.shape[0]
+        if self.mask_type == "simple":
+            attn_mask = torch.ones((batch_size, num_tokens, num_tokens), dtype=torch.bool, device=device)
+            attn_mask[:, :, num_patches:] = slot_mask.expand(-1, num_tokens, -1)
+        elif self.mask_type == "causal":
+            attn_mask = torch.zeros((batch_size, num_tokens, num_tokens), dtype=torch.bool, device=device)
+            # 1) patches can see each other
+            attn_mask[:, :num_patches, :num_patches] = True
+            # 2) pathes can not see the last few slots
+            slot_mask = slot_mask.expand(-1, num_patches, -1)
+            attn_mask[:, :num_patches, num_patches:] = slot_mask
+            # 3) remaining slots are causal to each other
+            causal_mask = torch.ones((num_slots, num_slots), dtype=torch.bool, device=device).tril(diagonal=0)
+            attn_mask[:, num_patches:, num_patches:] = causal_mask
+            # 4) only the first slot can see the patches
+            attn_mask[:, num_patches + 1:, :num_patches] = False
+        else:
+            raise NotImplementedError(f"Invalid mask_type: {self.mask_type}")
+        return attn_mask.unsqueeze(1) # (batch_size, 1, num_tokens, num_tokens)
+
+class DiffuseSlot(nn.Module):
+    def __init__(
+        self,
+        encoder="vit_base_patch16",
+        drop_path_rate=0.1,
+        enc_img_size=256,
+        enc_causal=True,
+        enc_use_mlp=False,
+        enc_hidden_dim=4096,
+        num_slots=16,
+        slot_dim=256,
+        slot_through=True,
+        norm_slots=False,
+        use_kl_loss=False,
+        kl_loss_weight=1e-6,
+        enable_nest=False,
+        enable_nest_after=-1,
+        nest_dist="geometric",
+        nest_rho=0.03,
+        nest_null_prob=0.1,
+        nest_allow_zero=False,
+        nest_one_slot_before=0,
+        coupled_sampling=False,
+        coupled_rho=-0.8,
+        dit_class_cond=False,
+        dit_mask_type="simple",
+        cond_method="adaln",
+        dit_model="DiT-B-4",
+        vae="stabilityai/sd-vae-ft-ema",
+        vae_path="pretrained_models/kl16.ckpt",
+        pretrained_dit=None,
+        pretrained_encoder=None,
+        freeze_dit=False,
+        freeze_vit_after=-1,
+        num_sampling_steps="ddim25",
+        ckpt_path=None,
+        ema_path=None,
+        use_repa=False,
+        repa_encoder="dinov2_vitb14",
+        repa_encoder_depth=8,
+        repa_loss_weight=1.0,
+        use_sit=False,
+        **kwargs,
+    ):
+        super().__init__()
+
+        z_dim = 0
+        self.use_repa = use_repa
+        self.repa_encoder_name = repa_encoder
+        self.repa_loss_weight = repa_loss_weight
+        self.use_sit = use_sit
+        if use_repa:
+            if "dinov2" in repa_encoder:
+                if "vitb" in repa_encoder or "vit_b" in repa_encoder or "vit-b" in repa_encoder:
+                    self.repa_encoder = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14')
+                elif "vitl" in repa_encoder or "vit_l" in repa_encoder or "vit-l" in repa_encoder:
+                    self.repa_encoder = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitl14')
+                else:
+                    raise ValueError(f"Invalid repa_encoder: {repa_encoder}")
+                self.repa_encoder.image_size = 224
+            elif "clip" in repa_encoder:
+                if "vitb" in repa_encoder or "vit_b" in repa_encoder or "vit-b" in repa_encoder or "vit-base" in repa_encoder:
+                    self.repa_encoder = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch16")
+                elif "vitl" in repa_encoder or "vit_l" in repa_encoder or "vit-l" in repa_encoder or "vit-large" in repa_encoder:
+                    self.repa_encoder = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14")
+                else:
+                    raise ValueError(f"Invalid repa_encoder: {repa_encoder}")
+                self.repa_encoder.embed_dim = self.repa_encoder.config.hidden_size
+                self.repa_encoder.image_size = self.repa_encoder.config.image_size
+            elif "siglip2" in repa_encoder:
+                if "vitb" in repa_encoder or "vit_b" in repa_encoder or "vit-b" in repa_encoder or "vit-base" in repa_encoder:
+                    self.repa_encoder = SiglipVisionModel.from_pretrained("google/siglip2-base-patch16-256")
+                elif "vitl" in repa_encoder or "vit_l" in repa_encoder or "vit-l" in repa_encoder or "vit-large" in repa_encoder:
+                    self.repa_encoder = SiglipVisionModel.from_pretrained("google/siglip2-large-patch16-256")
+                else:
+                    raise ValueError(f"Invalid repa_encoder: {repa_encoder}")
+                self.repa_encoder.embed_dim = self.repa_encoder.config.hidden_size
+                self.repa_encoder.image_size = self.repa_encoder.config.image_size
+            elif "siglip" in repa_encoder:
+                if "vitb" in repa_encoder or "vit_b" in repa_encoder or "vit-b" in repa_encoder or "vit-base" in repa_encoder:
+                    self.repa_encoder = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-256")
+                elif "vitl" in repa_encoder or "vit_l" in repa_encoder or "vit-l" in repa_encoder or "vit-large" in repa_encoder:
+                    self.repa_encoder = SiglipVisionModel.from_pretrained("google/siglip-large-patch16-256")
+                else:
+                    raise ValueError(f"Invalid repa_encoder: {repa_encoder}")
+                self.repa_encoder.embed_dim = self.repa_encoder.config.hidden_size
+                self.repa_encoder.image_size = self.repa_encoder.config.image_size
+            else:
+                raise ValueError(f"Invalid repa_encoder: {repa_encoder}")
+            for param in self.repa_encoder.parameters():
+                param.requires_grad = False
+            self.repa_encoder.eval()
+            z_dim = self.repa_encoder.embed_dim
+
+        # DiT part
+        if not use_sit:
+            self.diffusion = create_diffusion(timestep_respacing="")
+            self.gen_diffusion = create_diffusion(timestep_respacing=num_sampling_steps)
+        else:
+            self.transport = create_transport()
+            self.sampler = Sampler(self.transport)
+        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,
+            cond_method=cond_method,
+            class_cond=dit_class_cond,
+            mask_type=dit_mask_type,
+            use_repa=use_repa,
+            encoder_depth=repa_encoder_depth,
+            z_dim=z_dim,
+            learn_sigma=not use_sit,
+        )
+        self.dit_patch_size = self.dit.x_embedder.patch_size[0]
+        self.dit_num_patches = (self.dit_input_size // self.dit_patch_size) ** 2
+        self.pretrained_dit = pretrained_dit
+        if pretrained_dit is not None:
+            # now we load some pretrained model
+            dit_ckpt = torch.load(pretrained_dit, map_location="cpu")
+            msg = self.dit.load_state_dict(dit_ckpt, strict=False)
+            print("Load DiT from ckpt")
+            print(msg)
+        self.freeze_dit = freeze_dit
+        if freeze_dit:
+            assert pretrained_dit is not None, "pretrained_dit must be provided"
+            for param in self.dit.parameters():
+                param.requires_grad = False
+
+        if "mar" in vae:
+            from diffusers import AutoencoderKL
+            self.vae = AutoencoderKL.from_pretrained("xwen99/mar-vae-kl16")
+            self.scaling_factor = 0.2325
+        elif vae == "openai/consistency-decoder":
+            from diffusers import ConsistencyDecoderVAE
+            self.vae = ConsistencyDecoderVAE.from_pretrained(vae)
+            self.scaling_factor = 0.18215
+        else: # eg, "stabilityai/sd-vae-ft-ema"
+            from diffusers import AutoencoderKL
+            self.vae = AutoencoderKL.from_pretrained(vae)
+            self.scaling_factor = 0.18215
+        
+        self.vae.eval().requires_grad_(False)
+
+        # image encoder part
+        import paintmind.stage1.vision_transformers as vision_transformer
+
+        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,
+            slot_through=slot_through,
+            drop_path_rate=drop_path_rate,
+        )
+        self.num_slots = num_slots
+        self.norm_slots = norm_slots
+        self.use_kl_loss = use_kl_loss
+        self.kl_loss_weight = kl_loss_weight
+        self.num_channels = self.encoder.num_features
+        self.pretrained_encoder = pretrained_encoder
+        if pretrained_encoder is not None:
+            # __import__("ipdb").set_trace()
+            encoder_ckpt = torch.load(pretrained_encoder, map_location="cpu")
+            # drop pos_embed from ckpt
+            encoder_ckpt = {
+                k.replace("blocks.", "blocks.0."): v
+                for k, v in encoder_ckpt.items()
+                if not k.startswith("pos_embed")
+            }
+            msg = self.encoder.load_state_dict(encoder_ckpt, strict=False)
+            print("Load encoder from ckpt")
+            print(msg)
+            
+        if not enc_use_mlp:
+            self.encoder2slot = nn.Linear(self.num_channels, slot_dim * 2 if self.use_kl_loss else slot_dim)
+        else:
+            self.encoder2slot = nn.Sequential(
+                nn.Linear(self.num_channels, enc_hidden_dim),
+                nn.GELU(),
+                nn.Linear(enc_hidden_dim, slot_dim * 2 if self.use_kl_loss else slot_dim),
+            )
+
+        self.nested_sampler = NestedSampler(
+            num_slots,
+            rho=nest_rho,
+            nest_dist=nest_dist,
+            mask_type=dit_mask_type,
+            null_prob=nest_null_prob,
+            allow_zero=nest_allow_zero,
+            one_slot_before=nest_one_slot_before,
+        )
+        self.nest_allow_zero = nest_allow_zero
+        self.nest_rho = nest_rho
+        self.use_coupled_sampling = coupled_sampling
+        self.couple_sampling_rho = coupled_rho
+        self.enable_nest = enable_nest
+        self.enable_nest_after = enable_nest_after
+        self.freeze_vit_after = freeze_vit_after
+        self.current_epoch = 0
+    
+    def coupled_sampling(self, timestamps):
+        """
+        Convert timestamps to coupled num_slots values where higher timestamps 
+        tend to produce lower num_slots values.
+        
+        Args:
+            timestamps: Tensor of shape (batch_size,) with values in [0, 1000)
+        
+        Returns:
+            Tensor of shape (batch_size,) with values in [1, num_slots + 1)
+        """
+        # Normalize timestamps to [0, 1]
+        t_normalized = 1 - (timestamps.float() / timestamps.max())
+        # Scale to [1, num_slots + 1) and round to integers
+        adder = int(not self.nest_allow_zero)
+        scaled = adder + t_normalized * (self.num_slots + 1 - adder)
+        num_slots2use = scaled.long().clamp(adder, self.num_slots)
+        return num_slots2use
+
+    @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):
+        if "dinov2" in self.repa_encoder_name:
+            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)
+        elif "clip" in self.repa_encoder_name:
+            mean = torch.Tensor(CLIP_DEFAULT_MEAN).to(x.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            std = torch.Tensor(CLIP_DEFAULT_STD).to(x.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+        elif "siglip" in self.repa_encoder_name:
+            mean = torch.Tensor(SIGLIP_DEFAULT_MEAN).to(x.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            std = torch.Tensor(SIGLIP_DEFAULT_STD).to(x.device).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+        else:
+            raise ValueError(f"Invalid repa_encoder: {self.repa_encoder_name}")
+        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')
+        if "dinov2" in self.repa_encoder_name:
+            x = self.repa_encoder.forward_features(x)['x_norm_patchtokens']
+        else:
+            x = self.repa_encoder(x)["last_hidden_state"]
+        return x
+
+    def encode_slots(self, x):
+        if self.pretrained_encoder is not None:
+            x = F.interpolate(x, size=224, mode='bicubic')
+        slots = self.encoder(x, is_causal=self.enc_causal)
+        slots = self.encoder2slot(slots)
+        if self.norm_slots:
+            if not self.use_kl_loss:
+                slots_std = torch.std(slots, dim=-1, keepdim=True)
+                slots_mean = torch.mean(slots, dim=-1, keepdim=True)
+                slots = (slots - slots_mean) / slots_std # this works better than kl loss
+            else:
+                slots = DiagonalGaussianDistribution(slots)
+        return slots
+    
+    def forward_with_latents(self,
+                             x_vae,
+                             slots,
+                             z,
+                             targets=None,
+                             sample=False,
+                             epoch=None,
+                             inference_with_n_slots=-1, 
+                             cfg=1.0):
+        losses = {}
+        batch_size = x_vae.shape[0]
+        num_patches = self.dit_num_patches
+        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:
+            if self.use_coupled_sampling:
+                num_slots2use = self.coupled_sampling(t)
+            else:
+                num_slots2use = None
+            drop_mask = self.nested_sampler(
+                batch_size, num_patches, device, 
+                inference_with_n_slots=inference_with_n_slots, 
+                coupled_value=num_slots2use,
+                epoch=epoch
+            )
+        else:
+            drop_mask = None
+            
+        if sample:
+            return self.sample(slots if not self.use_kl_loss else slots.sample(), drop_mask=drop_mask, targets=targets, cfg=cfg)
+
+        model_kwargs = dict(autoenc_cond=slots if not self.use_kl_loss else slots.sample(), drop_mask=drop_mask, y=targets)
+        if not self.use_sit:
+            loss_dict = self.diffusion.training_losses(self.dit, x_vae, t, model_kwargs)
+        else:
+            loss_dict = self.transport.training_losses(self.dit, x_vae, model_kwargs)
+        diff_loss = loss_dict["loss"].mean()
+        losses["diff_loss"] = diff_loss
+
+        if self.use_kl_loss:
+            kl_loss = slots.kl()
+            losses["kl_loss"] = kl_loss.mean() * self.kl_loss_weight
+        
+        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])
+            
+            assert z_tilde.shape[-1] == z.shape[-1], f"Feature dimensions don't match: {z_tilde.shape} vs {z.shape}"
+            
+            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,
+                targets=None,
+                latents=None,
+                sample=False,
+                epoch=None,
+                inference_with_n_slots=-1,
+                cfg=1.0):
+
+        # it will be used in train() and decide whether to set the encoder to eval mode
+        if epoch is not None:
+            self.current_epoch = epoch
+
+        if latents is None:
+            x_vae = self.vae_encode(x)  # (N, C, H, W)
+        else:
+            x_vae = latents
+        
+        if self.use_repa:
+            z = self.repa_encode(x)
+        else:
+            z = None
+
+        slots = self.encode_slots(x)
+        return self.forward_with_latents(x_vae, slots, z, targets, sample, epoch, inference_with_n_slots, cfg)
+
+
+    @torch.no_grad()
+    def sample(self, slots, drop_mask=None, targets=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)
+            if targets is not None:
+                targets = torch.cat([targets, targets], 0)
+            model_kwargs = dict(autoenc_cond=slots, drop_mask=drop_mask, y=targets, cfg_scale=cfg)
+            sample_fn = self.dit.forward_with_cfg
+        else:
+            model_kwargs = dict(autoenc_cond=slots, drop_mask=drop_mask, y=targets)
+            sample_fn = self.dit.forward
+        # Sample images:
+        if not self.use_sit:
+            samples = self.gen_diffusion.p_sample_loop(
+                sample_fn,
+                z.shape,
+                z,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+                progress=False,
+                device=device,
+            )
+        else:
+            sde_sample_fn = self.sampler.sample_sde(diffusion_form="sigma")
+            samples = sde_sample_fn(z, sample_fn, **model_kwargs)[-1]
+        if cfg != 1.0:
+            samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+        samples = self.vae_decode(samples)
+        return samples
+
+    def cancel_gradients_encoder(self, epoch):
+        """Cancel gradients for encoder components after backward pass"""
+        if (epoch is not None 
+            and epoch >= self.freeze_vit_after 
+            and self.freeze_vit_after != -1):
+            # Directly access parameters from the modules
+            for p in self.encoder.parameters():
+                if p.grad is not None:
+                    p.grad = None
+            for p in self.encoder2slot.parameters():
+                if p.grad is not None:
+                    p.grad = None
+
+    def train(self, mode=True):
+        """Override train() to keep certain components in eval mode"""
+        super().train(mode)
+        # VAE should always be in eval mode
+        self.vae.eval()
+        
+        # Keep encoder in eval mode if frozen
+        if (self.freeze_vit_after != -1 and 
+            hasattr(self, 'current_epoch') and 
+            self.current_epoch >= self.freeze_vit_after):
+            self.encoder.eval()
+            self.encoder2slot.eval()
+            
+        # Keep DiT in eval mode if frozen
+        if self.freeze_dit:
+            self.dit.eval()
+            
+        return self

paintmind/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,
+    )

paintmind/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

paintmind/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)

paintmind/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)
+    

paintmind/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()

paintmind/stage1/diffusion_transfomers.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 paintmind.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,
+}

paintmind/stage1/fused_attention.py

@@ -0,0 +1,94 @@
+import math
+import torch
+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, # True indicates parts should take part in attention in this API
+            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
+
+
+class MultiHeadCrossAttention(nn.Module):
+    def __init__(self, d_model, num_heads, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        if d_model % num_heads != 0:
+            raise AssertionError(
+                "d_model (%d) must be divisible by num_heads (%d)"
+                % (d_model, num_heads)
+            )
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+
+        self.q_linear = nn.Linear(d_model, d_model)
+        self.kv_linear = nn.Linear(d_model, d_model * 2)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(d_model, d_model)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, cond, mask=None):
+        # query/value: img tokens; key: condition; mask: if padding tokens
+        B, N, C = x.shape
+
+        q = self.q_linear(x).view(-1, self.num_heads, self.head_dim)
+        kv = self.kv_linear(cond).view(-1, 2, self.num_heads, self.head_dim)
+        k, v = kv.unbind(1)
+
+        q = q.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        if mask is not None:
+            temp_mask = torch.ones(B, 1, q.size(-2), k.size(-2), dtype=torch.bool, device=q.device)
+            for i in range(B):
+                temp_mask[i, :, :, mask[i]:] = False
+            mask = temp_mask
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=mask,
+            dropout_p=self.attn_drop.p if self.training else 0.,
+        ).transpose(1, 2)
+
+        x = x.view(B, -1, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x

paintmind/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

paintmind/stage1/quantize.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+def l2norm(t):
+    return F.normalize(t, p = 2, dim = -1)
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_e, e_dim, beta=0.25, use_norm=True):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.use_norm = use_norm
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.normal_()
+
+    def forward(self, z):
+        if self.use_norm:
+            z = l2norm(z)
+        z_flattened = z.view(-1, self.e_dim)
+        if self.use_norm:
+            embedd_norm = l2norm(self.embedding.weight)
+        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+
+        if self.use_norm:
+            d = 2 - 2 * torch.einsum('bc, nc -> bn', z_flattened, embedd_norm)
+        else:
+            d = torch.sum(z_flattened**2, dim=1, keepdim=True) + \
+                torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+                torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight)
+
+        encoding_indices = torch.argmin(d, dim=1).view(*z.shape[:-1])
+        z_q = self.embedding(encoding_indices)
+        if self.use_norm:
+            z_q = l2norm(z_q)
+
+        # compute loss for embedding
+        loss = self.beta * torch.mean((z_q.detach()-z)**2) + torch.mean((z_q-z.detach())**2)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        return z_q, loss, encoding_indices
+
+    def decode_from_indice(self, indices):
+        z_q = self.embedding(indices)
+        if self.use_norm:
+            z_q = l2norm(z_q)
+        
+        return z_q
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters # [B, L, 2C], not [B, 2C, H, W]
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=2)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self):
+        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+                return 0.5 * torch.sum(torch.pow(self.mean, 2)
+                                       + self.var - 1.0 - self.logvar,
+                                       dim=[1, 2])
+            else:
+                return 0.5 * torch.sum(
+                    torch.pow(self.mean - other.mean, 2) / other.var
+                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
+                    dim=[1, 2])
+
+    def nll(self, sample, dims=[1,2]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(
+            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+            dim=dims)
+
+    def mode(self):
+        return self.mean

paintmind/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

paintmind/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

paintmind/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))

paintmind/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

paintmind/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)

paintmind/stage1/vision_transformers.py

@@ -0,0 +1,267 @@
+# Copyright (c) ByteDance, Inc. and its 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.
+
+"""
+Mostly copy-paste from DINO and timm library:
+https://github.com/facebookresearch/dino
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from functools import partial
+from paintmind.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, slot_through=True):
+        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 if slot_through else 0
+        self.slot_through = slot_through
+
+        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))
+        if self.slot_through:
+            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)
+        if self.slot_through:
+            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)
+
+        if self.slot_through:
+            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)
+
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_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)
+        if self.slot_through:
+            x = torch.cat((self.cls_token.expand(B, -1, -1), x, self.slot_embed.expand(B, -1, -1)), dim=1)
+        else:
+            x = torch.cat((self.cls_token.expand(B, -1, -1), x), 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 and self.slot_through:
+            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)
+        if self.slot_through:
+            outcome = x[:, -self.num_slots:] # return the slots
+        else:
+            outcome = x[:, 1:] # return the patches
+        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

paintmind/stage2/causaldit.py

@@ -0,0 +1,422 @@
+# 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.
+# --------------------------------------------------------
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from timm.models.vision_transformer import Mlp
+from paintmind.stage1.diffusion_transfomers import TimestepEmbedder, LabelEmbedder, FinalLayer, modulate
+from paintmind.stage1.diffusion import create_diffusion
+from paintmind.stage1.transport import create_transport, Sampler
+
+
+class GeneralizedCausalAttention(nn.Module):
+    def __init__(self, dim, num_heads, norm_layer=nn.LayerNorm):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.proj = nn.Linear(dim, dim)
+        self.q_norm = norm_layer(self.head_dim)
+        self.k_norm = norm_layer(self.head_dim)
+
+    def _forward_kv_cache(
+        self, 
+        x: torch.Tensor, 
+        layer_index: int,
+        kv_cache: dict,
+        update_kv_cache: bool = False,
+    ):
+        N, Lq = x.shape[:2]
+        qkv = self.qkv(x).reshape(N, Lq, 3, self.num_heads, self.head_dim)
+        q, curr_k, curr_v = qkv.permute(2, 0, 3, 1, 4).unbind(0) # N, nhead, Lq, dhead
+        q = self.q_norm(q)
+        curr_k = self.k_norm(curr_k)
+
+        if kv_cache[layer_index]["k"] is not None:
+            k = kv_cache[layer_index]["k"]
+            v = kv_cache[layer_index]["v"]
+            k = torch.cat((k, curr_k), dim=2)
+            v = torch.cat((v, curr_v), dim=2)
+        else:
+            k = curr_k
+            v = curr_v
+
+        if update_kv_cache:
+            kv_cache[layer_index]["k"] = k
+            kv_cache[layer_index]["v"] = v
+
+        return self._forward_sdpa(q, k, v, attn_mask=None)
+
+    def _forward(self, x, attn_mask):
+        N, L = x.shape[:2]
+        qkv = self.qkv(x).reshape(N, L, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) # N, nhead, L, dhead
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        return self._forward_sdpa(q, k, v, attn_mask)
+
+    def _forward_sdpa(self, q, k, v, attn_mask):
+        N, _, Lq, _ = q.shape
+        x = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+        )
+
+        x = x.transpose(1, 2).reshape(N, Lq, -1)
+        x = self.proj(x)
+        return x
+
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        attn_mask: Optional[torch.Tensor] = None, 
+        kv_cache: Optional[dict] = None,
+        layer_index: Optional[int] = None,
+        update_kv_cache: Optional[bool] = None,
+    ) -> torch.Tensor:
+        if kv_cache is not None:
+            return self._forward_kv_cache(
+                x, 
+                kv_cache=kv_cache, 
+                layer_index=layer_index, 
+                update_kv_cache=update_kv_cache
+            )
+        else:
+            return self._forward(x, attn_mask)
+
+
+class DiTBlock(nn.Module):
+    def __init__(
+        self, 
+        hidden_size, 
+        num_heads, 
+        mlp_ratio=4.0, 
+        norm_layer=nn.LayerNorm,
+        causal_fusion=False,
+        deep_supervision=False,
+        output_dim=None,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = GeneralizedCausalAttention(hidden_size, num_heads=num_heads, norm_layer=norm_layer)
+        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.causal_fusion = causal_fusion
+        if not causal_fusion:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+            )
+        self.deep_supervision = deep_supervision
+        if deep_supervision:
+            if not causal_fusion:
+                self.final_layer = FinalLayer(hidden_size, 1, output_dim)
+            else:
+                self.final_layer = nn.Sequential(
+                    nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6),
+                    nn.Linear(hidden_size, output_dim)
+                )
+
+    def forward_causal_dit(self, x, c, **kwargs):
+        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), **kwargs)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        if self.deep_supervision and self.training:
+            return x, self.final_layer(x, c)
+        else:
+            return x
+
+    def forward_causal_fusion(self, x, **kwargs):
+        x = x + self.attn(self.norm1(x), **kwargs)
+        x = x + self.mlp(self.norm2(x))
+        if self.deep_supervision and self.training:
+            return x, self.final_layer(x)
+        else:
+            return x
+    
+    def forward(self, x, c=None, **kwargs):
+        if self.causal_fusion:
+            return self.forward_causal_fusion(x, **kwargs)
+        else:
+            return self.forward_causal_dit(x, c, **kwargs)
+
+class CausalDiT(nn.Module):
+    def __init__(
+        self,
+        num_slots=16,
+        slot_dim=256,
+        hidden_size=1152,
+        depth=28,
+        num_heads=16,
+        mlp_ratio=4.0,
+        class_dropout_prob=0.1,
+        num_classes=1000,
+        num_sampling_steps='250',
+        use_si=False,
+        predict_xstart=False,
+        causal_fusion=False,
+        deep_supervision=False,
+        cls_token_num=0,
+        **kwargs
+    ):
+        super().__init__()
+        self.num_slots = num_slots
+        self.slot_dim = slot_dim
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.num_classes = num_classes
+        self.output_dim = slot_dim * 2 if not use_si else slot_dim
+        self.x_embedder = nn.Linear(slot_dim, hidden_size)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_slots, hidden_size))
+        blocks = [
+            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, causal_fusion=causal_fusion,
+                     deep_supervision=deep_supervision, output_dim=self.output_dim) for _ in range(depth - 1)
+        ]
+        blocks.append(DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, causal_fusion=causal_fusion))
+        self.blocks = nn.ModuleList(blocks)
+        self.cls_token_num = cls_token_num
+        self.causal_fusion = causal_fusion
+        self.deep_supervision = deep_supervision
+        if not causal_fusion:
+            self.final_layer = FinalLayer(hidden_size, 1, self.output_dim)
+        else:
+            self.cond_pos_embed = nn.Parameter(torch.zeros(1, cls_token_num, hidden_size))
+            self.final_layer = nn.Sequential(
+                nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6),
+                nn.Linear(hidden_size, self.output_dim)
+            )
+
+        self.initialize_weights()
+
+        self.use_si = use_si
+        if not use_si:
+            self.train_diffusion = create_diffusion(timestep_respacing="", predict_xstart=predict_xstart)
+            self.gen_diffusion = create_diffusion(timestep_respacing=num_sampling_steps, predict_xstart=predict_xstart)
+        else:
+            self.transport = create_transport()
+            self.sampler = Sampler(self.transport)
+
+
+    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)
+            if self.causal_fusion and isinstance(module, nn.LayerNorm):
+                if module.weight is not None:
+                    nn.init.constant_(module.weight, 1.0)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+        # Initialize pos_embed:
+        nn.init.normal_(self.pos_embed, std=0.02)
+        # Initialize label embedding table:
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        if self.causal_fusion:
+            nn.init.normal_(self.cond_pos_embed, 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 i, block in enumerate(self.blocks):
+            if not self.causal_fusion:
+                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+            if self.deep_supervision and i < len(self.blocks) - 1:
+                if not self.causal_fusion:
+                    nn.init.constant_(block.final_layer.adaLN_modulation[-1].weight, 0)
+                    nn.init.constant_(block.final_layer.adaLN_modulation[-1].bias, 0)
+                    nn.init.constant_(block.final_layer.linear.weight, 0)
+                    nn.init.constant_(block.final_layer.linear.bias, 0)
+                else:
+                    nn.init.constant_(block.final_layer[1].weight, 0)
+                    nn.init.constant_(block.final_layer[1].bias, 0)
+
+        # Zero-out output layers:
+        if not self.causal_fusion:
+            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)
+        else:
+            nn.init.constant_(self.final_layer[1].weight, 0)
+            nn.init.constant_(self.final_layer[1].bias, 0)
+
+    def forward_cache_update(self, x, kv_cache):
+        for idx, block in enumerate(self.blocks):
+            x = block(x, layer_index=idx, kv_cache=kv_cache, update_kv_cache=True)
+        return None
+
+    def _forward_inference(self, xn, t, y, kv_cache, pos_embed, attn_mask, context):
+        xn = xn.transpose(1, 2)
+        if context is not None:
+            xn = torch.cat([context, xn], dim=1)
+        xn = self.x_embedder(xn) + pos_embed
+        y = self.y_embedder(y, self.training)
+        t = self.t_embedder(t)
+        if not self.causal_fusion:
+            c = t + y
+        else:
+            y = y.unsqueeze(1).expand(-1, self.cls_token_num, -1) + self.cond_pos_embed
+            xn = torch.cat([y, xn], dim=1)
+            t = t.unsqueeze(1)
+            xn = xn + t
+            c = None
+
+        for idx, block in enumerate(self.blocks):
+            xn = block(xn, c, attn_mask=attn_mask, layer_index=idx, kv_cache=kv_cache, update_kv_cache=False)
+
+        xn = xn[:, -1].unsqueeze(1)
+        xn = self._forward_final_layer(xn, c)
+        return xn.transpose(1, 2)
+
+    def _forward_inference_with_cfg(self, xn, t, y, kv_cache, pos_embed, attn_mask, context, cfg_scale):
+        half = xn[: len(xn) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = self._forward_inference(combined, t, y, kv_cache, pos_embed, attn_mask, context)
+        eps, rest = model_out[:, :self.slot_dim], model_out[:, self.slot_dim:]
+        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)
+
+    def sample(self, y, pos_embed, attn_mask, context=None, cfg=1.0):
+        # diffusion loss sampling
+        device = y.device
+        if not cfg == 1.0:
+            noise = torch.randn(y.shape[0] // 2, self.slot_dim, 1, device=device)
+            noise = torch.cat([noise, noise], dim=0)
+            model_kwargs = dict(y=y, kv_cache=None, pos_embed=pos_embed, attn_mask=attn_mask, context=context, cfg_scale=cfg)
+            sample_fn = self._forward_inference_with_cfg
+        else:
+            noise = torch.randn(y.shape[0], self.slot_dim, 1, device=device)
+            model_kwargs = dict(y=y, kv_cache=None, pos_embed=pos_embed, attn_mask=attn_mask, context=context)
+            sample_fn = self._forward_inference
+
+        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,
+                device=device
+            )
+        else:
+            sde_sample_fn = self.sampler.sample_sde(diffusion_form="sigma")
+            sampled_token_latent = sde_sample_fn(noise, sample_fn, **model_kwargs)[-1]
+
+        return sampled_token_latent.transpose(1, 2)
+
+    def _forward_final_layer(self, x, c):
+        if not self.causal_fusion:
+            return self.final_layer(x, c)
+        else:
+            return self.final_layer(x)
+
+    def forward_train(self, xn, t, y, xc):
+        """
+        Args:
+            xn: noised latent
+            t: time step
+            y: condition
+            xc: clean latent
+        """
+
+        xc = self.x_embedder(xc.transpose(1, 2)) + self.pos_embed
+        xn = self.x_embedder(xn.transpose(1, 2)) + self.pos_embed
+
+        t = self.t_embedder(t)
+        y = self.y_embedder(y, self.training)
+        if not self.causal_fusion:
+            c = t + y
+        else:
+            y = y.unsqueeze(1).expand(-1, self.cls_token_num, -1) + self.cond_pos_embed
+            t = t.unsqueeze(1)
+            xc = torch.cat((y, xc), dim=1)
+            xn = xn + t
+            c = None
+
+        # forward transformer
+        x = torch.cat((xc, xn), dim=1)
+        attn_mask = get_attn_mask(self.cls_token_num, self.num_slots).to(x.device)
+
+        if self.deep_supervision and self.training:
+            xs = []
+            for block in self.blocks[:-1]:
+                x, x_hat = block(x, c=c, attn_mask=attn_mask)
+                xs.append(x_hat[:, -self.num_slots:])
+            x = self.blocks[-1](x, c=c, attn_mask=attn_mask)
+            x = self._forward_final_layer(x[:, -self.num_slots:], c)
+            xs.append(x)
+            # N, B, L, C -> B, C, L, N
+            return torch.stack(xs, dim=0).permute(1, 3, 2, 0)
+        else:
+            for block in self.blocks:
+                x = block(x, c=c, attn_mask=attn_mask)
+            x = x[:, -self.num_slots:]
+            return self._forward_final_layer(x, c).transpose(1, 2)
+    
+    def forward(self, slots, targets):
+        slots = slots.transpose(1, 2)
+        model_kwargs = dict(y=targets, xc=slots)
+        if not self.use_si:
+            t = torch.randint(0, self.train_diffusion.num_timesteps, (slots.shape[0],), device=slots.device)
+            loss_dict = self.train_diffusion.training_losses(self.forward_train, slots, t, model_kwargs)
+        else:
+            loss_dict = self.transport.training_losses(self.forward_train, slots, model_kwargs)
+        loss = loss_dict["loss"]
+        return loss.mean()
+
+
+def CausalDiT_L(**kwargs):
+    return CausalDiT(depth=24, hidden_size=1024, num_heads=16, **kwargs)
+
+
+def CausalDiT_XL(**kwargs):
+    return CausalDiT(depth=32, hidden_size=1280, num_heads=20, **kwargs)
+
+
+def CausalDiT_H(**kwargs):
+    return CausalDiT(depth=48, hidden_size=1408, num_heads=22, **kwargs)
+
+
+CausalDiT_models = {
+    "CausalDiT-L": CausalDiT_L, 
+    "CausalDiT-XL": CausalDiT_XL, 
+    "CausalDiT-H": CausalDiT_H
+}
+
+def get_attn_mask(context_len, sample_len):
+    padder_len = 1 if context_len <= 0 else context_len
+    seq_len = padder_len + sample_len * 2
+    attn_mask = torch.eye(seq_len, dtype=bool)
+    triangle = torch.ones(sample_len, sample_len, dtype=bool).tril()
+    attn_mask[-sample_len:, :sample_len] = triangle
+    if padder_len == 1:
+        return attn_mask[1:, 1:]
+    else:
+        return attn_mask
+
+# if context_len == 0
+# 100000
+# 010000
+# 001000
+# 000100
+# 100010
+# 110001

paintmind/stage2/diffloss.py

@@ -0,0 +1,314 @@
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+import math
+
+from paintmind.stage1.diffusion import create_diffusion
+from paintmind.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, deep_supervision=False, token_drop_prob=0.0, cond_method="adaln", decoupled_cfg=True):
+        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,
+            deep_supervision=deep_supervision,
+            token_drop_prob=token_drop_prob,
+            cond_method=cond_method,
+            decoupled_cfg=decoupled_cfg,
+        )
+        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,
+        out_channels=None,
+        deep_supervision=False,
+        cond_method="adaln",
+    ):
+        super().__init__()
+        self.channels = channels
+        self.deep_supervision = deep_supervision
+        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),
+        )
+
+        if self.deep_supervision:
+            self.final_layer = FinalLayer(channels, out_channels, cond_method=cond_method)
+
+    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
+            
+        if self.deep_supervision and self.training:
+            return x, self.final_layer(x, t, c)
+        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,
+        deep_supervision=False,
+        token_drop_prob=0.0,
+        cond_method="adaln",
+        decoupled_cfg=True,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.deep_supervision = deep_supervision
+        self.token_drop_prob = token_drop_prob
+        self.cond_method = cond_method
+        self.decoupled_cfg = decoupled_cfg
+        if decoupled_cfg and token_drop_prob > 0.0:
+            self.null_token = nn.Parameter(torch.zeros(1, z_channels))
+
+        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 = []
+        for i in range(num_res_blocks - 1):
+            res_blocks.append(ResBlock(model_channels, out_channels, deep_supervision, cond_method))
+        res_blocks.append(ResBlock(model_channels, cond_method=cond_method))
+        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)
+
+        if self.token_drop_prob > 0.0:
+            nn.init.normal_(self.null_token, 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)
+            if self.deep_supervision and i < len(self.res_blocks) - 1:
+                nn.init.constant_(block.final_layer.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(block.final_layer.adaLN_modulation[-1].bias, 0)
+                nn.init.constant_(block.final_layer.linear.weight, 0)
+                nn.init.constant_(block.final_layer.linear.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)
+        
+        # Apply token dropout if needed
+        if self.token_drop_prob > 0.0 and self.training:
+            drop_ids = torch.rand(c.shape[0], 1, device=c.device) < self.token_drop_prob
+            c = torch.where(drop_ids, self.null_token, c)
+        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
+
+        if self.deep_supervision and self.training:
+            xs = []
+            for block in self.res_blocks[:-1]:
+                x, x_hat = block(x, t_combined, c_for_concat)
+                xs.append(x_hat)
+            x = self.res_blocks[-1](x, t_combined, c_for_concat)
+            x = self.final_layer(x, t_combined, c_for_concat)
+            xs.append(x)
+            return torch.stack(xs, dim=-1)
+        else:
+            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)

paintmind/stage2/generate.py

@@ -0,0 +1,127 @@
+# 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, diff_cfg: float = 1.0, temperature: float = 1.0):
+    tokens = model(None, cond_idx, input_pos, cfg=cfg_scale, diff_cfg=diff_cfg, temperature=temperature)
+    return tokens.unsqueeze(1)
+
+
+def decode_one_token(model, x: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, diff_cfg: float = 1.0, 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, diff_cfg=diff_cfg, temperature=temperature)
+    return tokens
+
+
+def decode_n_tokens(
+    model, cur_token: torch.Tensor, input_pos: torch.Tensor, num_new_tokens: int, 
+    cfg_scale: float, diff_cfg: float = 1.0, temperature: float = 1.0, cfg_schedule = "constant", diff_cfg_schedule = "constant"):
+    new_tokens = []
+    for i in range(num_new_tokens):
+        cfg_iter = get_cfg(cfg_scale, i + 1, num_new_tokens + 1, cfg_schedule)
+        diff_cfg_iter = get_cfg(diff_cfg, i + 1, num_new_tokens + 1, diff_cfg_schedule)
+        next_token = decode_one_token(model, cur_token, input_pos, cfg_iter, diff_cfg=diff_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, diff_cfg: float = 1.0, temperature: float = 1.0, cfg_schedule = "constant", diff_cfg_schedule = "constant"):
+    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)
+    diff_cfg_iter = get_cfg(diff_cfg, 0, max_new_tokens, diff_cfg_schedule)
+    next_token = prefill(model, cond_combined, input_pos, cfg_iter, diff_cfg=diff_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, diff_cfg=diff_cfg, temperature=temperature, cfg_schedule=cfg_schedule, diff_cfg_schedule=diff_cfg_schedule)
+        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 == "inv_linear":
+        return 1 + (cfg - 1) * (total_step - cur_step - 1) / total_step
+    elif cfg_schedule == "constant":
+        return cfg
+    else:
+        raise NotImplementedError
+
+
+@torch.no_grad()
+def generate_causal_dit(model, cond, max_new_tokens, cfg_scale=1.0):
+    assert max_new_tokens == model.num_slots
+
+    batch_size = cond.shape[0]
+    device = cond.device
+
+    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
+    
+    cur_tokens = []
+    for i in range(max_new_tokens):
+        pos_embed = model.pos_embed[:, :i + 1].view(1, -1, model.hidden_size).expand(batch_size, -1, -1)
+        if cfg_scale > 1.0:
+            pos_embed = torch.cat([pos_embed, pos_embed], dim=0)
+
+        attn_mask = torch.ones(model.cls_token_num + i + 1, model.cls_token_num + i + 1, dtype=torch.bool).tril(diagonal=0).to(device)
+        
+        context = torch.cat(cur_tokens, dim=1) if len(cur_tokens) > 0 else None
+        if cfg_scale > 1.0 and context is not None:
+            context = torch.cat([context, context], dim=0)
+
+        next_token = model.sample(cond_combined, pos_embed, attn_mask, context, cfg_scale)
+        if cfg_scale > 1.0:
+            next_token, _ = next_token.chunk(2, dim=0)
+        cur_tokens.append(next_token.clone())
+    
+    seq = torch.cat(cur_tokens, dim=1)
+
+    return seq

paintmind/stage2/gpt.py

@@ -0,0 +1,451 @@
+# 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 dataclasses import dataclass
+from typing import Optional, List, Union
+
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from paintmind.stage1.vision_transformers import DropPath
+from paintmind.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,
+        deep_supervision: bool = False,
+        token_drop_prob: float = 0.0,
+        cond_method: str = "adaln",
+        decoupled_cfg: bool = True,
+        **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,
+            deep_supervision=deep_supervision,
+            token_drop_prob=token_drop_prob,
+            cond_method=cond_method,
+            decoupled_cfg=decoupled_cfg,
+        )
+        self.decoupled_cfg = decoupled_cfg
+        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,  # cond_idx_or_embed
+        input_pos:  Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        cfg: float = 1.0,
+        diff_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:
+            if self.decoupled_cfg:
+                h = self.forward_cfg(h[:, -1], cfg)
+                h = h + self.diffusion_pos_embed_learned[:, input_pos[-1] - self.cls_token_num + 1]
+                if diff_cfg > 1.0 and hasattr(self.diffloss.net, 'null_token'):
+                    null_token = self.diffloss.net.null_token.expand(h.shape[0], -1)
+                    h = torch.cat((h, null_token), dim=0)
+                else:
+                    diff_cfg = 1.0
+                next_tokens = self.diffloss.sample(h, temperature=temperature, cfg=diff_cfg)
+            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, 
+}

paintmind/utils/datasets.py

@@ -0,0 +1,77 @@
+import os
+import torch
+import torchvision
+import numpy as np
+import os.path as osp
+from glob import glob
+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(image_set, aug='randcrop', img_size=256):
+
+    t = []
+    if image_set == '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.Resize(img_size, interpolation=TF.InterpolationMode.BICUBIC, antialias=True))
+        t.append(TF.CenterCrop(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)

paintmind/utils/device_utils.py

@@ -0,0 +1,20 @@
+import gc
+import torch
+import importlib.util
+
+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")

paintmind/utils/logger.py

@@ -0,0 +1,170 @@
+from collections import defaultdict, deque
+import datetime
+import time
+import torch
+import torch.distributed as dist
+from paintmind.engine.misc import is_dist_avail_and_initialized, is_main_process
+from paintmind.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)))

paintmind/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

paintmind/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)        

paintmind/version.py

@@ -0,0 +1 @@
+__version__ = '0.0.0'

requirements.txt

@@ -0,0 +1,27 @@
+numpy==1.26.4
+sympy>=1.10
+accelerate
+datasets
+diffusers[torch]
+transformers
+safetensors
+smart_open
+dotwiz
+omegaconf
+tensorboard
+huggingface-hub
+einops
+lpips
+timm
+scipy
+scikit-learn
+scikit-image
+kornia
+torchtyping
+git+https://github.com/xwen99/torch-fidelity.git@master#egg=torch-fidelity
+open_clip_torch
+opencv-python-headless
+torchmetrics
+torchdiffeq
+lmdb
+triton==3.0.0

submitit_test.py

@@ -0,0 +1,290 @@
+# 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.
+# --------------------------------------------------------
+# A script to run multinode training with submitit.
+# --------------------------------------------------------
+
+import argparse
+import os.path as osp
+import submitit
+import itertools
+
+from omegaconf import OmegaConf
+from paintmind.engine.util import instantiate_from_config
+from paintmind.utils.device_utils import configure_compute_backend
+
+
+def parse_args():
+    parser = argparse.ArgumentParser("Submitit for accelerator training")
+    # Slurm configuration
+    parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node")
+    parser.add_argument("--nodes", default=1, type=int, help="Number of nodes to request")
+    parser.add_argument("--timeout", default=7000, type=int, help="Duration of the job, default 5 days")
+    parser.add_argument("--qos", default="normal", type=str, help="QOS to request")
+    parser.add_argument("--job_dir", default="", type=str, help="Job dir. Leave empty for automatic.")
+    parser.add_argument("--partition", default="h100-camera-train", type=str, help="Partition where to submit")
+    parser.add_argument("--exclude", default="", type=str, help="Exclude nodes from the partition")
+    parser.add_argument("--nodelist", default="", type=str, help="Nodelist to request")
+    parser.add_argument('--comment', default="", type=str, help="Comment to pass to scheduler")
+    
+    # Model and testing configuration
+    parser.add_argument('--model', type=str, nargs='+', default=[None], help="Path to model(s)")
+    parser.add_argument('--step', type=int, nargs='+', default=[250000], help="Step number(s)")
+    parser.add_argument('--cfg', type=str, default=None, help="Path to config file")
+    parser.add_argument('--dataset', type=str, default='imagenet', help="Dataset to use")
+    
+    # Legacy parameter (preserved for backward compatibility)
+    parser.add_argument('--cfg_value', type=float, nargs='+', default=[None],
+                       help='Legacy parameter for GPT classifier-free guidance scale')
+    
+    # CFG-related parameters - all with nargs='+' to support multiple values
+    parser.add_argument('--ae_cfg', type=float, nargs='+', default=[None], 
+                       help="Autoencoder classifier-free guidance scale")
+    parser.add_argument('--diff_cfg', type=float, nargs='+', default=[None], 
+                       help="Diffusion classifier-free guidance scale")
+    parser.add_argument('--cfg_schedule', type=str, nargs='+', default=[None], 
+                       help="CFG schedule type (e.g., constant, linear)")
+    parser.add_argument('--diff_cfg_schedule', type=str, nargs='+', default=[None], 
+                       help="Diffusion CFG schedule type (e.g., constant, inv_linear)")
+    parser.add_argument('--test_num_slots', type=int, nargs='+', default=[None], 
+                       help="Number of slots to use for inference")
+    parser.add_argument('--temperature', type=float, nargs='+', default=[None], 
+                       help="Temperature for sampling")
+    
+    return parser.parse_args()
+
+
+def load_config(model_path, cfg_path=None):
+    """Load configuration from file or model directory."""
+    if cfg_path is not None and osp.exists(cfg_path):
+        config_path = cfg_path
+    elif model_path and osp.exists(osp.join(model_path, 'config.yaml')):
+        config_path = osp.join(model_path, 'config.yaml')
+    else:
+        raise ValueError(f"No config file found at {model_path} or {cfg_path}")
+    
+    return OmegaConf.load(config_path)
+
+
+def setup_checkpoint_path(model_path, step, config):
+    """Set up the checkpoint path based on model and step."""
+    if model_path:
+        ckpt_path = osp.join(model_path, 'models', f'step{step}')
+        if not osp.exists(ckpt_path):
+            print(f"Skipping non-existent checkpoint: {ckpt_path}")
+            return None
+        if hasattr(config.trainer.params, 'model'):
+            config.trainer.params.model.params.ckpt_path = ckpt_path
+        else:
+            config.trainer.params.gpt_model.params.ckpt_path = ckpt_path
+    else:
+        result_folder = config.trainer.params.result_folder
+        ckpt_path = osp.join(result_folder, 'models', f'step{step}')
+        if hasattr(config.trainer.params, 'model'):
+            config.trainer.params.model.params.ckpt_path = ckpt_path
+        else:
+            config.trainer.params.gpt_model.params.ckpt_path = ckpt_path
+    
+    return ckpt_path
+
+
+def setup_test_config(config, use_coco=False):
+    """Set up common test configuration parameters."""
+    config.trainer.params.test_dataset = config.trainer.params.dataset
+    if not use_coco:
+        config.trainer.params.test_dataset.params.split = 'val'
+    else:
+        config.trainer.params.test_dataset.target = 'paintmind.utils.datasets.COCO'
+        config.trainer.params.test_dataset.params.root = './dataset/coco'
+        config.trainer.params.test_dataset.params.split = 'val2017'
+    config.trainer.params.test_only = True
+    config.trainer.params.compile = False
+    config.trainer.params.eval_fid = True
+    config.trainer.params.fid_stats = 'fid_stats/adm_in256_stats.npz'
+    if hasattr(config.trainer.params, 'model'):
+        config.trainer.params.model.params.num_sampling_steps = '250'
+    else:
+        config.trainer.params.ae_model.params.num_sampling_steps = '250'
+
+def apply_cfg_params(config, param_dict):
+    """Apply CFG-related parameters to the config."""
+    # Apply each parameter if it's not None
+    if param_dict.get('cfg_value') is not None:
+        config.trainer.params.cfg = param_dict['cfg_value']
+        print(f"Setting cfg to {param_dict['cfg_value']}")
+    
+    if param_dict.get('ae_cfg') is not None:
+        config.trainer.params.ae_cfg = param_dict['ae_cfg']
+        print(f"Setting ae_cfg to {param_dict['ae_cfg']}")
+        
+    if param_dict.get('diff_cfg') is not None:
+        config.trainer.params.diff_cfg = param_dict['diff_cfg']
+        print(f"Setting diff_cfg to {param_dict['diff_cfg']}")
+        
+    if param_dict.get('cfg_schedule') is not None:
+        config.trainer.params.cfg_schedule = param_dict['cfg_schedule']
+        print(f"Setting cfg_schedule to {param_dict['cfg_schedule']}")
+        
+    if param_dict.get('diff_cfg_schedule') is not None:
+        config.trainer.params.diff_cfg_schedule = param_dict['diff_cfg_schedule']
+        print(f"Setting diff_cfg_schedule to {param_dict['diff_cfg_schedule']}")
+        
+    if param_dict.get('test_num_slots') is not None:
+        config.trainer.params.test_num_slots = param_dict['test_num_slots']
+        print(f"Setting test_num_slots to {param_dict['test_num_slots']}")
+
+    if param_dict.get('temperature') is not None:
+        config.trainer.params.temperature = param_dict['temperature']
+        print(f"Setting temperature to {param_dict['temperature']}")
+
+
+def run_test(config):
+    """Instantiate trainer and run test."""
+    trainer = instantiate_from_config(config.trainer)
+    trainer.train()
+
+
+def generate_param_combinations(args):
+    """Generate all combinations of parameters from the provided arguments."""
+    # Create parameter grid for all combinations
+    param_grid = {
+        'cfg_value': [None] if args.cfg_value == [None] else args.cfg_value,
+        'ae_cfg': [None] if args.ae_cfg == [None] else args.ae_cfg,
+        'diff_cfg': [None] if args.diff_cfg == [None] else args.diff_cfg,
+        'cfg_schedule': [None] if args.cfg_schedule == [None] else args.cfg_schedule,
+        'diff_cfg_schedule': [None] if args.diff_cfg_schedule == [None] else args.diff_cfg_schedule,
+        'test_num_slots': [None] if args.test_num_slots == [None] else args.test_num_slots,
+        'temperature': [None] if args.temperature == [None] else args.temperature
+    }
+    
+    # Get all parameter names that have non-None values
+    active_params = [k for k, v in param_grid.items() if v != [None]]
+    
+    if not active_params:
+        # If no parameters are specified, yield a dict with all None values
+        yield {k: None for k in param_grid.keys()}
+        return
+    
+    # Generate all combinations of active parameters
+    active_values = [param_grid[k] for k in active_params]
+    for combination in itertools.product(*active_values):
+        param_dict = {k: None for k in param_grid.keys()}  # Start with all None
+        for i, param_name in enumerate(active_params):
+            param_dict[param_name] = combination[i]
+        yield param_dict
+
+
+class Trainer(object):
+    def __init__(self, args):
+        self.args = args
+
+    def __call__(self):
+        """Main entry point for the submitit job."""
+        self._setup_gpu_args()
+        configure_compute_backend()
+        self._run_tests()
+
+    def _run_tests(self):
+        """Run tests for all specified models and steps."""
+        for step in self.args.step:
+            for model in self.args.model:
+                print(f"Testing model: {model} at step: {step}")
+                
+                # Load configuration
+                config = load_config(model, self.args.cfg)
+                
+                # Setup checkpoint path
+                ckpt_path = setup_checkpoint_path(model, step, config)
+                if ckpt_path is None:
+                    continue
+                
+                use_coco = self.args.dataset == 'coco' or self.args.dataset == 'COCO'
+                # Setup test configuration
+                setup_test_config(config, use_coco)
+                
+                # Generate and apply all parameter combinations
+                for param_dict in generate_param_combinations(self.args):
+                    # Create a copy of the config for each parameter combination
+                    current_config = OmegaConf.create(OmegaConf.to_container(config, resolve=True))
+                    
+                    # Print parameter combination
+                    param_str = ", ".join([f"{k}={v}" for k, v in param_dict.items() if v is not None])
+                    print(f"Testing with parameters: {param_str}")
+                    
+                    # Apply parameters and run test
+                    apply_cfg_params(current_config, param_dict)
+                    run_test(current_config)
+
+    def _setup_gpu_args(self):
+        """Set up GPU and distributed environment variables."""
+        import submitit
+
+        print("Exporting PyTorch distributed environment variables")
+        dist_env = submitit.helpers.TorchDistributedEnvironment().export(set_cuda_visible_devices=False)
+        print(f"Master: {dist_env.master_addr}:{dist_env.master_port}")
+        print(f"Rank: {dist_env.rank}")
+        print(f"World size: {dist_env.world_size}")
+        print(f"Local rank: {dist_env.local_rank}")
+        print(f"Local world size: {dist_env.local_world_size}")
+
+        job_env = submitit.JobEnvironment()
+        self.args.output_dir = str(self.args.output_dir).replace("%j", str(job_env.job_id))
+        self.args.log_dir = self.args.output_dir
+        print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+
+
+def main():
+    """Main function to set up and submit the job."""
+    args = parse_args()
+
+    # Determine job directory
+    if args.cfg is not None and osp.exists(args.cfg):
+        config = OmegaConf.load(args.cfg)
+    elif osp.exists(osp.join(args.model[0], 'config.yaml')):
+        config = OmegaConf.load(osp.join(args.model[0], 'config.yaml'))
+    else:
+        raise ValueError(f"No config file found at {args.model[0]} or {args.cfg}")
+    
+    args.job_dir = config.trainer.params.result_folder
+
+    # Set up the executor
+    executor = submitit.AutoExecutor(folder=args.job_dir, slurm_max_num_timeout=30)
+
+    # Configure slurm parameters
+    slurm_kwargs = {
+        'slurm_signal_delay_s': 120,
+        'slurm_qos': args.qos
+    }
+    
+    if args.comment:
+        slurm_kwargs['slurm_comment'] = args.comment
+    if args.exclude:
+        slurm_kwargs['slurm_exclude'] = args.exclude
+    if args.nodelist:
+        slurm_kwargs['slurm_nodelist'] = args.nodelist
+
+    # Update executor parameters
+    executor.update_parameters(
+        gpus_per_node=args.ngpus,
+        tasks_per_node=args.ngpus,  # one task per GPU
+        nodes=args.nodes,
+        timeout_min=args.timeout,
+        slurm_partition=args.partition,
+        name="fid",
+        **slurm_kwargs
+    )
+
+    args.output_dir = args.job_dir
+
+    # Submit the job
+    trainer = Trainer(args)
+    job = executor.submit(trainer)
+
+    print("Submitted job_id:", job.job_id)
+
+
+if __name__ == "__main__":
+    main()

submitit_train.py

@@ -0,0 +1,148 @@
+# 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.
+# --------------------------------------------------------
+# A script to run multinode training with submitit.
+# --------------------------------------------------------
+
+import argparse
+import os
+import submitit
+
+from omegaconf import OmegaConf
+from paintmind.engine.util import instantiate_from_config
+from paintmind.utils.device_utils import configure_compute_backend
+
+def parse_args():
+    parser = argparse.ArgumentParser("Submitit for accelerator training")
+    parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node")
+    parser.add_argument("--nodes", default=2, type=int, help="Number of nodes to request")
+    parser.add_argument("--timeout", default=7000, type=int, help="Duration of the job, default 5 days")
+    parser.add_argument("--qos", default="normal", type=str, help="QOS to request")
+    parser.add_argument("--job_dir", default="", type=str, help="Job dir. Leave empty for automatic.")
+
+    parser.add_argument("--partition", default="h100-camera-train", type=str, help="Partition where to submit")
+    parser.add_argument("--exclude", default="", type=str, help="Exclude nodes from the partition")
+    parser.add_argument("--nodelist", default="", type=str, help="Nodelist to request")
+    parser.add_argument('--comment', default="", type=str, help="Comment to pass to scheduler")
+    parser.add_argument('--cfg', type=str, default='configs/dit_imagenet_400ep.yaml', help='accelerator configs')
+    return parser.parse_args()
+
+
+class Trainer(object):
+    def __init__(self, args, config):
+        self.args = args
+        self.config = config
+
+    def __call__(self):
+        self._setup_gpu_args()
+        configure_compute_backend()
+        trainer = instantiate_from_config(self.config.trainer)
+        trainer.train(self.config)
+
+    def checkpoint(self):
+        import os
+        import submitit
+
+        model_dir = os.path.join(self.args.output_dir, "models")
+        if os.path.exists(model_dir):
+            # Get all step folders
+            step_folders = [d for d in os.listdir(model_dir) if d.startswith("step")]
+            if step_folders:
+                # Extract step numbers and find max
+                steps = [int(f.replace("step", "")) for f in step_folders]
+                max_step = max(steps)
+                # Set ckpt path to the latest step folder
+                self.config.trainer.params.model.params.ckpt_path = os.path.join(model_dir, f"step{max_step}")
+        print("Requeuing ", self.args, self.config)
+        empty_trainer = type(self)(self.args, self.config)
+        return submitit.helpers.DelayedSubmission(empty_trainer)
+
+    def _setup_gpu_args(self):
+        import submitit
+
+         # print_env()
+        print("exporting PyTorch distributed environment variables")
+        dist_env = submitit.helpers.TorchDistributedEnvironment().export(set_cuda_visible_devices=False)
+        print(f"master: {dist_env.master_addr}:{dist_env.master_port}")
+        print(f"rank: {dist_env.rank}")
+        print(f"world size: {dist_env.world_size}")
+        print(f"local rank: {dist_env.local_rank}")
+        print(f"local world size: {dist_env.local_world_size}")
+        # print_env()
+
+        # os.environ["NCCL_DEBUG"] = "INFO"
+        os.environ["NCCL_P2P_DISABLE"] = "0"
+        os.environ["NCCL_IB_DISABLE"] = "0"
+
+        job_env = submitit.JobEnvironment()
+        self.args.output_dir = str(self.args.output_dir).replace("%j", str(job_env.job_id))
+        self.args.log_dir = self.args.output_dir
+        self.config.trainer.params.result_folder = self.args.output_dir
+        self.config.trainer.params.log_dir = os.path.join(self.args.output_dir, "logs")
+        # self.args.gpu = job_env.local_rank
+        # self.args.rank = job_env.global_rank
+        # self.args.world_size = job_env.num_tasks
+        print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")
+
+
+def main():
+    args = parse_args()
+    cfg_file = args.cfg
+    assert os.path.exists(cfg_file)
+    config = OmegaConf.load(cfg_file)
+
+    if config.trainer.params.result_folder is None:
+        if args.job_dir == "":
+            args.job_dir = "./output/%j"
+        
+        config.trainer.params.result_folder = args.job_dir
+        config.trainer.params.log_dir = os.path.join(args.job_dir, "logs")
+    else:
+        args.job_dir = config.trainer.params.result_folder
+
+    # Note that the folder will depend on the job_id, to easily track experiments
+    executor = submitit.AutoExecutor(folder=args.job_dir, slurm_max_num_timeout=30)
+
+    num_gpus_per_node = args.ngpus
+    nodes = args.nodes
+    timeout_min = args.timeout
+    qos = args.qos
+
+    partition = args.partition
+    kwargs = {}
+    if args.comment:
+        kwargs['slurm_comment'] = args.comment
+    if args.exclude:
+        kwargs["slurm_exclude"] = args.exclude
+    if args.nodelist:
+        kwargs["slurm_nodelist"] = args.nodelist
+
+    executor.update_parameters(
+        mem_gb=40 * num_gpus_per_node,
+        gpus_per_node=num_gpus_per_node,
+        tasks_per_node=num_gpus_per_node,  # one task per GPU
+        # cpus_per_task=16,
+        nodes=nodes,
+        timeout_min=timeout_min,  # max is 60 * 72
+        # Below are cluster dependent parameters
+        slurm_partition=partition,
+        slurm_signal_delay_s=120,
+        slurm_qos=qos,
+        **kwargs
+    )
+
+    executor.update_parameters(name="sar")
+
+    args.output_dir = args.job_dir
+
+    trainer = Trainer(args, config)
+    job = executor.submit(trainer)
+
+    print("Submitted job_id:", job.job_id)
+
+
+if __name__ == "__main__":
+    main()