Backup final base/rotation-layer and latest conv-layer experiments

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+experiments/bs256_lr1e-4/validation/base_ema_euler10_seed20260815.png filter=lfs diff=lfs merge=lfs -text
+experiments/bs256_lr1e-4/validation/conv-layer_ema_euler10_seed20260815.png filter=lfs diff=lfs merge=lfs -text
+experiments/bs256_lr1e-4/validation/rot-layer_ema_euler10_seed20260815.png filter=lfs diff=lfs merge=lfs -text

checkpoints/bs256_lr1e-4/base/4003200.pt

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

checkpoints/bs256_lr1e-4/conv-layer/1950000.pt

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

checkpoints/bs256_lr1e-4/rotation-layer/4003200.pt

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

code_snapshot/convert_imagenet_parquet_to_imagefolder.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""Convert the Hugging Face ImageNet-1K parquet export to ImageFolder."""
+
+import argparse
+import json
+import os
+from pathlib import Path
+
+import pyarrow.parquet as pq
+
+
+def load_classes(path: Path):
+    namespace = {}
+    exec(path.read_text(), namespace)
+    return namespace["IMAGENET2012_CLASSES"]
+
+
+def convert_split(src_dir: Path, dst_dir: Path, split: str, classes, wnids):
+    files = sorted((src_dir / "data").glob(f"{split}-*.parquet"))
+    if not files:
+        raise FileNotFoundError(f"No parquet files found for {split}")
+
+    counts = [0] * len(classes)
+    for parquet_path in files:
+        parquet = pq.ParquetFile(parquet_path)
+        for batch in parquet.iter_batches(batch_size=256, columns=["image", "label"]):
+            rows = batch.to_pylist()
+            for row in rows:
+                label = int(row["label"])
+                if not 0 <= label < len(classes):
+                    raise ValueError(f"Invalid {split} label {label} in {parquet_path}")
+                image = row["image"]
+                image_bytes = image["bytes"]
+                if image_bytes is None:
+                    raise ValueError(f"Missing image bytes in {parquet_path}")
+                original_name = os.path.basename(image.get("path") or "image.JPEG")
+                # The source paths are unique, but prefixing the label keeps the
+                # generated names deterministic even if a future export changes.
+                filename = f"{counts[label]:06d}_{original_name}"
+                output = dst_dir / split / wnids[label] / filename
+                output.parent.mkdir(parents=True, exist_ok=True)
+                if not output.exists():
+                    output.write_bytes(image_bytes)
+                counts[label] += 1
+        print(f"{split}: processed {parquet_path.name}", flush=True)
+    return sum(counts), counts
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("src", type=Path)
+    parser.add_argument("dst", type=Path)
+    parser.add_argument("classes", type=Path)
+    args = parser.parse_args()
+
+    classes = load_classes(args.classes)
+    wnids = list(classes.keys())
+    args.dst.mkdir(parents=True, exist_ok=True)
+    class_to_idx = {wnid: i for i, wnid in enumerate(wnids)}
+    (args.dst / "class_to_idx.json").write_text(
+        json.dumps(class_to_idx, indent=2) + "\n"
+    )
+    (args.dst / "classes.json").write_text(
+        json.dumps({wnid: name for wnid, name in classes.items()}, indent=2) + "\n"
+    )
+
+    summary = {}
+    for split in ("train", "validation"):
+        total, counts = convert_split(args.src, args.dst, split, classes, wnids)
+        summary[split] = {"total": total, "per_class": counts}
+    (args.dst / "conversion_summary.json").write_text(
+        json.dumps(summary, indent=2) + "\n"
+    )
+    print(json.dumps({k: v["total"] for k, v in summary.items()}))
+
+
+if __name__ == "__main__":
+    main()

code_snapshot/models.py

@@ -0,0 +1,369 @@
+# 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, Attention, Mlp
+
+
+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)
+        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 SiT Model                                #
+#################################################################################
+
+class SiTBlock(nn.Module):
+    """
+    A SiT 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):
+        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))
+        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 SiT.
+    """
+    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 SiT(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([
+            SiTBlock(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 SiT 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 SiT.
+        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)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        return x
+
+    def forward_with_cfg(self, x, t, y, cfg_scale):
+        """
+        Forward pass of SiT, but also batches the unconSiTional 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
+
+
+#################################################################################
+#                                   SiT Configs                                  #
+#################################################################################
+
+def SiT_XL_2(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_XL_4(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_XL_8(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_L_2(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_L_4(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_L_8(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_B_2(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs)
+
+def SiT_B_4(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs)
+
+def SiT_B_8(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs)
+
+def SiT_S_2(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs)
+
+def SiT_S_4(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs)
+
+def SiT_S_8(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs)
+
+
+SiT_models = {
+    'SiT-XL/2': SiT_XL_2,  'SiT-XL/4': SiT_XL_4,  'SiT-XL/8': SiT_XL_8,
+    'SiT-L/2':  SiT_L_2,   'SiT-L/4':  SiT_L_4,   'SiT-L/8':  SiT_L_8,
+    'SiT-B/2':  SiT_B_2,   'SiT-B/4':  SiT_B_4,   'SiT-B/8':  SiT_B_8,
+    'SiT-S/2':  SiT_S_2,   'SiT-S/4':  SiT_S_4,   'SiT-S/8':  SiT_S_8,
+}

code_snapshot/models_conv.py

@@ -0,0 +1,422 @@
+# 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 attention_conv import Attention, AttentionConv
+
+
+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)
+        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 SiT Model                                #
+#################################################################################
+class SiTBlock(nn.Module):
+    """
+    A SiT 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):
+        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))
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+class SiTBlockConv(nn.Module):
+    """
+    A SiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0,
+                 stride=2, kernel_size=2,
+                 **block_kwargs):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = AttentionConv(hidden_size, num_heads=num_heads, qkv_bias=True,
+                              stride=stride, kernel_size=kernel_size, **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):
+        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))
+        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 SiT.
+    """
+    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 SiT(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,
+        stride_list = [1,1,2,2,2,4,4,2,2,2,1,1],
+        kernel_size_list = [1,1,2,2,2,4,4,2,2,2,1,1]
+    ):
+        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([
+        #     SiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)
+        # ])
+
+        # --------- #
+        self.blocks = nn.ModuleList()
+        self.stride_list = stride_list
+        self.kernel_size_list = kernel_size_list
+        assert len(stride_list) == len(kernel_size_list), "stride_list and kernel_size_list must have the same length"
+        assert len(stride_list) == depth
+        print("Initializing LightningDiT with the following block configurations:")
+        print("---------------------------------------")
+        print("stride_list:", stride_list)
+        print("kernel_size_list:", kernel_size_list)
+        print("---------------------------------------")
+        for i in range(depth):
+            stride = stride_list[i]
+            kernel_size = kernel_size_list[i]
+            if stride == 1 and kernel_size == 1:
+                block = SiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio)
+                print(f"idx={i}: add full block")
+            else:
+                block = SiTBlockConv(hidden_size, num_heads, mlp_ratio=mlp_ratio,
+                                     stride=stride, kernel_size=kernel_size)
+                print(f"idx={i}: stride={stride}, kernel_size={kernel_size}")
+            self.blocks.append(block)
+        # --------- #
+
+        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 SiT 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 SiT.
+        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)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        return x
+
+    def forward_with_cfg(self, x, t, y, cfg_scale):
+        """
+        Forward pass of SiT, but also batches the unconSiTional 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
+
+
+#################################################################################
+#                                   SiT Configs                                  #
+#################################################################################
+
+def SiT_XL_2(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_XL_4(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_XL_8(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_L_2(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_L_4(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_L_8(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_B_2(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs)
+
+def SiT_B_4(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs)
+
+def SiT_B_8(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs)
+
+def SiT_S_2(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs)
+
+def SiT_S_4(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs)
+
+def SiT_S_8(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs)
+
+
+SiT_models = {
+    'SiT-XL/2': SiT_XL_2,  'SiT-XL/4': SiT_XL_4,  'SiT-XL/8': SiT_XL_8,
+    'SiT-L/2':  SiT_L_2,   'SiT-L/4':  SiT_L_4,   'SiT-L/8':  SiT_L_8,
+    'SiT-B/2':  SiT_B_2,   'SiT-B/4':  SiT_B_4,   'SiT-B/8':  SiT_B_8,
+    'SiT-S/2':  SiT_S_2,   'SiT-S/4':  SiT_S_4,   'SiT-S/8':  SiT_S_8,
+}

code_snapshot/models_rot_layer.py

@@ -0,0 +1,379 @@
+# 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 attention_has_causal import Attention
+
+def rot180(x):
+    # print("x.shape:", x.shape) #(128,256,1152)
+    B, N, C = x.shape
+    H = W = int(N ** 0.5)
+    x = x.reshape(B, H, W, C)
+    x = torch.rot90(x, k=2, dims=(2,1))
+    x = x.reshape(B, N, C)
+    return x
+
+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)
+        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 SiT Model                                #
+#################################################################################
+
+class SiTBlock(nn.Module):
+    """
+    A SiT 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):
+        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), is_causal=True)
+        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 SiT.
+    """
+    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 SiT(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([
+            SiTBlock(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 SiT 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 SiT.
+        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)
+            x = rot180(x)                      # (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)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        return x
+
+    def forward_with_cfg(self, x, t, y, cfg_scale):
+        """
+        Forward pass of SiT, but also batches the unconSiTional 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
+
+
+#################################################################################
+#                                   SiT Configs                                  #
+#################################################################################
+
+def SiT_XL_2(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_XL_4(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_XL_8(**kwargs):
+    return SiT(depth=28, hidden_size=1152, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_L_2(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=2, num_heads=16, **kwargs)
+
+def SiT_L_4(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=4, num_heads=16, **kwargs)
+
+def SiT_L_8(**kwargs):
+    return SiT(depth=24, hidden_size=1024, patch_size=8, num_heads=16, **kwargs)
+
+def SiT_B_2(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=2, num_heads=12, **kwargs)
+
+def SiT_B_4(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=4, num_heads=12, **kwargs)
+
+def SiT_B_8(**kwargs):
+    return SiT(depth=12, hidden_size=768, patch_size=8, num_heads=12, **kwargs)
+
+def SiT_S_2(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=2, num_heads=6, **kwargs)
+
+def SiT_S_4(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=4, num_heads=6, **kwargs)
+
+def SiT_S_8(**kwargs):
+    return SiT(depth=12, hidden_size=384, patch_size=8, num_heads=6, **kwargs)
+
+
+SiT_models = {
+    'SiT-XL/2': SiT_XL_2,  'SiT-XL/4': SiT_XL_4,  'SiT-XL/8': SiT_XL_8,
+    'SiT-L/2':  SiT_L_2,   'SiT-L/4':  SiT_L_4,   'SiT-L/8':  SiT_L_8,
+    'SiT-B/2':  SiT_B_2,   'SiT-B/4':  SiT_B_4,   'SiT-B/8':  SiT_B_8,
+    'SiT-S/2':  SiT_S_2,   'SiT-S/4':  SiT_S_4,   'SiT-S/8':  SiT_S_8,
+}

code_snapshot/plot_fid_training_curves.py

@@ -0,0 +1,266 @@
+#!/usr/bin/env python3
+"""Plot aligned CFG=1 50k PyTorch-FID histories for the SiT-S/2 runs."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import html
+import math
+import os
+from pathlib import Path
+
+from PIL import Image, ImageDraw, ImageFont
+
+
+STEPS_PER_EPOCH = 5004
+DEFAULT_BASE_HISTORY = Path(
+    "/home/nvidia/SiT-official-bs256/results-official-bs256-400k/"
+    "000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv"
+)
+DEFAULT_ROT_HISTORY = Path(
+    "/home/nvidia/SiT-rot-layer-bs256/results-200ep/"
+    "SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv"
+)
+DEFAULT_CONV_HISTORY = Path(
+    "/data/nvidia/SiT-conv-layer-bs256/results-800ep/"
+    "SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv"
+)
+
+SERIES = (
+    ("Base", "#2563eb"),
+    ("Rotation-layer", "#dc2626"),
+    ("Conv-layer", "#059669"),
+)
+
+
+def _read_history(path: Path) -> list[dict[str, object]]:
+    if not path.is_file():
+        return []
+    by_step: dict[int, dict[str, object]] = {}
+    with path.open(newline="") as handle:
+        for row in csv.DictReader(handle, delimiter="\t"):
+            if row.get("status") != "ok" or float(row.get("cfg", "nan")) != 1.0:
+                continue
+            step = int(row["step"])
+            by_step[step] = {
+                "step": step,
+                "epoch": step / STEPS_PER_EPOCH,
+                "fid": float(row["fid"]),
+                "cfg": float(row["cfg"]),
+                "num_png": int(row["num_png"]),
+                "checkpoint": row["checkpoint"],
+                "timestamp_utc": row["timestamp_utc"],
+                "source": str(path),
+            }
+    return [by_step[step] for step in sorted(by_step)]
+
+
+def _font(size: int, bold: bool = False) -> ImageFont.ImageFont:
+    name = "DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf"
+    path = Path("/usr/share/fonts/truetype/dejavu") / name
+    try:
+        return ImageFont.truetype(str(path), size=size)
+    except OSError:
+        return ImageFont.load_default()
+
+
+def _bounds(histories: dict[str, list[dict[str, object]]]) -> tuple[float, float]:
+    values = [float(row["fid"]) for rows in histories.values() for row in rows]
+    if not values:
+        return 0.0, 100.0
+    low = 5.0 * math.floor((min(values) - 2.0) / 5.0)
+    high = 5.0 * math.ceil((max(values) + 2.0) / 5.0)
+    return max(0.0, low), max(low + 5.0, high)
+
+
+def _write_csv(path: Path, histories: dict[str, list[dict[str, object]]]) -> None:
+    fields = (
+        "model", "step", "epoch", "fid", "cfg", "num_png",
+        "checkpoint", "timestamp_utc", "source",
+    )
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=fields)
+        writer.writeheader()
+        for model, _ in SERIES:
+            for row in histories.get(model, []):
+                writer.writerow({"model": model, **row})
+
+
+def _write_png(path: Path, histories: dict[str, list[dict[str, object]]]) -> None:
+    width, height = 1600, 900
+    left, top, right, bottom = 125, 150, 1540, 785
+    y_min, y_max = _bounds(histories)
+    image = Image.new("RGB", (width, height), "#f8fafc")
+    draw = ImageDraw.Draw(image)
+    draw.rounded_rectangle((55, 45, width - 55, height - 45), radius=22,
+                           fill="white", outline="#dbe3ee", width=2)
+    draw.text((left, 67), "SiT-S/2 ImageNet 256: FID vs Training Progress",
+              font=_font(34, bold=True), fill="#0f172a")
+    draw.text((left, 113),
+              "CFG=1 · PyTorch-FID · 50,176 images · global batch 256 · lr 1e-4 · lower is better",
+              font=_font(19), fill="#475569")
+
+    def px(epoch: float) -> float:
+        return left + (epoch / 800.0) * (right - left)
+
+    def py(fid: float) -> float:
+        return bottom - ((fid - y_min) / (y_max - y_min)) * (bottom - top)
+
+    draw.rectangle((left, top, right, bottom), fill="#ffffff", outline="#94a3b8", width=2)
+    for epoch in range(0, 801, 100):
+        x = px(epoch)
+        draw.line((x, top, x, bottom), fill="#e2e8f0", width=1)
+        label = str(epoch)
+        box = draw.textbbox((0, 0), label, font=_font(16))
+        draw.text((x - (box[2] - box[0]) / 2, bottom + 12), label,
+                  font=_font(16), fill="#475569")
+    y_tick = y_min
+    while y_tick <= y_max + 1e-9:
+        y = py(y_tick)
+        draw.line((left, y, right, y), fill="#e2e8f0", width=1)
+        label = f"{y_tick:g}"
+        box = draw.textbbox((0, 0), label, font=_font(16))
+        draw.text((left - 18 - (box[2] - box[0]), y - 10), label,
+                  font=_font(16), fill="#475569")
+        y_tick += 5.0
+
+    draw.text(((left + right) / 2 - 95, bottom + 54),
+              "Training epoch (5,004 steps/epoch)", font=_font(18), fill="#334155")
+    draw.text((28, (top + bottom) / 2 + 50), "PyTorch FID", font=_font(18),
+              fill="#334155")
+
+    for model, color in SERIES:
+        rows = histories.get(model, [])
+        if not rows:
+            continue
+        points = [(px(float(row["epoch"])), py(float(row["fid"]))) for row in rows]
+        if len(points) > 1:
+            draw.line(points, fill=color, width=5, joint="curve")
+        radius = 4 if len(rows) <= 20 else 3
+        for x, y in points:
+            draw.ellipse((x - radius, y - radius, x + radius, y + radius),
+                         fill="white", outline=color, width=3)
+
+    legend_x, legend_y = right - 445, top + 22
+    present = [(model, color, histories[model]) for model, color in SERIES if histories.get(model)]
+    legend_h = 26 + 44 * len(present)
+    draw.rounded_rectangle((legend_x, legend_y, right - 20, legend_y + legend_h), radius=12,
+                           fill="#ffffff", outline="#cbd5e1", width=2)
+    for index, (model, color, rows) in enumerate(present):
+        y = legend_y + 21 + index * 44
+        draw.line((legend_x + 20, y, legend_x + 68, y), fill=color, width=5)
+        draw.ellipse((legend_x + 40, y - 4, legend_x + 48, y + 4), fill="white",
+                     outline=color, width=2)
+        latest = rows[-1]
+        label = f"{model}  ·  latest {float(latest['fid']):.3f} @ {float(latest['epoch']):.1f} ep"
+        draw.text((legend_x + 82, y - 13), label, font=_font(17, bold=True), fill="#1e293b")
+
+    image.save(path, optimize=True)
+
+
+def _write_svg(path: Path, histories: dict[str, list[dict[str, object]]]) -> None:
+    width, height = 1600, 900
+    left, top, right, bottom = 125, 150, 1540, 785
+    y_min, y_max = _bounds(histories)
+
+    def px(epoch: float) -> float:
+        return left + (epoch / 800.0) * (right - left)
+
+    def py(fid: float) -> float:
+        return bottom - ((fid - y_min) / (y_max - y_min)) * (bottom - top)
+
+    lines = [
+        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
+        '<rect width="100%" height="100%" fill="#f8fafc"/>',
+        '<rect x="55" y="45" width="1490" height="810" rx="22" fill="white" stroke="#dbe3ee" stroke-width="2"/>',
+        '<text x="125" y="95" font-family="DejaVu Sans" font-size="34" font-weight="700" fill="#0f172a">SiT-S/2 ImageNet 256: FID vs Training Progress</text>',
+        '<text x="125" y="130" font-family="DejaVu Sans" font-size="19" fill="#475569">CFG=1 · PyTorch-FID · 50,176 images · global batch 256 · lr 1e-4 · lower is better</text>',
+        f'<rect x="{left}" y="{top}" width="{right-left}" height="{bottom-top}" fill="white" stroke="#94a3b8" stroke-width="2"/>',
+    ]
+    for epoch in range(0, 801, 100):
+        x = px(epoch)
+        lines.extend([
+            f'<line x1="{x:.2f}" y1="{top}" x2="{x:.2f}" y2="{bottom}" stroke="#e2e8f0"/>',
+            f'<text x="{x:.2f}" y="{bottom+32}" text-anchor="middle" font-family="DejaVu Sans" font-size="16" fill="#475569">{epoch}</text>',
+        ])
+    y_tick = y_min
+    while y_tick <= y_max + 1e-9:
+        y = py(y_tick)
+        lines.extend([
+            f'<line x1="{left}" y1="{y:.2f}" x2="{right}" y2="{y:.2f}" stroke="#e2e8f0"/>',
+            f'<text x="{left-18}" y="{y+6:.2f}" text-anchor="end" font-family="DejaVu Sans" font-size="16" fill="#475569">{y_tick:g}</text>',
+        ])
+        y_tick += 5.0
+    lines.extend([
+        f'<text x="{(left+right)/2:.2f}" y="{bottom+74}" text-anchor="middle" font-family="DejaVu Sans" font-size="18" fill="#334155">Training epoch (5,004 steps/epoch)</text>',
+        f'<text x="35" y="{(top+bottom)/2:.2f}" text-anchor="middle" transform="rotate(-90 35 {(top+bottom)/2:.2f})" font-family="DejaVu Sans" font-size="18" fill="#334155">PyTorch FID</text>',
+    ])
+    for model, color in SERIES:
+        rows = histories.get(model, [])
+        if not rows:
+            continue
+        points = " ".join(f"{px(float(row['epoch'])):.2f},{py(float(row['fid'])):.2f}" for row in rows)
+        lines.append(f'<polyline points="{points}" fill="none" stroke="{color}" stroke-width="5" stroke-linejoin="round" stroke-linecap="round"/>')
+        radius = 4 if len(rows) <= 20 else 3
+        for row in rows:
+            lines.append(f'<circle cx="{px(float(row["epoch"])):.2f}" cy="{py(float(row["fid"])):.2f}" r="{radius}" fill="white" stroke="{color}" stroke-width="3"/>')
+    present = [(model, color, histories[model]) for model, color in SERIES if histories.get(model)]
+    legend_x, legend_y = right - 445, top + 22
+    legend_h = 26 + 44 * len(present)
+    lines.append(f'<rect x="{legend_x}" y="{legend_y}" width="425" height="{legend_h}" rx="12" fill="white" stroke="#cbd5e1" stroke-width="2"/>')
+    for index, (model, color, rows) in enumerate(present):
+        y = legend_y + 21 + index * 44
+        latest = rows[-1]
+        label = html.escape(f"{model} · latest {float(latest['fid']):.3f} @ {float(latest['epoch']):.1f} ep")
+        lines.extend([
+            f'<line x1="{legend_x+20}" y1="{y}" x2="{legend_x+68}" y2="{y}" stroke="{color}" stroke-width="5"/>',
+            f'<circle cx="{legend_x+44}" cy="{y}" r="4" fill="white" stroke="{color}" stroke-width="2"/>',
+            f'<text x="{legend_x+82}" y="{y+6}" font-family="DejaVu Sans" font-size="17" font-weight="700" fill="#1e293b">{label}</text>',
+        ])
+    lines.append("</svg>")
+    path.write_text("\n".join(lines) + "\n")
+
+
+def generate_plot(
+    output_dir: str | os.PathLike[str],
+    base_history: str | os.PathLike[str] = DEFAULT_BASE_HISTORY,
+    rot_history: str | os.PathLike[str] = DEFAULT_ROT_HISTORY,
+    conv_history: str | os.PathLike[str] = DEFAULT_CONV_HISTORY,
+) -> dict[str, Path]:
+    output = Path(output_dir)
+    output.mkdir(parents=True, exist_ok=True)
+    histories = {
+        "Base": _read_history(Path(base_history)),
+        "Rotation-layer": _read_history(Path(rot_history)),
+        "Conv-layer": _read_history(Path(conv_history)),
+    }
+    if not any(histories.values()):
+        raise RuntimeError("No completed CFG=1 FID records were found")
+    paths = {
+        "csv": output / "fid_cfg1_50k_training_curves.csv",
+        "png": output / "fid_cfg1_50k_training_curves.png",
+        "svg": output / "fid_cfg1_50k_training_curves.svg",
+    }
+    _write_csv(paths["csv"], histories)
+    _write_png(paths["png"], histories)
+    _write_svg(paths["svg"], histories)
+    return paths
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--base-history", default=str(DEFAULT_BASE_HISTORY))
+    parser.add_argument("--rot-history", default=str(DEFAULT_ROT_HISTORY))
+    parser.add_argument("--conv-history", default=str(DEFAULT_CONV_HISTORY))
+    args = parser.parse_args()
+    paths = generate_plot(
+        args.output_dir, args.base_history, args.rot_history, args.conv_history
+    )
+    for kind, path in paths.items():
+        print(f"{kind.upper()}: {path}")
+
+
+if __name__ == "__main__":
+    main()

code_snapshot/run_train.sh

@@ -0,0 +1,13 @@
+torchrun \
+--nnodes=1 \
+--nproc_per_node=8 \
+train.py \
+--model SiT-S/2 \
+--epochs=800 \
+--data-path /home/nvidia/datasets/imagenet-1k/train \
+--wandb \
+--global-batch-size=1024 \
+--run-name 005-SiT-S-2-Linear-velocity-None \
+--ckpt /home/nvidia/SiT-Complementary/results/005-SiT-S-2-Linear-velocity-None/checkpoints/0550000.pt
+
+# batch_size x 4, lr x 2

code_snapshot/run_train_conv.sh

@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /home/nvidia/SiT-Complementary
+
+export WANDB_KEY
+WANDB_KEY="$(python -c 'import netrc; print(netrc.netrc().authenticators("api.wandb.ai")[2])')"
+
+# Match the base and rotation-layer runs while keeping a separate W&B run.
+export WANDB_MODE=offline
+export WANDB_DIR=/data/nvidia/SiT-conv-layer-bs256/wandb
+export SIT_FID_COMPARISON_OUTPUT_DIR=/home/nvidia/SiT-comparisons/bs256-lr1e-4-800ep
+mkdir -p "$WANDB_DIR"
+
+exec torchrun \
+  --nnodes=1 \
+  --nproc_per_node=8 \
+  train_conv.py \
+  --model SiT-S/2 \
+  --epochs 800 \
+  --data-path /home/nvidia/datasets/imagenet-1k/train \
+  --results-dir /data/nvidia/SiT-conv-layer-bs256/results-800ep \
+  --global-batch-size 256 \
+  --learning-rate 0.0001 \
+  --global-seed 0 \
+  --vae ema \
+  --num-workers 4 \
+  --log-every 100 \
+  --ckpt-every 50000 \
+  --sample-every 10000 \
+  --cfg-scale 4.0 \
+  --run-name SiT-S-2-ConvLayer-bs256-lr1e-4-800ep \
+  --fid-every-checkpoint \
+  --fid-every 250000 \
+  --fid-num-samples 50000 \
+  --fid-reference /home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz \
+  --fid-history /data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv \
+  --fid-per-proc-batch-size 64 \
+  --fid-inception-batch-size 128 \
+  --fid-num-workers 8 \
+  --fid-sampling-steps 250 \
+  --fid-seed 0 \
+  --fid-stop-consecutive-increases 3 \
+  --fid-stop-min-absolute-rise 0.25 \
+  --fid-stop-min-relative-rise 0.005 \
+  --wandb

code_snapshot/run_train_rot_layer.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /home/nvidia/SiT-Complementary
+
+export WANDB_KEY
+WANDB_KEY="$(python -c 'import netrc; print(netrc.netrc().authenticators("api.wandb.ai")[2])')"
+
+# Match the base run's low-overhead W&B recording setup. This is a new run
+# because its name (and therefore deterministic W&B run ID) is unique.
+export WANDB_MODE=offline
+export WANDB_DIR=/home/nvidia/SiT-rot-layer-bs256/wandb
+mkdir -p "$WANDB_DIR"
+
+exec torchrun \
+  --nnodes=1 \
+  --nproc_per_node=8 \
+  train_rot_layer.py \
+  --model SiT-S/2 \
+  --epochs 800 \
+  --data-path /home/nvidia/datasets/imagenet-1k/train \
+  --results-dir /home/nvidia/SiT-rot-layer-bs256/results-200ep \
+  --global-batch-size 256 \
+  --learning-rate 0.0001 \
+  --global-seed 0 \
+  --vae ema \
+  --num-workers 4 \
+  --log-every 100 \
+  --ckpt-every 50000 \
+  --sample-every 10000 \
+  --cfg-scale 4.0 \
+  --run-name SiT-S-2-RotLayer-bs256-lr1e-4-200ep \
+  --ckpt /home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000800.pt \
+  --fid-every-checkpoint \
+  --fid-every 250000 \
+  --fid-num-samples 50000 \
+  --fid-reference /home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz \
+  --fid-history /home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv \
+  --fid-per-proc-batch-size 64 \
+  --fid-inception-batch-size 128 \
+  --fid-num-workers 8 \
+  --fid-sampling-steps 250 \
+  --fid-seed 0 \
+  --fid-stop-consecutive-increases 3 \
+  --fid-stop-min-absolute-rise 0.25 \
+  --fid-stop-min-relative-rise 0.005 \
+  --wandb

code_snapshot/sample.py

@@ -0,0 +1,145 @@
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Sample new images from a pre-trained SiT.
+"""
+import torch
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+from torchvision.utils import save_image
+from diffusers.models import AutoencoderKL
+from download import find_model
+from models import SiT_models
+from train_utils import parse_ode_args, parse_sde_args, parse_transport_args
+from transport import create_transport, Sampler
+import argparse
+import sys
+from time import time
+
+
+def main(mode, args):
+    # Setup PyTorch:
+    torch.manual_seed(args.seed)
+    torch.set_grad_enabled(False)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if args.ckpt is None:
+        assert args.model == "SiT-XL/2", "Only SiT-XL/2 models are available for auto-download."
+        assert args.image_size in [256, 512]
+        assert args.num_classes == 1000
+        assert args.image_size == 256, "512x512 models are not yet available for auto-download." # remove this line when 512x512 models are available
+        learn_sigma = args.image_size == 256
+    else:
+        # train.py uses the model default learn_sigma=True.
+        learn_sigma = True
+
+    # Load model:
+    latent_size = args.image_size // 8
+    model = SiT_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes,
+        learn_sigma=learn_sigma,
+    ).to(device)
+    # Auto-download a pre-trained model or load a custom SiT checkpoint from train.py:
+    ckpt_path = args.ckpt or f"SiT-XL-2-{args.image_size}x{args.image_size}.pt"
+    state_dict = find_model(ckpt_path)
+    model.load_state_dict(state_dict)
+    model.eval()  # important!
+    transport = create_transport(
+        args.path_type,
+        args.prediction,
+        args.loss_weight,
+        args.train_eps,
+        args.sample_eps
+    )
+    sampler = Sampler(transport)
+    if mode == "ODE":
+        if args.likelihood:
+            assert args.cfg_scale == 1, "Likelihood is incompatible with guidance"
+            sample_fn = sampler.sample_ode_likelihood(
+                sampling_method=args.sampling_method,
+                num_steps=args.num_sampling_steps,
+                atol=args.atol,
+                rtol=args.rtol,
+            )
+        else:
+            sample_fn = sampler.sample_ode(
+                sampling_method=args.sampling_method,
+                num_steps=args.num_sampling_steps,
+                atol=args.atol,
+                rtol=args.rtol,
+                reverse=args.reverse
+            )
+            
+    elif mode == "SDE":
+        sample_fn = sampler.sample_sde(
+            sampling_method=args.sampling_method,
+            diffusion_form=args.diffusion_form,
+            diffusion_norm=args.diffusion_norm,
+            last_step=args.last_step,
+            last_step_size=args.last_step_size,
+            num_steps=args.num_sampling_steps,
+        )
+    
+
+    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+
+    # Labels to condition the model with (feel free to change):
+    class_labels = [207, 360, 387, 974, 88, 979, 417, 279]
+    
+    # Create sampling noise:
+    n = len(class_labels)
+    z = torch.randn(n, 4, latent_size, latent_size, device=device)
+    y = torch.tensor(class_labels, device=device)
+
+    # Setup classifier-free guidance:
+    z = torch.cat([z, z], 0)
+    y_null = torch.tensor([1000] * n, device=device)
+    y = torch.cat([y, y_null], 0)
+    model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
+
+    # Sample images:
+    start_time = time()
+    samples = sample_fn(z, model.forward_with_cfg, **model_kwargs)[-1]
+    samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+    samples = vae.decode(samples / 0.18215).sample
+    print(f"Sampling took {time() - start_time:.2f} seconds.")
+
+    # Save and display images:
+    save_image(samples, "sample.png", nrow=4, normalize=True, value_range=(-1, 1))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    if len(sys.argv) < 2:
+        print("Usage: program.py <mode> [options]")
+        sys.exit(1)
+    
+    mode = sys.argv[1]
+
+    assert mode[:2] != "--", "Usage: program.py <mode> [options]"
+    assert mode in ["ODE", "SDE"], "Invalid mode. Please choose 'ODE' or 'SDE'"
+    
+    parser.add_argument("--model", type=str, choices=list(SiT_models.keys()), default="SiT-XL/2")
+    parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="mse")
+    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--cfg-scale", type=float, default=4.0)
+    parser.add_argument("--num-sampling-steps", type=int, default=250)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--ckpt", type=str, default=None,
+                        help="Optional path to a SiT checkpoint (default: auto-download a pre-trained SiT-XL/2 model).")
+
+
+    parse_transport_args(parser)
+    if mode == "ODE":
+        parse_ode_args(parser)
+        # Further processing for ODE
+    elif mode == "SDE":
+        parse_sde_args(parser)
+        # Further processing for SDE
+    
+    args = parser.parse_known_args()[0]
+    main(mode, args)

code_snapshot/sample_ddp.py

@@ -0,0 +1,249 @@
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Samples a large number of images from a pre-trained SiT model using DDP.
+Subsequently saves a .npz file that can be used to compute FID and other
+evaluation metrics via the ADM repo: https://github.com/openai/guided-diffusion/tree/main/evaluations
+
+For a simple single-GPU/CPU sampling script, see sample.py.
+"""
+import importlib
+import os
+
+import torch
+import torch.distributed as dist
+
+MODEL_MODULE_NAME = os.environ.get("SIT_MODEL_MODULE", "models")
+model_module = importlib.import_module(MODEL_MODULE_NAME)
+SiT_models = model_module.SiT_models
+MODEL_IMPLEMENTATION_PATH = os.path.realpath(model_module.__file__)
+from download import find_model
+from transport import create_transport, Sampler
+from diffusers.models import AutoencoderKL
+from train_utils import parse_ode_args, parse_sde_args, parse_transport_args
+from tqdm import tqdm
+from PIL import Image
+import numpy as np
+import math
+import argparse
+import sys
+
+
+def create_npz_from_sample_folder(sample_dir, num=50_000):
+    """
+    Builds a single .npz file from a folder of .png samples.
+    """
+    samples = []
+    for i in tqdm(range(num), desc="Building .npz file from samples"):
+        sample_pil = Image.open(f"{sample_dir}/{i:06d}.png")
+        sample_np = np.asarray(sample_pil).astype(np.uint8)
+        samples.append(sample_np)
+    samples = np.stack(samples)
+    assert samples.shape == (num, samples.shape[1], samples.shape[2], 3)
+    npz_path = f"{sample_dir}.npz"
+    np.savez(npz_path, arr_0=samples)
+    print(f"Saved .npz file to {npz_path} [shape={samples.shape}].")
+    return npz_path
+
+
+def main(mode, args):
+    """
+    Run sampling.
+    """
+    torch.backends.cuda.matmul.allow_tf32 = args.tf32  # True: fast but may lead to some small numerical differences
+    assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage"
+    expected_model_module = os.environ.get("SIT_EXPECTED_MODEL_MODULE")
+    if expected_model_module and MODEL_MODULE_NAME != expected_model_module:
+        raise RuntimeError(
+            f"Expected model module {expected_model_module!r}, but loaded "
+            f"{MODEL_MODULE_NAME!r} from {MODEL_IMPLEMENTATION_PATH}"
+        )
+    torch.set_grad_enabled(False)
+
+    # Setup DDP:
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    device = rank % torch.cuda.device_count()
+    seed = args.global_seed * dist.get_world_size() + rank
+    torch.manual_seed(seed)
+    torch.cuda.set_device(device)
+    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
+    if rank == 0:
+        print(f"Model implementation: {MODEL_MODULE_NAME} ({MODEL_IMPLEMENTATION_PATH})")
+
+    if args.ckpt is None:
+        assert args.model == "SiT-XL/2", "Only SiT-XL/2 models are available for auto-download."
+        assert args.image_size in [256, 512]
+        assert args.num_classes == 1000
+        assert args.image_size == 256, "512x512 models are not yet available for auto-download." # remove this line when 512x512 models are available
+        learn_sigma = args.image_size == 256
+    else:
+        # train.py constructs custom checkpoints with the model default
+        # learn_sigma=True, so preserve that architecture for strict loading.
+        learn_sigma = True
+
+    # Load model:
+    latent_size = args.image_size // 8
+    model = SiT_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes,
+        learn_sigma=learn_sigma,
+    ).to(device)
+    # Auto-download a pre-trained model or load a custom SiT checkpoint from train.py:
+    ckpt_path = args.ckpt or f"SiT-XL-2-{args.image_size}x{args.image_size}.pt"
+    state_dict = find_model(ckpt_path)
+    model.load_state_dict(state_dict)
+    model.eval()  # important!
+    
+    
+    transport = create_transport(
+        args.path_type,
+        args.prediction,
+        args.loss_weight,
+        args.train_eps,
+        args.sample_eps
+    )
+    sampler = Sampler(transport)
+    if mode == "ODE":
+        if args.likelihood:
+            assert args.cfg_scale == 1, "Likelihood is incompatible with guidance"
+            sample_fn = sampler.sample_ode_likelihood(
+                sampling_method=args.sampling_method,
+                num_steps=args.num_sampling_steps,
+                atol=args.atol,
+                rtol=args.rtol,
+            )
+        else:
+            sample_fn = sampler.sample_ode(
+                sampling_method=args.sampling_method,
+                num_steps=args.num_sampling_steps,
+                atol=args.atol,
+                rtol=args.rtol,
+                reverse=args.reverse
+            )
+    elif mode == "SDE":
+        sample_fn = sampler.sample_sde(
+            sampling_method=args.sampling_method,
+            diffusion_form=args.diffusion_form,
+            diffusion_norm=args.diffusion_norm,
+            last_step=args.last_step,
+            last_step_size=args.last_step_size,
+            num_steps=args.num_sampling_steps,
+        )
+    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+    assert args.cfg_scale >= 1.0, "In almost all cases, cfg_scale be >= 1.0"
+    using_cfg = args.cfg_scale > 1.0
+
+    # Create folder to save samples:
+    model_string_name = args.model.replace("/", "-")
+    ckpt_string_name = os.path.basename(args.ckpt).replace(".pt", "") if args.ckpt else "pretrained"
+    if mode == "ODE":
+        folder_name = f"{model_string_name}-{ckpt_string_name}-" \
+                  f"cfg-{args.cfg_scale}-{args.per_proc_batch_size}-"\
+                  f"{mode}-{args.num_sampling_steps}-{args.sampling_method}"
+    elif mode == "SDE":
+        folder_name = f"{model_string_name}-{ckpt_string_name}-" \
+                    f"cfg-{args.cfg_scale}-{args.per_proc_batch_size}-"\
+                    f"{mode}-{args.num_sampling_steps}-{args.sampling_method}-"\
+                    f"{args.diffusion_form}-{args.last_step}-{args.last_step_size}"
+    sample_folder_dir = f"{args.sample_dir}/{folder_name}"
+    if rank == 0:
+        os.makedirs(sample_folder_dir, exist_ok=True)
+        print(f"Saving .png samples at {sample_folder_dir}")
+    dist.barrier()
+
+    # Figure out how many samples we need to generate on each GPU and how many iterations we need to run:
+    n = args.per_proc_batch_size
+    global_batch_size = n * dist.get_world_size()
+    # To make things evenly-divisible, we'll sample a bit more than we need and then discard the extra samples:
+    num_samples = len([name for name in os.listdir(sample_folder_dir) if (os.path.isfile(os.path.join(sample_folder_dir, name)) and ".png" in name)])
+    total_samples = int(math.ceil(args.num_fid_samples / global_batch_size) * global_batch_size)
+    if rank == 0:
+        print(f"Total number of images that will be sampled: {total_samples}")
+    assert total_samples % dist.get_world_size() == 0, "total_samples must be divisible by world_size"
+    samples_needed_this_gpu = int(total_samples // dist.get_world_size())
+    assert samples_needed_this_gpu % n == 0, "samples_needed_this_gpu must be divisible by the per-GPU batch size"
+    iterations = int(samples_needed_this_gpu // n)
+    done_iterations = int( int(num_samples // dist.get_world_size()) // n)
+    pbar = range(iterations)
+    pbar = tqdm(pbar) if rank == 0 else pbar
+    total = 0
+    
+    for i in pbar:
+        # Sample inputs:
+        z = torch.randn(n, model.in_channels, latent_size, latent_size, device=device)
+        y = torch.randint(0, args.num_classes, (n,), device=device)
+        
+        # Setup classifier-free guidance:
+        if using_cfg:
+            z = torch.cat([z, z], 0)
+            y_null = torch.tensor([1000] * n, device=device)
+            y = torch.cat([y, y_null], 0)
+            model_kwargs = dict(y=y, cfg_scale=args.cfg_scale)
+            model_fn = model.forward_with_cfg
+        else:
+            model_kwargs = dict(y=y)
+            model_fn = model.forward
+
+        samples = sample_fn(z, model_fn, **model_kwargs)[-1]
+        if using_cfg:
+            samples, _ = samples.chunk(2, dim=0)  # Remove null class samples
+
+        samples = vae.decode(samples / 0.18215).sample
+        samples = torch.clamp(127.5 * samples + 128.0, 0, 255).permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+
+        # Save samples to disk as individual .png files
+        for i, sample in enumerate(samples):
+            index = i * dist.get_world_size() + rank + total
+            Image.fromarray(sample).save(f"{sample_folder_dir}/{index:06d}.png")
+        total += global_batch_size
+        dist.barrier()
+
+    # Make sure all processes have finished saving their samples before attempting to convert to .npz
+    dist.barrier()
+    if rank == 0:
+        create_npz_from_sample_folder(sample_folder_dir, args.num_fid_samples)
+        print("Done.")
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+
+    if len(sys.argv) < 2:
+        print("Usage: program.py <mode> [options]")
+        sys.exit(1)
+    
+    mode = sys.argv[1]
+    
+    assert mode[:2] != "--", "Usage: program.py <mode> [options]"
+    assert mode in ["ODE", "SDE"], "Invalid mode. Please choose 'ODE' or 'SDE'"
+
+    parser.add_argument("--model", type=str, choices=list(SiT_models.keys()), default="SiT-XL/2")
+    parser.add_argument("--vae",  type=str, choices=["ema", "mse"], default="ema")
+    parser.add_argument("--sample-dir", type=str, default="samples")
+    parser.add_argument("--per-proc-batch-size", type=int, default=64)
+    parser.add_argument("--num-fid-samples", type=int, default=50_000)
+    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--cfg-scale",  type=float, default=1.0)
+    parser.add_argument("--num-sampling-steps", type=int, default=250)
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--tf32", action=argparse.BooleanOptionalAction, default=True,
+                        help="By default, use TF32 matmuls. This massively accelerates sampling on Ampere GPUs.")
+    parser.add_argument("--ckpt", type=str, default=None,
+                        help="Optional path to a SiT checkpoint (default: auto-download a pre-trained SiT-XL/2 model).")
+
+    parse_transport_args(parser)
+    if mode == "ODE":
+        parse_ode_args(parser)
+        # Further processing for ODE
+    elif mode == "SDE":
+        parse_sde_args(parser)
+        # Further processing for SDE
+
+    args = parser.parse_known_args()[0]
+    main(mode, args)

code_snapshot/server_changes.patch

@@ -0,0 +1,1424 @@
+diff --git a/README.md b/README.md
+index 00362d2..edcb40b 100644
+--- a/README.md
++++ b/README.md
+@@ -19,6 +19,14 @@ This repository contains:
+ * ⚡️ Pre-trained class-conditional SiT models trained on ImageNet 256x256
+ * 🛸 A SiT [training script](train.py) using PyTorch DDP
+ 
++## Experiment backup and resume
++
++The project-specific recovery procedure for the base, rotation-layer, and
++convolution-layer experiments is documented in [docs/RESUME_GUIDE.md](docs/RESUME_GUIDE.md).
++It starts from downloading gated ImageNet-1K, recreates the ImageFolder layout,
++downloads the backed-up checkpoints/results from Hugging Face, verifies a
++checkpoint by sampling, and resumes the matching model implementation.
++
+ ## Setup
+ 
+ First, download and set up the repo:
+@@ -166,4 +174,3 @@ versus 2.06 in the paper).
+ ## License
+ This project is under the MIT license. See [LICENSE](LICENSE.txt) for details.
+ 
+-
+diff --git a/run_train.sh b/run_train.sh
+index 1d30f42..0f1ac15 100644
+--- a/run_train.sh
++++ b/run_train.sh
+@@ -3,9 +3,11 @@ torchrun \
+ --nproc_per_node=8 \
+ train.py \
+ --model SiT-S/2 \
+---epochs=400 \
+---data-path /home/jiayou.zhang/hom/personal/imagenet_dataset/images/train \
++--epochs=800 \
++--data-path /home/nvidia/datasets/imagenet-1k/train \
+ --wandb \
+---global-batch-size=1024
++--global-batch-size=1024 \
++--run-name 005-SiT-S-2-Linear-velocity-None \
++--ckpt /home/nvidia/SiT-Complementary/results/005-SiT-S-2-Linear-velocity-None/checkpoints/0550000.pt
+ 
+ # batch_size x 4, lr x 2
+diff --git a/run_train_conv.sh b/run_train_conv.sh
+old mode 100644
+new mode 100755
+index a3fa493..98b68ee
+--- a/run_train_conv.sh
++++ b/run_train_conv.sh
+@@ -1,11 +1,46 @@
+-torchrun \
+---nnodes=1 \
+---nproc_per_node=8 \
+-train_conv.py \
+---model SiT-S/2 \
+---epochs=200 \
+---data-path /home/jiayou.zhang/hom/personal/imagenet_dataset/images/train \
+---wandb \
+---global-batch-size=1024
+-
+-# batch_size x 4, lr x 2
++#!/usr/bin/env bash
++set -euo pipefail
++
++cd /home/nvidia/SiT-Complementary
++
++export WANDB_KEY
++WANDB_KEY="$(python -c 'import netrc; print(netrc.netrc().authenticators("api.wandb.ai")[2])')"
++
++# Match the base and rotation-layer runs while keeping a separate W&B run.
++export WANDB_MODE=offline
++export WANDB_DIR=/data/nvidia/SiT-conv-layer-bs256/wandb
++export SIT_FID_COMPARISON_OUTPUT_DIR=/home/nvidia/SiT-comparisons/bs256-lr1e-4-800ep
++mkdir -p "$WANDB_DIR"
++
++exec torchrun \
++  --nnodes=1 \
++  --nproc_per_node=8 \
++  train_conv.py \
++  --model SiT-S/2 \
++  --epochs 800 \
++  --data-path /home/nvidia/datasets/imagenet-1k/train \
++  --results-dir /data/nvidia/SiT-conv-layer-bs256/results-800ep \
++  --global-batch-size 256 \
++  --learning-rate 0.0001 \
++  --global-seed 0 \
++  --vae ema \
++  --num-workers 4 \
++  --log-every 100 \
++  --ckpt-every 50000 \
++  --sample-every 10000 \
++  --cfg-scale 4.0 \
++  --run-name SiT-S-2-ConvLayer-bs256-lr1e-4-800ep \
++  --fid-every-checkpoint \
++  --fid-every 250000 \
++  --fid-num-samples 50000 \
++  --fid-reference /home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz \
++  --fid-history /data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv \
++  --fid-per-proc-batch-size 64 \
++  --fid-inception-batch-size 128 \
++  --fid-num-workers 8 \
++  --fid-sampling-steps 250 \
++  --fid-seed 0 \
++  --fid-stop-consecutive-increases 3 \
++  --fid-stop-min-absolute-rise 0.25 \
++  --fid-stop-min-relative-rise 0.005 \
++  --wandb
+diff --git a/run_train_rot_layer.sh b/run_train_rot_layer.sh
+old mode 100644
+new mode 100755
+index e5bc80a..7396223
+--- a/run_train_rot_layer.sh
++++ b/run_train_rot_layer.sh
+@@ -1,11 +1,47 @@
+-torchrun \
+---nnodes=1 \
+---nproc_per_node=8 \
+-train_rot_layer.py \
+---model SiT-S/2 \
+---epochs=200 \
+---data-path /home/jiayou.zhang/hom/personal/imagenet_dataset/images/train \
+---wandb \
+---global-batch-size=1024
+-
+-# batch_size x 4, lr x 2
++#!/usr/bin/env bash
++set -euo pipefail
++
++cd /home/nvidia/SiT-Complementary
++
++export WANDB_KEY
++WANDB_KEY="$(python -c 'import netrc; print(netrc.netrc().authenticators("api.wandb.ai")[2])')"
++
++# Match the base run's low-overhead W&B recording setup. This is a new run
++# because its name (and therefore deterministic W&B run ID) is unique.
++export WANDB_MODE=offline
++export WANDB_DIR=/home/nvidia/SiT-rot-layer-bs256/wandb
++mkdir -p "$WANDB_DIR"
++
++exec torchrun \
++  --nnodes=1 \
++  --nproc_per_node=8 \
++  train_rot_layer.py \
++  --model SiT-S/2 \
++  --epochs 800 \
++  --data-path /home/nvidia/datasets/imagenet-1k/train \
++  --results-dir /home/nvidia/SiT-rot-layer-bs256/results-200ep \
++  --global-batch-size 256 \
++  --learning-rate 0.0001 \
++  --global-seed 0 \
++  --vae ema \
++  --num-workers 4 \
++  --log-every 100 \
++  --ckpt-every 50000 \
++  --sample-every 10000 \
++  --cfg-scale 4.0 \
++  --run-name SiT-S-2-RotLayer-bs256-lr1e-4-200ep \
++  --ckpt /home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000800.pt \
++  --fid-every-checkpoint \
++  --fid-every 250000 \
++  --fid-num-samples 50000 \
++  --fid-reference /home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz \
++  --fid-history /home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv \
++  --fid-per-proc-batch-size 64 \
++  --fid-inception-batch-size 128 \
++  --fid-num-workers 8 \
++  --fid-sampling-steps 250 \
++  --fid-seed 0 \
++  --fid-stop-consecutive-increases 3 \
++  --fid-stop-min-absolute-rise 0.25 \
++  --fid-stop-min-relative-rise 0.005 \
++  --wandb
+diff --git a/sample.py b/sample.py
+index 8bd86b5..2abb502 100644
+--- a/sample.py
++++ b/sample.py
+@@ -31,7 +31,8 @@ def main(mode, args):
+         assert args.image_size == 256, "512x512 models are not yet available for auto-download." # remove this line when 512x512 models are available
+         learn_sigma = args.image_size == 256
+     else:
+-        learn_sigma = False
++        # train.py uses the model default learn_sigma=True.
++        learn_sigma = True
+ 
+     # Load model:
+     latent_size = args.image_size // 8
+diff --git a/sample_ddp.py b/sample_ddp.py
+index 346b846..20d8202 100644
+--- a/sample_ddp.py
++++ b/sample_ddp.py
+@@ -8,15 +8,21 @@ evaluation metrics via the ADM repo: https://github.com/openai/guided-diffusion/
+ 
+ For a simple single-GPU/CPU sampling script, see sample.py.
+ """
++import importlib
++import os
++
+ import torch
+ import torch.distributed as dist
+-from models_rot_head import SiT_models
++
++MODEL_MODULE_NAME = os.environ.get("SIT_MODEL_MODULE", "models")
++model_module = importlib.import_module(MODEL_MODULE_NAME)
++SiT_models = model_module.SiT_models
++MODEL_IMPLEMENTATION_PATH = os.path.realpath(model_module.__file__)
+ from download import find_model
+ from transport import create_transport, Sampler
+ from diffusers.models import AutoencoderKL
+ from train_utils import parse_ode_args, parse_sde_args, parse_transport_args
+ from tqdm import tqdm
+-import os
+ from PIL import Image
+ import numpy as np
+ import math
+@@ -47,6 +53,12 @@ def main(mode, args):
+     """
+     torch.backends.cuda.matmul.allow_tf32 = args.tf32  # True: fast but may lead to some small numerical differences
+     assert torch.cuda.is_available(), "Sampling with DDP requires at least one GPU. sample.py supports CPU-only usage"
++    expected_model_module = os.environ.get("SIT_EXPECTED_MODEL_MODULE")
++    if expected_model_module and MODEL_MODULE_NAME != expected_model_module:
++        raise RuntimeError(
++            f"Expected model module {expected_model_module!r}, but loaded "
++            f"{MODEL_MODULE_NAME!r} from {MODEL_IMPLEMENTATION_PATH}"
++        )
+     torch.set_grad_enabled(False)
+ 
+     # Setup DDP:
+@@ -57,6 +69,8 @@ def main(mode, args):
+     torch.manual_seed(seed)
+     torch.cuda.set_device(device)
+     print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
++    if rank == 0:
++        print(f"Model implementation: {MODEL_MODULE_NAME} ({MODEL_IMPLEMENTATION_PATH})")
+ 
+     if args.ckpt is None:
+         assert args.model == "SiT-XL/2", "Only SiT-XL/2 models are available for auto-download."
+@@ -65,6 +79,8 @@ def main(mode, args):
+         assert args.image_size == 256, "512x512 models are not yet available for auto-download." # remove this line when 512x512 models are available
+         learn_sigma = args.image_size == 256
+     else:
++        # train.py constructs custom checkpoints with the model default
++        # learn_sigma=True, so preserve that architecture for strict loading.
+         learn_sigma = True
+ 
+     # Load model:
+@@ -123,7 +139,7 @@ def main(mode, args):
+     model_string_name = args.model.replace("/", "-")
+     ckpt_string_name = os.path.basename(args.ckpt).replace(".pt", "") if args.ckpt else "pretrained"
+     if mode == "ODE":
+-        folder_name = f"{model_string_name}-rot-head-{ckpt_string_name}-" \
++        folder_name = f"{model_string_name}-{ckpt_string_name}-" \
+                   f"cfg-{args.cfg_scale}-{args.per_proc_batch_size}-"\
+                   f"{mode}-{args.num_sampling_steps}-{args.sampling_method}"
+     elif mode == "SDE":
+diff --git a/train.py b/train.py
+index ece1e65..a24028f 100644
+--- a/train.py
++++ b/train.py
+@@ -21,10 +21,19 @@ from copy import deepcopy
+ from glob import glob
+ from time import time
+ import argparse
++import csv
++import importlib
+ import logging
++import math
+ import os
+-
+-from models import SiT_models
++import re
++import shutil
++from itertools import islice
++
++MODEL_MODULE_NAME = os.environ.get("SIT_MODEL_MODULE", "models")
++model_module = importlib.import_module(MODEL_MODULE_NAME)
++SiT_models = model_module.SiT_models
++MODEL_IMPLEMENTATION_PATH = os.path.realpath(model_module.__file__)
+ from download import find_model
+ from transport import create_transport, Sampler
+ from diffusers.models import AutoencoderKL
+@@ -82,6 +91,20 @@ def create_logger(logging_dir):
+     return logger
+ 
+ 
++class SkipBatchSampler:
++    """Skip already-consumed batches without loading or transforming their images."""
++
++    def __init__(self, batch_sampler, skip):
++        self.batch_sampler = batch_sampler
++        self.skip = skip
++
++    def __iter__(self):
++        return islice(iter(self.batch_sampler), self.skip, None)
++
++    def __len__(self):
++        return max(0, len(self.batch_sampler) - self.skip)
++
++
+ def center_crop_arr(pil_image, image_size):
+     """
+     Center cropping implementation from ADM.
+@@ -103,6 +126,177 @@ def center_crop_arr(pil_image, image_size):
+     return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+ 
+ 
++@torch.no_grad()
++def evaluate_checkpoint_fid(ema, vae, transport_sampler, args, train_steps,
++                            device, rank, logger, experiment_dir):
++    """Evaluate EMA with CFG=1 while preserving the training RNG trajectory."""
++    world_size = dist.get_world_size()
++    local_batch = args.fid_per_proc_batch_size
++    global_batch = local_batch * world_size
++    total_samples = math.ceil(args.fid_num_samples / global_batch) * global_batch
++    history_path = args.fid_history or os.path.join(experiment_dir, "fid_cfg1_50k.tsv")
++    sample_dir = os.path.join(
++        experiment_dir, "fid_cfg1_work", f"{train_steps:07d}"
++    )
++
++    # A completed record is reusable after a restart. Rank 0 decides and tells
++    # every worker, so all ranks take the same collective path.
++    already_done = False
++    if rank == 0 and os.path.isfile(history_path):
++        with open(history_path, newline="") as f:
++            for row in csv.DictReader(f, delimiter="\t"):
++                if int(row["step"]) == train_steps and row["status"] == "ok":
++                    already_done = True
++                    break
++    done_tensor = torch.tensor(int(already_done), device=device)
++    dist.broadcast(done_tensor, src=0)
++    if done_tensor.item():
++        logger.info(f"Reusing recorded CFG=1 FID for checkpoint {train_steps:07d}")
++        return False
++
++    cpu_rng_state = torch.get_rng_state()
++    cuda_rng_state = torch.cuda.get_rng_state(device)
++    torch.manual_seed(args.fid_seed * world_size + rank)
++    torch.cuda.manual_seed(args.fid_seed * world_size + rank)
++
++    if rank == 0:
++        os.makedirs(sample_dir, exist_ok=True)
++        # A prior interrupted attempt may contain a partial sample set.
++        for name in os.listdir(sample_dir):
++            if name.endswith(".png"):
++                os.remove(os.path.join(sample_dir, name))
++        logger.info(
++            f"Evaluating checkpoint {train_steps:07d}: CFG=1, "
++            f"requested={args.fid_num_samples:,}, actual={total_samples:,}"
++        )
++    dist.barrier()
++
++    sample_fn = transport_sampler.sample_ode(num_steps=args.fid_sampling_steps)
++    latent_size = args.image_size // 8
++    iterations = total_samples // global_batch
++    for batch_index in range(iterations):
++        z = torch.randn(local_batch, 4, latent_size, latent_size, device=device)
++        y = torch.randint(0, args.num_classes, (local_batch,), device=device)
++        samples = sample_fn(z, ema.forward, y=y)[-1]
++        samples = vae.decode(samples / 0.18215).sample
++        samples = torch.clamp(127.5 * samples + 128.0, 0, 255)
++        samples = samples.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
++        for local_index, sample in enumerate(samples):
++            image_index = batch_index * global_batch + local_index * world_size + rank
++            Image.fromarray(sample).save(os.path.join(sample_dir, f"{image_index:06d}.png"))
++        if batch_index % 10 == 0:
++            dist.barrier()
++    dist.barrier()
++
++    fid_value = 0.0
++    stop_requested = False
++    previous_step = None
++    previous_fid = None
++    if rank == 0:
++        from pytorch_fid.fid_score import calculate_fid_given_paths
++
++        fid_value = float(calculate_fid_given_paths(
++            [args.fid_reference, sample_dir],
++            batch_size=args.fid_inception_batch_size,
++            device="cuda:0",
++            dims=2048,
++            num_workers=args.fid_num_workers,
++        ))
++
++        prior_rows = []
++        if os.path.isfile(history_path):
++            with open(history_path, newline="") as f:
++                prior_rows = [
++                    row for row in csv.DictReader(f, delimiter="\t")
++                    if row["status"] == "ok" and int(row["step"]) < train_steps
++                ]
++        trend_rows = sorted(
++            (
++                (int(row["step"]), float(row["fid"]))
++                for row in prior_rows
++            ),
++            key=lambda item: item[0],
++        )
++        trend_rows.append((train_steps, fid_value))
++        required_points = args.fid_stop_consecutive_increases + 1
++        recent_trend = trend_rows[-required_points:]
++        if prior_rows:
++            previous_step, previous_fid = trend_rows[-2]
++        if len(recent_trend) == required_points:
++            consecutive_increases = all(
++                right_fid > left_fid
++                for (_, left_fid), (_, right_fid)
++                in zip(recent_trend, recent_trend[1:])
++            )
++            cumulative_rise = recent_trend[-1][1] - recent_trend[0][1]
++            required_rise = max(
++                args.fid_stop_min_absolute_rise,
++                recent_trend[0][1] * args.fid_stop_min_relative_rise,
++            )
++            stop_requested = consecutive_increases and cumulative_rise >= required_rise
++
++        os.makedirs(os.path.dirname(history_path), exist_ok=True)
++        needs_header = not os.path.isfile(history_path) or os.path.getsize(history_path) == 0
++        with open(history_path, "a", newline="") as f:
++            writer = csv.writer(f, delimiter="\t", lineterminator="\n")
++            if needs_header:
++                writer.writerow([
++                    "step", "checkpoint", "status", "fid", "cfg",
++                    "num_requested", "num_png", "seed", "timestamp_utc"
++                ])
++            from datetime import datetime, timezone
++            writer.writerow([
++                train_steps,
++                os.path.join(experiment_dir, "checkpoints", f"{train_steps:07d}.pt"),
++                "ok", repr(fid_value), "1.0", args.fid_num_samples,
++                total_samples, args.fid_seed,
++                datetime.now(timezone.utc).isoformat(),
++            ])
++
++        logger.info(f"Checkpoint {train_steps:07d} CFG=1 PyTorch FID: {fid_value:.9f}")
++        comparison_output_dir = os.environ.get("SIT_FID_COMPARISON_OUTPUT_DIR")
++        if comparison_output_dir:
++            try:
++                from tools.plot_fid_training_curves import generate_plot
++                generated = generate_plot(
++                    comparison_output_dir,
++                    conv_history=history_path,
++                )
++                logger.info(
++                    f"Updated FID comparison plot: {generated['png']}"
++                )
++            except Exception:
++                # A reporting artifact must never interrupt model training.
++                logger.exception("Could not update the FID comparison plot")
++        if args.wandb:
++            wandb_utils.log({"eval/fid_cfg1_50k": fid_value}, step=train_steps)
++        if stop_requested:
++            marker = os.path.join(experiment_dir, "FID_REGRESSION_STOPPED")
++            with open(marker, "w") as f:
++                f.write(
++                    f"sustained FID regression over {args.fid_stop_consecutive_increases} "
++                    f"consecutive checkpoints: step {recent_trend[0][0]} "
++                    f"FID {recent_trend[0][1]:.9f} -> step {train_steps} "
++                    f"FID {fid_value:.9f}\n"
++                )
++            logger.error(
++                f"FID increased for {args.fid_stop_consecutive_increases} consecutive "
++                f"checkpoints, from {recent_trend[0][1]:.9f} at step "
++                f"{recent_trend[0][0]} to {fid_value:.9f}; stopping after "
++                f"checkpoint {train_steps:07d}."
++            )
++        shutil.rmtree(sample_dir)
++
++    result = torch.tensor([fid_value, float(stop_requested)], device=device)
++    dist.broadcast(result, src=0)
++    dist.barrier()
++
++    # Evaluation must not perturb the random stream used by resumed training.
++    torch.set_rng_state(cpu_rng_state)
++    torch.cuda.set_rng_state(cuda_rng_state, device)
++    return bool(result[1].item())
++
++
+ #################################################################################
+ #                                  Training Loop                                #
+ #################################################################################
+@@ -112,6 +306,37 @@ def main(args):
+     Trains a new SiT model.
+     """
+     assert torch.cuda.is_available(), "Training currently requires at least one GPU."
++    expected_model_module = os.environ.get("SIT_EXPECTED_MODEL_MODULE")
++    if expected_model_module and MODEL_MODULE_NAME != expected_model_module:
++        raise RuntimeError(
++            f"Expected model module {expected_model_module!r}, but loaded "
++            f"{MODEL_MODULE_NAME!r} from {MODEL_IMPLEMENTATION_PATH}"
++        )
++
++    # Load resume metadata before creating the output directory or WandB run. The
++    # checkpoint hyperparameters remain authoritative; only runtime location,
++    # target epoch, and logging options may be overridden by the command line.
++    resume_checkpoint = None
++    resume_step = 0
++    if args.ckpt is not None:
++        runtime_args = args
++        resume_checkpoint = torch.load(args.ckpt, map_location="cpu", weights_only=False)
++        checkpoint_args = resume_checkpoint["args"]
++        runtime_names = (
++            "data_path", "results_dir", "epochs", "wandb", "ckpt", "run_name",
++            "fid_every_checkpoint", "fid_every", "fid_num_samples", "fid_reference",
++            "fid_history", "fid_per_proc_batch_size", "fid_inception_batch_size",
++            "fid_num_workers", "fid_sampling_steps", "fid_seed",
++            "fid_stop_consecutive_increases", "fid_stop_min_absolute_rise",
++            "fid_stop_min_relative_rise",
++        )
++        for name in runtime_names:
++            setattr(checkpoint_args, name, getattr(runtime_args, name))
++        args = checkpoint_args
++        match = re.fullmatch(r"(\d+)\.pt", os.path.basename(args.ckpt))
++        if match is None:
++            raise ValueError("Cannot infer the training step from checkpoint filename; expected NNNNNNN.pt")
++        resume_step = int(match.group(1))
+ 
+     # Setup DDP:
+     dist.init_process_group("nccl")
+@@ -129,13 +354,16 @@ def main(args):
+         os.makedirs(args.results_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+         experiment_index = len(glob(f"{args.results_dir}/*"))
+         model_string_name = args.model.replace("/", "-")  # e.g., SiT-XL/2 --> SiT-XL-2 (for naming folders)
+-        experiment_name = f"{experiment_index:03d}-{model_string_name}-" \
+-                        f"{args.path_type}-{args.prediction}-{args.loss_weight}"
++        experiment_name = args.run_name or (f"{experiment_index:03d}-{model_string_name}-" \
++                        f"{args.path_type}-{args.prediction}-{args.loss_weight}")
+         experiment_dir = f"{args.results_dir}/{experiment_name}"  # Create an experiment folder
+         checkpoint_dir = f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
+         os.makedirs(checkpoint_dir, exist_ok=True)
+         logger = create_logger(experiment_dir)
+         logger.info(f"Experiment directory created at {experiment_dir}")
++        logger.info(
++            f"Model implementation: {MODEL_MODULE_NAME} ({MODEL_IMPLEMENTATION_PATH})"
++        )
+ 
+         entity = os.environ["ENTITY"]
+         project = os.environ["PROJECT"]
+@@ -143,6 +371,10 @@ def main(args):
+             wandb_utils.initialize(args, entity, experiment_name, project)
+     else:
+         logger = create_logger(None)
++        experiment_dir = None
++    path_objects = [experiment_dir]
++    dist.broadcast_object_list(path_objects, src=0)
++    experiment_dir = path_objects[0]
+ 
+     # Create model:
+     assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
+@@ -155,14 +387,6 @@ def main(args):
+     # Note that parameter initialization is done within the SiT constructor
+     ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+ 
+-    if args.ckpt is not None:
+-        ckpt_path = args.ckpt
+-        state_dict = find_model(ckpt_path)
+-        model.load_state_dict(state_dict["model"])
+-        ema.load_state_dict(state_dict["ema"])
+-        opt.load_state_dict(state_dict["opt"])
+-        args = state_dict["args"]
+-
+     requires_grad(ema, False)
+     
+     model = DDP(model.to(device), device_ids=[device])
+@@ -178,7 +402,16 @@ def main(args):
+     logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+ 
+     # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+-    opt = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0)
++    opt = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0)
++    logger.info(
++        f"Optimizer: AdamW(lr={args.learning_rate:g}, weight_decay=0, "
++        "betas=(0.9, 0.999))"
++    )
++    if resume_checkpoint is not None:
++        model.module.load_state_dict(resume_checkpoint["model"])
++        ema.load_state_dict(resume_checkpoint["ema"])
++        opt.load_state_dict(resume_checkpoint["opt"])
++        logger.info(f"Resumed model, EMA, and optimizer from {args.ckpt} at step {resume_step:,}")
+ 
+     # Setup data:
+     transform = transforms.Compose([
+@@ -207,12 +440,13 @@ def main(args):
+     logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
+ 
+     # Prepare models for training:
+-    update_ema(ema, model.module, decay=0)  # Ensure EMA is initialized with synced weights
++    if resume_checkpoint is None:
++        update_ema(ema, model.module, decay=0)  # Initialize EMA only for a new run.
+     model.train()  # important! This enables embedding dropout for classifier-free guidance
+     ema.eval()  # EMA model should always be in eval mode
+ 
+     # Variables for monitoring/logging purposes:
+-    train_steps = 0
++    train_steps = resume_step
+     log_steps = 0
+     running_loss = 0
+     start_time = time()
+@@ -235,11 +469,38 @@ def main(args):
+         sample_model_kwargs = dict(y=ys)
+         model_fn = ema.forward
+ 
+-    logger.info(f"Training for {args.epochs} epochs...")
+-    for epoch in range(args.epochs):
++    steps_per_epoch = len(loader)
++    target_steps = args.epochs * steps_per_epoch
++    start_epoch, batches_to_skip = divmod(train_steps, steps_per_epoch)
++    logger.info(
++        f"Training to {args.epochs} total epochs ({target_steps:,} steps); "
++        f"starting at epoch {start_epoch}, batch {batches_to_skip}, step {train_steps:,}."
++    )
++    # Establish a same-protocol 50k baseline for the resume checkpoint before
++    # comparing later checkpoints against it. A recorded baseline is reused.
++    if args.fid_every_checkpoint and resume_checkpoint is not None:
++        if evaluate_checkpoint_fid(
++            ema, vae, transport_sampler, args, train_steps,
++            device, rank, logger, experiment_dir,
++        ):
++            logger.error("Resume checkpoint already violates the recorded FID trend; exiting.")
++            cleanup()
++            return
++        start_time = time()
++
++    stop_requested = False
++    for epoch in range(start_epoch, args.epochs):
+         sampler.set_epoch(epoch)
+         logger.info(f"Beginning epoch {epoch}...")
+-        for x, y in loader:
++        epoch_loader = loader
++        if epoch == start_epoch and batches_to_skip:
++            epoch_loader = DataLoader(
++                dataset,
++                batch_sampler=SkipBatchSampler(loader.batch_sampler, batches_to_skip),
++                num_workers=args.num_workers,
++                pin_memory=True,
++            )
++        for x, y in epoch_loader:
+             x = x.to(device)
+             y = y.to(device)
+             with torch.no_grad():
+@@ -290,6 +551,14 @@ def main(args):
+                     torch.save(checkpoint, checkpoint_path)
+                     logger.info(f"Saved checkpoint to {checkpoint_path}")
+                 dist.barrier()
++                if args.fid_every_checkpoint and train_steps % args.fid_every == 0:
++                    stop_requested = evaluate_checkpoint_fid(
++                        ema, vae, transport_sampler, args, train_steps,
++                        device, rank, logger, experiment_dir,
++                    )
++                    start_time = time()
++                    if stop_requested:
++                        break
+             
+             if train_steps % args.sample_every == 0 and train_steps > 0:
+                 logger.info("Generating EMA samples...")
+@@ -308,6 +577,21 @@ def main(args):
+                     wandb_utils.log_image(out_samples, train_steps)
+                 logging.info("Generating EMA samples done.")
+ 
++        if stop_requested:
++            break
++
++    if rank == 0 and not stop_requested and train_steps % args.ckpt_every != 0:
++        checkpoint = {
++            "model": model.module.state_dict(),
++            "ema": ema.state_dict(),
++            "opt": opt.state_dict(),
++            "args": args,
++        }
++        checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
++        torch.save(checkpoint, checkpoint_path)
++        logger.info(f"Saved final checkpoint to {checkpoint_path}")
++    dist.barrier()
++
+     model.eval()  # important! This disables randomized embedding dropout
+     # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+ 
+@@ -325,6 +609,7 @@ if __name__ == "__main__":
+     parser.add_argument("--num-classes", type=int, default=1000)
+     parser.add_argument("--epochs", type=int, default=1400)
+     parser.add_argument("--global-batch-size", type=int, default=256)
++    parser.add_argument("--learning-rate", type=float, default=1e-4)
+     parser.add_argument("--global-seed", type=int, default=0)
+     parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="ema")  # Choice doesn't affect training
+     parser.add_argument("--num-workers", type=int, default=4)
+@@ -335,6 +620,27 @@ if __name__ == "__main__":
+     parser.add_argument("--wandb", action="store_true")
+     parser.add_argument("--ckpt", type=str, default=None,
+                         help="Optional path to a custom SiT checkpoint")
++    parser.add_argument("--run-name", type=str, default=None,
++                        help="Experiment directory and WandB run name (useful when resuming)")
++    parser.add_argument("--fid-every-checkpoint", action="store_true",
++                        help="Run periodic CFG=1 PyTorch FID checks and stop on sustained regression")
++    parser.add_argument("--fid-every", type=int, default=50_000,
++                        help="Training-step interval between FID checks")
++    parser.add_argument("--fid-num-samples", type=int, default=50_000)
++    parser.add_argument("--fid-reference", type=str,
++                        default="/home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz")
++    parser.add_argument("--fid-history", type=str, default=None)
++    parser.add_argument("--fid-per-proc-batch-size", type=int, default=64)
++    parser.add_argument("--fid-inception-batch-size", type=int, default=128)
++    parser.add_argument("--fid-num-workers", type=int, default=8)
++    parser.add_argument("--fid-sampling-steps", type=int, default=250)
++    parser.add_argument("--fid-seed", type=int, default=0)
++    parser.add_argument("--fid-stop-consecutive-increases", type=int, default=3,
++                        help="Stop only after this many consecutive checkpoint FID increases")
++    parser.add_argument("--fid-stop-min-absolute-rise", type=float, default=0.25,
++                        help="Minimum cumulative absolute FID rise required to stop")
++    parser.add_argument("--fid-stop-min-relative-rise", type=float, default=0.005,
++                        help="Minimum cumulative relative FID rise required to stop")
+ 
+     parse_transport_args(parser)
+     args = parser.parse_args()
+diff --git a/train_conv.py b/train_conv.py
+old mode 100644
+new mode 100755
+index ee1640b..9989900
+--- a/train_conv.py
++++ b/train_conv.py
+@@ -1,341 +1,13 @@
+-# This source code is licensed under the license found in the
+-# LICENSE file in the root directory of this source tree.
++#!/usr/bin/env python3
++"""Run the shared SiT trainer with the convolutional-layer model implementation."""
+ 
+-"""
+-A minimal training script for SiT using PyTorch DDP.
+-"""
+-import torch
+-# the first flag below was False when we tested this script but True makes A100 training a lot faster:
+-torch.backends.cuda.matmul.allow_tf32 = True
+-torch.backends.cudnn.allow_tf32 = True
+-import torch.distributed as dist
+-from torch.nn.parallel import DistributedDataParallel as DDP
+-from torch.utils.data import DataLoader
+-from torch.utils.data.distributed import DistributedSampler
+-from torchvision.datasets import ImageFolder
+-from torchvision import transforms
+-import numpy as np
+-from collections import OrderedDict
+-from PIL import Image
+-from copy import deepcopy
+-from glob import glob
+-from time import time
+-import argparse
+-import logging
+ import os
++import runpy
+ 
+-from models_conv import SiT_models
+-from download import find_model
+-from transport import create_transport, Sampler
+-from diffusers.models import AutoencoderKL
+-from train_utils import parse_transport_args
+-import wandb_utils
+ 
++# Select models_conv.py before train.py imports the model registry, and require
++# the shared trainer to fail closed if another model module is selected.
++os.environ["SIT_MODEL_MODULE"] = "models_conv"
++os.environ["SIT_EXPECTED_MODEL_MODULE"] = "models_conv"
+ 
+-#################################################################################
+-#                             Training Helper Functions                         #
+-#################################################################################
+-
+-@torch.no_grad()
+-def update_ema(ema_model, model, decay=0.9999):
+-    """
+-    Step the EMA model towards the current model.
+-    """
+-    ema_params = OrderedDict(ema_model.named_parameters())
+-    model_params = OrderedDict(model.named_parameters())
+-
+-    for name, param in model_params.items():
+-        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+-        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+-
+-
+-def requires_grad(model, flag=True):
+-    """
+-    Set requires_grad flag for all parameters in a model.
+-    """
+-    for p in model.parameters():
+-        p.requires_grad = flag
+-
+-
+-def cleanup():
+-    """
+-    End DDP training.
+-    """
+-    dist.destroy_process_group()
+-
+-
+-def create_logger(logging_dir):
+-    """
+-    Create a logger that writes to a log file and stdout.
+-    """
+-    if dist.get_rank() == 0:  # real logger
+-        logging.basicConfig(
+-            level=logging.INFO,
+-            format='[\033[34m%(asctime)s\033[0m] %(message)s',
+-            datefmt='%Y-%m-%d %H:%M:%S',
+-            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+-        )
+-        logger = logging.getLogger(__name__)
+-    else:  # dummy logger (does nothing)
+-        logger = logging.getLogger(__name__)
+-        logger.addHandler(logging.NullHandler())
+-    return logger
+-
+-
+-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])
+-
+-
+-#################################################################################
+-#                                  Training Loop                                #
+-#################################################################################
+-
+-def main(args):
+-    """
+-    Trains a new SiT model.
+-    """
+-    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
+-
+-    # Setup DDP:
+-    dist.init_process_group("nccl")
+-    assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
+-    rank = dist.get_rank()
+-    device = rank % torch.cuda.device_count()
+-    seed = args.global_seed * dist.get_world_size() + rank
+-    torch.manual_seed(seed)
+-    torch.cuda.set_device(device)
+-    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
+-    local_batch_size = int(args.global_batch_size // dist.get_world_size())
+-
+-    # Setup an experiment folder:
+-    if rank == 0:
+-        os.makedirs(args.results_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+-        experiment_index = len(glob(f"{args.results_dir}/*"))
+-        model_string_name = args.model.replace("/", "-")  # e.g., SiT-XL/2 --> SiT-XL-2 (for naming folders)
+-        experiment_name = f"{experiment_index:03d}-{model_string_name}-conv-" \
+-                        f"{args.path_type}-{args.prediction}-{args.loss_weight}"
+-        experiment_dir = f"{args.results_dir}/{experiment_name}"  # Create an experiment folder
+-        checkpoint_dir = f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
+-        os.makedirs(checkpoint_dir, exist_ok=True)
+-        logger = create_logger(experiment_dir)
+-        logger.info(f"Experiment directory created at {experiment_dir}")
+-
+-        entity = os.environ["ENTITY"]
+-        project = os.environ["PROJECT"]
+-        if args.wandb:
+-            wandb_utils.initialize(args, entity, experiment_name, project)
+-    else:
+-        logger = create_logger(None)
+-
+-    # Create model:
+-    assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
+-    latent_size = args.image_size // 8
+-    model = SiT_models[args.model](
+-        input_size=latent_size,
+-        num_classes=args.num_classes
+-    )
+-
+-    # Note that parameter initialization is done within the SiT constructor
+-    ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+-
+-    if args.ckpt is not None:
+-        ckpt_path = args.ckpt
+-        state_dict = find_model(ckpt_path)
+-        model.load_state_dict(state_dict["model"])
+-        ema.load_state_dict(state_dict["ema"])
+-        opt.load_state_dict(state_dict["opt"])
+-        args = state_dict["args"]
+-
+-    requires_grad(ema, False)
+-    
+-    model = DDP(model.to(device), device_ids=[device])
+-    transport = create_transport(
+-        args.path_type,
+-        args.prediction,
+-        args.loss_weight,
+-        args.train_eps,
+-        args.sample_eps
+-    )  # default: velocity; 
+-    transport_sampler = Sampler(transport)
+-    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+-    logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+-
+-    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+-    opt = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0)
+-
+-    # Setup data:
+-    transform = transforms.Compose([
+-        transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
+-        transforms.RandomHorizontalFlip(),
+-        transforms.ToTensor(),
+-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+-    ])
+-    dataset = ImageFolder(args.data_path, transform=transform)
+-    sampler = DistributedSampler(
+-        dataset,
+-        num_replicas=dist.get_world_size(),
+-        rank=rank,
+-        shuffle=True,
+-        seed=args.global_seed
+-    )
+-    loader = DataLoader(
+-        dataset,
+-        batch_size=local_batch_size,
+-        shuffle=False,
+-        sampler=sampler,
+-        num_workers=args.num_workers,
+-        pin_memory=True,
+-        drop_last=True
+-    )
+-    logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
+-
+-    # Prepare models for training:
+-    update_ema(ema, model.module, decay=0)  # Ensure EMA is initialized with synced weights
+-    model.train()  # important! This enables embedding dropout for classifier-free guidance
+-    ema.eval()  # EMA model should always be in eval mode
+-
+-    # Variables for monitoring/logging purposes:
+-    train_steps = 0
+-    log_steps = 0
+-    running_loss = 0
+-    start_time = time()
+-
+-    # Labels to condition the model with (feel free to change):
+-    ys = torch.randint(1000, size=(local_batch_size,), device=device)
+-    use_cfg = args.cfg_scale > 1.0
+-    # Create sampling noise:
+-    n = ys.size(0)
+-    zs = torch.randn(n, 4, latent_size, latent_size, device=device)
+-
+-    # Setup classifier-free guidance:
+-    if use_cfg:
+-        zs = torch.cat([zs, zs], 0)
+-        y_null = torch.tensor([1000] * n, device=device)
+-        ys = torch.cat([ys, y_null], 0)
+-        sample_model_kwargs = dict(y=ys, cfg_scale=args.cfg_scale)
+-        model_fn = ema.forward_with_cfg
+-    else:
+-        sample_model_kwargs = dict(y=ys)
+-        model_fn = ema.forward
+-
+-    logger.info(f"Training for {args.epochs} epochs...")
+-    for epoch in range(args.epochs):
+-        sampler.set_epoch(epoch)
+-        logger.info(f"Beginning epoch {epoch}...")
+-        for x, y in loader:
+-            x = x.to(device)
+-            y = y.to(device)
+-            with torch.no_grad():
+-                # Map input images to latent space + normalize latents:
+-                x = vae.encode(x).latent_dist.sample().mul_(0.18215)
+-            model_kwargs = dict(y=y)
+-            loss_dict = transport.training_losses(model, x, model_kwargs)
+-            loss = loss_dict["loss"].mean()
+-            opt.zero_grad()
+-            loss.backward()
+-            opt.step()
+-            update_ema(ema, model.module)
+-
+-            # Log loss values:
+-            running_loss += loss.item()
+-            log_steps += 1
+-            train_steps += 1
+-            if train_steps % args.log_every == 0:
+-                # Measure training speed:
+-                torch.cuda.synchronize()
+-                end_time = time()
+-                steps_per_sec = log_steps / (end_time - start_time)
+-                # Reduce loss history over all processes:
+-                avg_loss = torch.tensor(running_loss / log_steps, device=device)
+-                dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+-                avg_loss = avg_loss.item() / dist.get_world_size()
+-                logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
+-                if args.wandb:
+-                    wandb_utils.log(
+-                        { "train loss": avg_loss, "train steps/sec": steps_per_sec },
+-                        step=train_steps
+-                    )
+-                # Reset monitoring variables:
+-                running_loss = 0
+-                log_steps = 0
+-                start_time = time()
+-
+-            # Save SiT checkpoint:
+-            if train_steps % args.ckpt_every == 0 and train_steps > 0:
+-                if rank == 0:
+-                    checkpoint = {
+-                        "model": model.module.state_dict(),
+-                        "ema": ema.state_dict(),
+-                        "opt": opt.state_dict(),
+-                        "args": args
+-                    }
+-                    checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+-                    torch.save(checkpoint, checkpoint_path)
+-                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+-                dist.barrier()
+-            
+-            if train_steps % args.sample_every == 0 and train_steps > 0:
+-                logger.info("Generating EMA samples...")
+-                with torch.no_grad():
+-                    sample_fn = transport_sampler.sample_ode() # default to ode sampling
+-                    samples = sample_fn(zs, model_fn, **sample_model_kwargs)[-1]
+-                    dist.barrier()
+-
+-                    if use_cfg: #remove null samples
+-                        samples, _ = samples.chunk(2, dim=0)
+-                    samples = vae.decode(samples / 0.18215).sample
+-                    out_samples = torch.zeros((args.global_batch_size, 3, args.image_size, args.image_size), device=device)
+-                    dist.all_gather_into_tensor(out_samples, samples)
+-
+-                if args.wandb:
+-                    wandb_utils.log_image(out_samples, train_steps)
+-                logging.info("Generating EMA samples done.")
+-
+-    model.eval()  # important! This disables randomized embedding dropout
+-    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+-
+-    logger.info("Done!")
+-    cleanup()
+-
+-
+-if __name__ == "__main__":
+-    # Default args here will train SiT-XL/2 with the hyperparameters we used in our paper (except training iters).
+-    parser = argparse.ArgumentParser()
+-    parser.add_argument("--data-path", type=str, required=True)
+-    parser.add_argument("--results-dir", type=str, default="results")
+-    parser.add_argument("--model", type=str, choices=list(SiT_models.keys()), default="SiT-XL/2")
+-    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
+-    parser.add_argument("--num-classes", type=int, default=1000)
+-    parser.add_argument("--epochs", type=int, default=1400)
+-    parser.add_argument("--global-batch-size", type=int, default=256)
+-    parser.add_argument("--global-seed", type=int, default=0)
+-    parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="ema")  # Choice doesn't affect training
+-    parser.add_argument("--num-workers", type=int, default=4)
+-    parser.add_argument("--log-every", type=int, default=100)
+-    parser.add_argument("--ckpt-every", type=int, default=50_000)
+-    parser.add_argument("--sample-every", type=int, default=10_000)
+-    parser.add_argument("--cfg-scale", type=float, default=4.0)
+-    parser.add_argument("--wandb", action="store_true")
+-    parser.add_argument("--ckpt", type=str, default=None,
+-                        help="Optional path to a custom SiT checkpoint")
+-
+-    parse_transport_args(parser)
+-    args = parser.parse_args()
+-    main(args)
++runpy.run_module("train", run_name="__main__")
+diff --git a/train_rot_layer.py b/train_rot_layer.py
+old mode 100644
+new mode 100755
+index 3b40586..8fdc85a
+--- a/train_rot_layer.py
++++ b/train_rot_layer.py
+@@ -1,341 +1,14 @@
+-# This source code is licensed under the license found in the
+-# LICENSE file in the root directory of this source tree.
++#!/usr/bin/env python3
++"""Run the shared SiT trainer with the rotation-layer model implementation."""
+ 
+-"""
+-A minimal training script for SiT using PyTorch DDP.
+-"""
+-import torch
+-# the first flag below was False when we tested this script but True makes A100 training a lot faster:
+-torch.backends.cuda.matmul.allow_tf32 = True
+-torch.backends.cudnn.allow_tf32 = True
+-import torch.distributed as dist
+-from torch.nn.parallel import DistributedDataParallel as DDP
+-from torch.utils.data import DataLoader
+-from torch.utils.data.distributed import DistributedSampler
+-from torchvision.datasets import ImageFolder
+-from torchvision import transforms
+-import numpy as np
+-from collections import OrderedDict
+-from PIL import Image
+-from copy import deepcopy
+-from glob import glob
+-from time import time
+-import argparse
+-import logging
+ import os
++import runpy
+ 
+-from models_rot_layer import SiT_models
+-from download import find_model
+-from transport import create_transport, Sampler
+-from diffusers.models import AutoencoderKL
+-from train_utils import parse_transport_args
+-import wandb_utils
+ 
++# train.py imports its model registry at module load time. Set both variables
++# first, and require the shared trainer to fail closed if another registry is
++# ever selected accidentally.
++os.environ["SIT_MODEL_MODULE"] = "models_rot_layer"
++os.environ["SIT_EXPECTED_MODEL_MODULE"] = "models_rot_layer"
+ 
+-#################################################################################
+-#                             Training Helper Functions                         #
+-#################################################################################
+-
+-@torch.no_grad()
+-def update_ema(ema_model, model, decay=0.9999):
+-    """
+-    Step the EMA model towards the current model.
+-    """
+-    ema_params = OrderedDict(ema_model.named_parameters())
+-    model_params = OrderedDict(model.named_parameters())
+-
+-    for name, param in model_params.items():
+-        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+-        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+-
+-
+-def requires_grad(model, flag=True):
+-    """
+-    Set requires_grad flag for all parameters in a model.
+-    """
+-    for p in model.parameters():
+-        p.requires_grad = flag
+-
+-
+-def cleanup():
+-    """
+-    End DDP training.
+-    """
+-    dist.destroy_process_group()
+-
+-
+-def create_logger(logging_dir):
+-    """
+-    Create a logger that writes to a log file and stdout.
+-    """
+-    if dist.get_rank() == 0:  # real logger
+-        logging.basicConfig(
+-            level=logging.INFO,
+-            format='[\033[34m%(asctime)s\033[0m] %(message)s',
+-            datefmt='%Y-%m-%d %H:%M:%S',
+-            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+-        )
+-        logger = logging.getLogger(__name__)
+-    else:  # dummy logger (does nothing)
+-        logger = logging.getLogger(__name__)
+-        logger.addHandler(logging.NullHandler())
+-    return logger
+-
+-
+-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])
+-
+-
+-#################################################################################
+-#                                  Training Loop                                #
+-#################################################################################
+-
+-def main(args):
+-    """
+-    Trains a new SiT model.
+-    """
+-    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
+-
+-    # Setup DDP:
+-    dist.init_process_group("nccl")
+-    assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
+-    rank = dist.get_rank()
+-    device = rank % torch.cuda.device_count()
+-    seed = args.global_seed * dist.get_world_size() + rank
+-    torch.manual_seed(seed)
+-    torch.cuda.set_device(device)
+-    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
+-    local_batch_size = int(args.global_batch_size // dist.get_world_size())
+-
+-    # Setup an experiment folder:
+-    if rank == 0:
+-        os.makedirs(args.results_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+-        experiment_index = len(glob(f"{args.results_dir}/*"))
+-        model_string_name = args.model.replace("/", "-")  # e.g., SiT-XL/2 --> SiT-XL-2 (for naming folders)
+-        experiment_name = f"{experiment_index:03d}-{model_string_name}-rot-layer-" \
+-                        f"{args.path_type}-{args.prediction}-{args.loss_weight}"
+-        experiment_dir = f"{args.results_dir}/{experiment_name}"  # Create an experiment folder
+-        checkpoint_dir = f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
+-        os.makedirs(checkpoint_dir, exist_ok=True)
+-        logger = create_logger(experiment_dir)
+-        logger.info(f"Experiment directory created at {experiment_dir}")
+-
+-        entity = os.environ["ENTITY"]
+-        project = os.environ["PROJECT"]
+-        if args.wandb:
+-            wandb_utils.initialize(args, entity, experiment_name, project)
+-    else:
+-        logger = create_logger(None)
+-
+-    # Create model:
+-    assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
+-    latent_size = args.image_size // 8
+-    model = SiT_models[args.model](
+-        input_size=latent_size,
+-        num_classes=args.num_classes
+-    )
+-
+-    # Note that parameter initialization is done within the SiT constructor
+-    ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+-
+-    if args.ckpt is not None:
+-        ckpt_path = args.ckpt
+-        state_dict = find_model(ckpt_path)
+-        model.load_state_dict(state_dict["model"])
+-        ema.load_state_dict(state_dict["ema"])
+-        opt.load_state_dict(state_dict["opt"])
+-        args = state_dict["args"]
+-
+-    requires_grad(ema, False)
+-    
+-    model = DDP(model.to(device), device_ids=[device])
+-    transport = create_transport(
+-        args.path_type,
+-        args.prediction,
+-        args.loss_weight,
+-        args.train_eps,
+-        args.sample_eps
+-    )  # default: velocity; 
+-    transport_sampler = Sampler(transport)
+-    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+-    logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+-
+-    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+-    opt = torch.optim.AdamW(model.parameters(), lr=2e-4, weight_decay=0)
+-
+-    # Setup data:
+-    transform = transforms.Compose([
+-        transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
+-        transforms.RandomHorizontalFlip(),
+-        transforms.ToTensor(),
+-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+-    ])
+-    dataset = ImageFolder(args.data_path, transform=transform)
+-    sampler = DistributedSampler(
+-        dataset,
+-        num_replicas=dist.get_world_size(),
+-        rank=rank,
+-        shuffle=True,
+-        seed=args.global_seed
+-    )
+-    loader = DataLoader(
+-        dataset,
+-        batch_size=local_batch_size,
+-        shuffle=False,
+-        sampler=sampler,
+-        num_workers=args.num_workers,
+-        pin_memory=True,
+-        drop_last=True
+-    )
+-    logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
+-
+-    # Prepare models for training:
+-    update_ema(ema, model.module, decay=0)  # Ensure EMA is initialized with synced weights
+-    model.train()  # important! This enables embedding dropout for classifier-free guidance
+-    ema.eval()  # EMA model should always be in eval mode
+-
+-    # Variables for monitoring/logging purposes:
+-    train_steps = 0
+-    log_steps = 0
+-    running_loss = 0
+-    start_time = time()
+-
+-    # Labels to condition the model with (feel free to change):
+-    ys = torch.randint(1000, size=(local_batch_size,), device=device)
+-    use_cfg = args.cfg_scale > 1.0
+-    # Create sampling noise:
+-    n = ys.size(0)
+-    zs = torch.randn(n, 4, latent_size, latent_size, device=device)
+-
+-    # Setup classifier-free guidance:
+-    if use_cfg:
+-        zs = torch.cat([zs, zs], 0)
+-        y_null = torch.tensor([1000] * n, device=device)
+-        ys = torch.cat([ys, y_null], 0)
+-        sample_model_kwargs = dict(y=ys, cfg_scale=args.cfg_scale)
+-        model_fn = ema.forward_with_cfg
+-    else:
+-        sample_model_kwargs = dict(y=ys)
+-        model_fn = ema.forward
+-
+-    logger.info(f"Training for {args.epochs} epochs...")
+-    for epoch in range(args.epochs):
+-        sampler.set_epoch(epoch)
+-        logger.info(f"Beginning epoch {epoch}...")
+-        for x, y in loader:
+-            x = x.to(device)
+-            y = y.to(device)
+-            with torch.no_grad():
+-                # Map input images to latent space + normalize latents:
+-                x = vae.encode(x).latent_dist.sample().mul_(0.18215)
+-            model_kwargs = dict(y=y)
+-            loss_dict = transport.training_losses(model, x, model_kwargs)
+-            loss = loss_dict["loss"].mean()
+-            opt.zero_grad()
+-            loss.backward()
+-            opt.step()
+-            update_ema(ema, model.module)
+-
+-            # Log loss values:
+-            running_loss += loss.item()
+-            log_steps += 1
+-            train_steps += 1
+-            if train_steps % args.log_every == 0:
+-                # Measure training speed:
+-                torch.cuda.synchronize()
+-                end_time = time()
+-                steps_per_sec = log_steps / (end_time - start_time)
+-                # Reduce loss history over all processes:
+-                avg_loss = torch.tensor(running_loss / log_steps, device=device)
+-                dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+-                avg_loss = avg_loss.item() / dist.get_world_size()
+-                logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
+-                if args.wandb:
+-                    wandb_utils.log(
+-                        { "train loss": avg_loss, "train steps/sec": steps_per_sec },
+-                        step=train_steps
+-                    )
+-                # Reset monitoring variables:
+-                running_loss = 0
+-                log_steps = 0
+-                start_time = time()
+-
+-            # Save SiT checkpoint:
+-            if train_steps % args.ckpt_every == 0 and train_steps > 0:
+-                if rank == 0:
+-                    checkpoint = {
+-                        "model": model.module.state_dict(),
+-                        "ema": ema.state_dict(),
+-                        "opt": opt.state_dict(),
+-                        "args": args
+-                    }
+-                    checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+-                    torch.save(checkpoint, checkpoint_path)
+-                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+-                dist.barrier()
+-            
+-            if train_steps % args.sample_every == 0 and train_steps > 0:
+-                logger.info("Generating EMA samples...")
+-                with torch.no_grad():
+-                    sample_fn = transport_sampler.sample_ode() # default to ode sampling
+-                    samples = sample_fn(zs, model_fn, **sample_model_kwargs)[-1]
+-                    dist.barrier()
+-
+-                    if use_cfg: #remove null samples
+-                        samples, _ = samples.chunk(2, dim=0)
+-                    samples = vae.decode(samples / 0.18215).sample
+-                    out_samples = torch.zeros((args.global_batch_size, 3, args.image_size, args.image_size), device=device)
+-                    dist.all_gather_into_tensor(out_samples, samples)
+-
+-                if args.wandb:
+-                    wandb_utils.log_image(out_samples, train_steps)
+-                logging.info("Generating EMA samples done.")
+-
+-    model.eval()  # important! This disables randomized embedding dropout
+-    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+-
+-    logger.info("Done!")
+-    cleanup()
+-
+-
+-if __name__ == "__main__":
+-    # Default args here will train SiT-XL/2 with the hyperparameters we used in our paper (except training iters).
+-    parser = argparse.ArgumentParser()
+-    parser.add_argument("--data-path", type=str, required=True)
+-    parser.add_argument("--results-dir", type=str, default="results")
+-    parser.add_argument("--model", type=str, choices=list(SiT_models.keys()), default="SiT-XL/2")
+-    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
+-    parser.add_argument("--num-classes", type=int, default=1000)
+-    parser.add_argument("--epochs", type=int, default=1400)
+-    parser.add_argument("--global-batch-size", type=int, default=256)
+-    parser.add_argument("--global-seed", type=int, default=0)
+-    parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="ema")  # Choice doesn't affect training
+-    parser.add_argument("--num-workers", type=int, default=4)
+-    parser.add_argument("--log-every", type=int, default=100)
+-    parser.add_argument("--ckpt-every", type=int, default=50_000)
+-    parser.add_argument("--sample-every", type=int, default=10_000)
+-    parser.add_argument("--cfg-scale", type=float, default=4.0)
+-    parser.add_argument("--wandb", action="store_true")
+-    parser.add_argument("--ckpt", type=str, default=None,
+-                        help="Optional path to a custom SiT checkpoint")
+-
+-    parse_transport_args(parser)
+-    args = parser.parse_args()
+-    main(args)
++runpy.run_module("train", run_name="__main__")

code_snapshot/train.py

@@ -0,0 +1,647 @@
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+A minimal training script for SiT using PyTorch DDP.
+"""
+import torch
+# the first flag below was False when we tested this script but True makes A100 training a lot faster:
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from torchvision.datasets import ImageFolder
+from torchvision import transforms
+import numpy as np
+from collections import OrderedDict
+from PIL import Image
+from copy import deepcopy
+from glob import glob
+from time import time
+import argparse
+import csv
+import importlib
+import logging
+import math
+import os
+import re
+import shutil
+from itertools import islice
+
+MODEL_MODULE_NAME = os.environ.get("SIT_MODEL_MODULE", "models")
+model_module = importlib.import_module(MODEL_MODULE_NAME)
+SiT_models = model_module.SiT_models
+MODEL_IMPLEMENTATION_PATH = os.path.realpath(model_module.__file__)
+from download import find_model
+from transport import create_transport, Sampler
+from diffusers.models import AutoencoderKL
+from train_utils import parse_transport_args
+import wandb_utils
+
+
+#################################################################################
+#                             Training Helper Functions                         #
+#################################################################################
+
+@torch.no_grad()
+def update_ema(ema_model, model, decay=0.9999):
+    """
+    Step the EMA model towards the current model.
+    """
+    ema_params = OrderedDict(ema_model.named_parameters())
+    model_params = OrderedDict(model.named_parameters())
+
+    for name, param in model_params.items():
+        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
+        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
+
+
+def requires_grad(model, flag=True):
+    """
+    Set requires_grad flag for all parameters in a model.
+    """
+    for p in model.parameters():
+        p.requires_grad = flag
+
+
+def cleanup():
+    """
+    End DDP training.
+    """
+    dist.destroy_process_group()
+
+
+def create_logger(logging_dir):
+    """
+    Create a logger that writes to a log file and stdout.
+    """
+    if dist.get_rank() == 0:  # real logger
+        logging.basicConfig(
+            level=logging.INFO,
+            format='[\033[34m%(asctime)s\033[0m] %(message)s',
+            datefmt='%Y-%m-%d %H:%M:%S',
+            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
+        )
+        logger = logging.getLogger(__name__)
+    else:  # dummy logger (does nothing)
+        logger = logging.getLogger(__name__)
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+class SkipBatchSampler:
+    """Skip already-consumed batches without loading or transforming their images."""
+
+    def __init__(self, batch_sampler, skip):
+        self.batch_sampler = batch_sampler
+        self.skip = skip
+
+    def __iter__(self):
+        return islice(iter(self.batch_sampler), self.skip, None)
+
+    def __len__(self):
+        return max(0, len(self.batch_sampler) - self.skip)
+
+
+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])
+
+
+@torch.no_grad()
+def evaluate_checkpoint_fid(ema, vae, transport_sampler, args, train_steps,
+                            device, rank, logger, experiment_dir):
+    """Evaluate EMA with CFG=1 while preserving the training RNG trajectory."""
+    world_size = dist.get_world_size()
+    local_batch = args.fid_per_proc_batch_size
+    global_batch = local_batch * world_size
+    total_samples = math.ceil(args.fid_num_samples / global_batch) * global_batch
+    history_path = args.fid_history or os.path.join(experiment_dir, "fid_cfg1_50k.tsv")
+    sample_dir = os.path.join(
+        experiment_dir, "fid_cfg1_work", f"{train_steps:07d}"
+    )
+
+    # A completed record is reusable after a restart. Rank 0 decides and tells
+    # every worker, so all ranks take the same collective path.
+    already_done = False
+    if rank == 0 and os.path.isfile(history_path):
+        with open(history_path, newline="") as f:
+            for row in csv.DictReader(f, delimiter="\t"):
+                if int(row["step"]) == train_steps and row["status"] == "ok":
+                    already_done = True
+                    break
+    done_tensor = torch.tensor(int(already_done), device=device)
+    dist.broadcast(done_tensor, src=0)
+    if done_tensor.item():
+        logger.info(f"Reusing recorded CFG=1 FID for checkpoint {train_steps:07d}")
+        return False
+
+    cpu_rng_state = torch.get_rng_state()
+    cuda_rng_state = torch.cuda.get_rng_state(device)
+    torch.manual_seed(args.fid_seed * world_size + rank)
+    torch.cuda.manual_seed(args.fid_seed * world_size + rank)
+
+    if rank == 0:
+        os.makedirs(sample_dir, exist_ok=True)
+        # A prior interrupted attempt may contain a partial sample set.
+        for name in os.listdir(sample_dir):
+            if name.endswith(".png"):
+                os.remove(os.path.join(sample_dir, name))
+        logger.info(
+            f"Evaluating checkpoint {train_steps:07d}: CFG=1, "
+            f"requested={args.fid_num_samples:,}, actual={total_samples:,}"
+        )
+    dist.barrier()
+
+    sample_fn = transport_sampler.sample_ode(num_steps=args.fid_sampling_steps)
+    latent_size = args.image_size // 8
+    iterations = total_samples // global_batch
+    for batch_index in range(iterations):
+        z = torch.randn(local_batch, 4, latent_size, latent_size, device=device)
+        y = torch.randint(0, args.num_classes, (local_batch,), device=device)
+        samples = sample_fn(z, ema.forward, y=y)[-1]
+        samples = vae.decode(samples / 0.18215).sample
+        samples = torch.clamp(127.5 * samples + 128.0, 0, 255)
+        samples = samples.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
+        for local_index, sample in enumerate(samples):
+            image_index = batch_index * global_batch + local_index * world_size + rank
+            Image.fromarray(sample).save(os.path.join(sample_dir, f"{image_index:06d}.png"))
+        if batch_index % 10 == 0:
+            dist.barrier()
+    dist.barrier()
+
+    fid_value = 0.0
+    stop_requested = False
+    previous_step = None
+    previous_fid = None
+    if rank == 0:
+        from pytorch_fid.fid_score import calculate_fid_given_paths
+
+        fid_value = float(calculate_fid_given_paths(
+            [args.fid_reference, sample_dir],
+            batch_size=args.fid_inception_batch_size,
+            device="cuda:0",
+            dims=2048,
+            num_workers=args.fid_num_workers,
+        ))
+
+        prior_rows = []
+        if os.path.isfile(history_path):
+            with open(history_path, newline="") as f:
+                prior_rows = [
+                    row for row in csv.DictReader(f, delimiter="\t")
+                    if row["status"] == "ok" and int(row["step"]) < train_steps
+                ]
+        trend_rows = sorted(
+            (
+                (int(row["step"]), float(row["fid"]))
+                for row in prior_rows
+            ),
+            key=lambda item: item[0],
+        )
+        trend_rows.append((train_steps, fid_value))
+        required_points = args.fid_stop_consecutive_increases + 1
+        recent_trend = trend_rows[-required_points:]
+        if prior_rows:
+            previous_step, previous_fid = trend_rows[-2]
+        if len(recent_trend) == required_points:
+            consecutive_increases = all(
+                right_fid > left_fid
+                for (_, left_fid), (_, right_fid)
+                in zip(recent_trend, recent_trend[1:])
+            )
+            cumulative_rise = recent_trend[-1][1] - recent_trend[0][1]
+            required_rise = max(
+                args.fid_stop_min_absolute_rise,
+                recent_trend[0][1] * args.fid_stop_min_relative_rise,
+            )
+            stop_requested = consecutive_increases and cumulative_rise >= required_rise
+
+        os.makedirs(os.path.dirname(history_path), exist_ok=True)
+        needs_header = not os.path.isfile(history_path) or os.path.getsize(history_path) == 0
+        with open(history_path, "a", newline="") as f:
+            writer = csv.writer(f, delimiter="\t", lineterminator="\n")
+            if needs_header:
+                writer.writerow([
+                    "step", "checkpoint", "status", "fid", "cfg",
+                    "num_requested", "num_png", "seed", "timestamp_utc"
+                ])
+            from datetime import datetime, timezone
+            writer.writerow([
+                train_steps,
+                os.path.join(experiment_dir, "checkpoints", f"{train_steps:07d}.pt"),
+                "ok", repr(fid_value), "1.0", args.fid_num_samples,
+                total_samples, args.fid_seed,
+                datetime.now(timezone.utc).isoformat(),
+            ])
+
+        logger.info(f"Checkpoint {train_steps:07d} CFG=1 PyTorch FID: {fid_value:.9f}")
+        comparison_output_dir = os.environ.get("SIT_FID_COMPARISON_OUTPUT_DIR")
+        if comparison_output_dir:
+            try:
+                from tools.plot_fid_training_curves import generate_plot
+                generated = generate_plot(
+                    comparison_output_dir,
+                    conv_history=history_path,
+                )
+                logger.info(
+                    f"Updated FID comparison plot: {generated['png']}"
+                )
+            except Exception:
+                # A reporting artifact must never interrupt model training.
+                logger.exception("Could not update the FID comparison plot")
+        if args.wandb:
+            wandb_utils.log({"eval/fid_cfg1_50k": fid_value}, step=train_steps)
+        if stop_requested:
+            marker = os.path.join(experiment_dir, "FID_REGRESSION_STOPPED")
+            with open(marker, "w") as f:
+                f.write(
+                    f"sustained FID regression over {args.fid_stop_consecutive_increases} "
+                    f"consecutive checkpoints: step {recent_trend[0][0]} "
+                    f"FID {recent_trend[0][1]:.9f} -> step {train_steps} "
+                    f"FID {fid_value:.9f}\n"
+                )
+            logger.error(
+                f"FID increased for {args.fid_stop_consecutive_increases} consecutive "
+                f"checkpoints, from {recent_trend[0][1]:.9f} at step "
+                f"{recent_trend[0][0]} to {fid_value:.9f}; stopping after "
+                f"checkpoint {train_steps:07d}."
+            )
+        shutil.rmtree(sample_dir)
+
+    result = torch.tensor([fid_value, float(stop_requested)], device=device)
+    dist.broadcast(result, src=0)
+    dist.barrier()
+
+    # Evaluation must not perturb the random stream used by resumed training.
+    torch.set_rng_state(cpu_rng_state)
+    torch.cuda.set_rng_state(cuda_rng_state, device)
+    return bool(result[1].item())
+
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):
+    """
+    Trains a new SiT model.
+    """
+    assert torch.cuda.is_available(), "Training currently requires at least one GPU."
+    expected_model_module = os.environ.get("SIT_EXPECTED_MODEL_MODULE")
+    if expected_model_module and MODEL_MODULE_NAME != expected_model_module:
+        raise RuntimeError(
+            f"Expected model module {expected_model_module!r}, but loaded "
+            f"{MODEL_MODULE_NAME!r} from {MODEL_IMPLEMENTATION_PATH}"
+        )
+
+    # Load resume metadata before creating the output directory or WandB run. The
+    # checkpoint hyperparameters remain authoritative; only runtime location,
+    # target epoch, and logging options may be overridden by the command line.
+    resume_checkpoint = None
+    resume_step = 0
+    if args.ckpt is not None:
+        runtime_args = args
+        resume_checkpoint = torch.load(args.ckpt, map_location="cpu", weights_only=False)
+        checkpoint_args = resume_checkpoint["args"]
+        runtime_names = (
+            "data_path", "results_dir", "epochs", "wandb", "ckpt", "run_name",
+            "fid_every_checkpoint", "fid_every", "fid_num_samples", "fid_reference",
+            "fid_history", "fid_per_proc_batch_size", "fid_inception_batch_size",
+            "fid_num_workers", "fid_sampling_steps", "fid_seed",
+            "fid_stop_consecutive_increases", "fid_stop_min_absolute_rise",
+            "fid_stop_min_relative_rise",
+        )
+        for name in runtime_names:
+            setattr(checkpoint_args, name, getattr(runtime_args, name))
+        args = checkpoint_args
+        match = re.fullmatch(r"(\d+)\.pt", os.path.basename(args.ckpt))
+        if match is None:
+            raise ValueError("Cannot infer the training step from checkpoint filename; expected NNNNNNN.pt")
+        resume_step = int(match.group(1))
+
+    # Setup DDP:
+    dist.init_process_group("nccl")
+    assert args.global_batch_size % dist.get_world_size() == 0, f"Batch size must be divisible by world size."
+    rank = dist.get_rank()
+    device = rank % torch.cuda.device_count()
+    seed = args.global_seed * dist.get_world_size() + rank
+    torch.manual_seed(seed)
+    torch.cuda.set_device(device)
+    print(f"Starting rank={rank}, seed={seed}, world_size={dist.get_world_size()}.")
+    local_batch_size = int(args.global_batch_size // dist.get_world_size())
+
+    # Setup an experiment folder:
+    if rank == 0:
+        os.makedirs(args.results_dir, exist_ok=True)  # Make results folder (holds all experiment subfolders)
+        experiment_index = len(glob(f"{args.results_dir}/*"))
+        model_string_name = args.model.replace("/", "-")  # e.g., SiT-XL/2 --> SiT-XL-2 (for naming folders)
+        experiment_name = args.run_name or (f"{experiment_index:03d}-{model_string_name}-" \
+                        f"{args.path_type}-{args.prediction}-{args.loss_weight}")
+        experiment_dir = f"{args.results_dir}/{experiment_name}"  # Create an experiment folder
+        checkpoint_dir = f"{experiment_dir}/checkpoints"  # Stores saved model checkpoints
+        os.makedirs(checkpoint_dir, exist_ok=True)
+        logger = create_logger(experiment_dir)
+        logger.info(f"Experiment directory created at {experiment_dir}")
+        logger.info(
+            f"Model implementation: {MODEL_MODULE_NAME} ({MODEL_IMPLEMENTATION_PATH})"
+        )
+
+        entity = os.environ["ENTITY"]
+        project = os.environ["PROJECT"]
+        if args.wandb:
+            wandb_utils.initialize(args, entity, experiment_name, project)
+    else:
+        logger = create_logger(None)
+        experiment_dir = None
+    path_objects = [experiment_dir]
+    dist.broadcast_object_list(path_objects, src=0)
+    experiment_dir = path_objects[0]
+
+    # Create model:
+    assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
+    latent_size = args.image_size // 8
+    model = SiT_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes
+    )
+
+    # Note that parameter initialization is done within the SiT constructor
+    ema = deepcopy(model).to(device)  # Create an EMA of the model for use after training
+
+    requires_grad(ema, False)
+    
+    model = DDP(model.to(device), device_ids=[device])
+    transport = create_transport(
+        args.path_type,
+        args.prediction,
+        args.loss_weight,
+        args.train_eps,
+        args.sample_eps
+    )  # default: velocity; 
+    transport_sampler = Sampler(transport)
+    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}").to(device)
+    logger.info(f"SiT Parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Setup optimizer (we used default Adam betas=(0.9, 0.999) and a constant learning rate of 1e-4 in our paper):
+    opt = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=0)
+    logger.info(
+        f"Optimizer: AdamW(lr={args.learning_rate:g}, weight_decay=0, "
+        "betas=(0.9, 0.999))"
+    )
+    if resume_checkpoint is not None:
+        model.module.load_state_dict(resume_checkpoint["model"])
+        ema.load_state_dict(resume_checkpoint["ema"])
+        opt.load_state_dict(resume_checkpoint["opt"])
+        logger.info(f"Resumed model, EMA, and optimizer from {args.ckpt} at step {resume_step:,}")
+
+    # Setup data:
+    transform = transforms.Compose([
+        transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, args.image_size)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+    ])
+    dataset = ImageFolder(args.data_path, transform=transform)
+    sampler = DistributedSampler(
+        dataset,
+        num_replicas=dist.get_world_size(),
+        rank=rank,
+        shuffle=True,
+        seed=args.global_seed
+    )
+    loader = DataLoader(
+        dataset,
+        batch_size=local_batch_size,
+        shuffle=False,
+        sampler=sampler,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True
+    )
+    logger.info(f"Dataset contains {len(dataset):,} images ({args.data_path})")
+
+    # Prepare models for training:
+    if resume_checkpoint is None:
+        update_ema(ema, model.module, decay=0)  # Initialize EMA only for a new run.
+    model.train()  # important! This enables embedding dropout for classifier-free guidance
+    ema.eval()  # EMA model should always be in eval mode
+
+    # Variables for monitoring/logging purposes:
+    train_steps = resume_step
+    log_steps = 0
+    running_loss = 0
+    start_time = time()
+
+    # Labels to condition the model with (feel free to change):
+    ys = torch.randint(1000, size=(local_batch_size,), device=device)
+    use_cfg = args.cfg_scale > 1.0
+    # Create sampling noise:
+    n = ys.size(0)
+    zs = torch.randn(n, 4, latent_size, latent_size, device=device)
+
+    # Setup classifier-free guidance:
+    if use_cfg:
+        zs = torch.cat([zs, zs], 0)
+        y_null = torch.tensor([1000] * n, device=device)
+        ys = torch.cat([ys, y_null], 0)
+        sample_model_kwargs = dict(y=ys, cfg_scale=args.cfg_scale)
+        model_fn = ema.forward_with_cfg
+    else:
+        sample_model_kwargs = dict(y=ys)
+        model_fn = ema.forward
+
+    steps_per_epoch = len(loader)
+    target_steps = args.epochs * steps_per_epoch
+    start_epoch, batches_to_skip = divmod(train_steps, steps_per_epoch)
+    logger.info(
+        f"Training to {args.epochs} total epochs ({target_steps:,} steps); "
+        f"starting at epoch {start_epoch}, batch {batches_to_skip}, step {train_steps:,}."
+    )
+    # Establish a same-protocol 50k baseline for the resume checkpoint before
+    # comparing later checkpoints against it. A recorded baseline is reused.
+    if args.fid_every_checkpoint and resume_checkpoint is not None:
+        if evaluate_checkpoint_fid(
+            ema, vae, transport_sampler, args, train_steps,
+            device, rank, logger, experiment_dir,
+        ):
+            logger.error("Resume checkpoint already violates the recorded FID trend; exiting.")
+            cleanup()
+            return
+        start_time = time()
+
+    stop_requested = False
+    for epoch in range(start_epoch, args.epochs):
+        sampler.set_epoch(epoch)
+        logger.info(f"Beginning epoch {epoch}...")
+        epoch_loader = loader
+        if epoch == start_epoch and batches_to_skip:
+            epoch_loader = DataLoader(
+                dataset,
+                batch_sampler=SkipBatchSampler(loader.batch_sampler, batches_to_skip),
+                num_workers=args.num_workers,
+                pin_memory=True,
+            )
+        for x, y in epoch_loader:
+            x = x.to(device)
+            y = y.to(device)
+            with torch.no_grad():
+                # Map input images to latent space + normalize latents:
+                x = vae.encode(x).latent_dist.sample().mul_(0.18215)
+            model_kwargs = dict(y=y)
+            loss_dict = transport.training_losses(model, x, model_kwargs)
+            loss = loss_dict["loss"].mean()
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            update_ema(ema, model.module)
+
+            # Log loss values:
+            running_loss += loss.item()
+            log_steps += 1
+            train_steps += 1
+            if train_steps % args.log_every == 0:
+                # Measure training speed:
+                torch.cuda.synchronize()
+                end_time = time()
+                steps_per_sec = log_steps / (end_time - start_time)
+                # Reduce loss history over all processes:
+                avg_loss = torch.tensor(running_loss / log_steps, device=device)
+                dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+                avg_loss = avg_loss.item() / dist.get_world_size()
+                logger.info(f"(step={train_steps:07d}) Train Loss: {avg_loss:.4f}, Train Steps/Sec: {steps_per_sec:.2f}")
+                if args.wandb:
+                    wandb_utils.log(
+                        { "train loss": avg_loss, "train steps/sec": steps_per_sec },
+                        step=train_steps
+                    )
+                # Reset monitoring variables:
+                running_loss = 0
+                log_steps = 0
+                start_time = time()
+
+            # Save SiT checkpoint:
+            if train_steps % args.ckpt_every == 0 and train_steps > 0:
+                if rank == 0:
+                    checkpoint = {
+                        "model": model.module.state_dict(),
+                        "ema": ema.state_dict(),
+                        "opt": opt.state_dict(),
+                        "args": args
+                    }
+                    checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+                    torch.save(checkpoint, checkpoint_path)
+                    logger.info(f"Saved checkpoint to {checkpoint_path}")
+                dist.barrier()
+                if args.fid_every_checkpoint and train_steps % args.fid_every == 0:
+                    stop_requested = evaluate_checkpoint_fid(
+                        ema, vae, transport_sampler, args, train_steps,
+                        device, rank, logger, experiment_dir,
+                    )
+                    start_time = time()
+                    if stop_requested:
+                        break
+            
+            if train_steps % args.sample_every == 0 and train_steps > 0:
+                logger.info("Generating EMA samples...")
+                with torch.no_grad():
+                    sample_fn = transport_sampler.sample_ode() # default to ode sampling
+                    samples = sample_fn(zs, model_fn, **sample_model_kwargs)[-1]
+                    dist.barrier()
+
+                    if use_cfg: #remove null samples
+                        samples, _ = samples.chunk(2, dim=0)
+                    samples = vae.decode(samples / 0.18215).sample
+                    out_samples = torch.zeros((args.global_batch_size, 3, args.image_size, args.image_size), device=device)
+                    dist.all_gather_into_tensor(out_samples, samples)
+
+                if args.wandb:
+                    wandb_utils.log_image(out_samples, train_steps)
+                logging.info("Generating EMA samples done.")
+
+        if stop_requested:
+            break
+
+    if rank == 0 and not stop_requested and train_steps % args.ckpt_every != 0:
+        checkpoint = {
+            "model": model.module.state_dict(),
+            "ema": ema.state_dict(),
+            "opt": opt.state_dict(),
+            "args": args,
+        }
+        checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
+        torch.save(checkpoint, checkpoint_path)
+        logger.info(f"Saved final checkpoint to {checkpoint_path}")
+    dist.barrier()
+
+    model.eval()  # important! This disables randomized embedding dropout
+    # do any sampling/FID calculation/etc. with ema (or model) in eval mode ...
+
+    logger.info("Done!")
+    cleanup()
+
+
+if __name__ == "__main__":
+    # Default args here will train SiT-XL/2 with the hyperparameters we used in our paper (except training iters).
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str, required=True)
+    parser.add_argument("--results-dir", type=str, default="results")
+    parser.add_argument("--model", type=str, choices=list(SiT_models.keys()), default="SiT-XL/2")
+    parser.add_argument("--image-size", type=int, choices=[256, 512], default=256)
+    parser.add_argument("--num-classes", type=int, default=1000)
+    parser.add_argument("--epochs", type=int, default=1400)
+    parser.add_argument("--global-batch-size", type=int, default=256)
+    parser.add_argument("--learning-rate", type=float, default=1e-4)
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="ema")  # Choice doesn't affect training
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--log-every", type=int, default=100)
+    parser.add_argument("--ckpt-every", type=int, default=50_000)
+    parser.add_argument("--sample-every", type=int, default=10_000)
+    parser.add_argument("--cfg-scale", type=float, default=4.0)
+    parser.add_argument("--wandb", action="store_true")
+    parser.add_argument("--ckpt", type=str, default=None,
+                        help="Optional path to a custom SiT checkpoint")
+    parser.add_argument("--run-name", type=str, default=None,
+                        help="Experiment directory and WandB run name (useful when resuming)")
+    parser.add_argument("--fid-every-checkpoint", action="store_true",
+                        help="Run periodic CFG=1 PyTorch FID checks and stop on sustained regression")
+    parser.add_argument("--fid-every", type=int, default=50_000,
+                        help="Training-step interval between FID checks")
+    parser.add_argument("--fid-num-samples", type=int, default=50_000)
+    parser.add_argument("--fid-reference", type=str,
+                        default="/home/nvidia/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz")
+    parser.add_argument("--fid-history", type=str, default=None)
+    parser.add_argument("--fid-per-proc-batch-size", type=int, default=64)
+    parser.add_argument("--fid-inception-batch-size", type=int, default=128)
+    parser.add_argument("--fid-num-workers", type=int, default=8)
+    parser.add_argument("--fid-sampling-steps", type=int, default=250)
+    parser.add_argument("--fid-seed", type=int, default=0)
+    parser.add_argument("--fid-stop-consecutive-increases", type=int, default=3,
+                        help="Stop only after this many consecutive checkpoint FID increases")
+    parser.add_argument("--fid-stop-min-absolute-rise", type=float, default=0.25,
+                        help="Minimum cumulative absolute FID rise required to stop")
+    parser.add_argument("--fid-stop-min-relative-rise", type=float, default=0.005,
+                        help="Minimum cumulative relative FID rise required to stop")
+
+    parse_transport_args(parser)
+    args = parser.parse_args()
+    main(args)

code_snapshot/train_conv.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+"""Run the shared SiT trainer with the convolutional-layer model implementation."""
+
+import os
+import runpy
+
+
+# Select models_conv.py before train.py imports the model registry, and require
+# the shared trainer to fail closed if another model module is selected.
+os.environ["SIT_MODEL_MODULE"] = "models_conv"
+os.environ["SIT_EXPECTED_MODEL_MODULE"] = "models_conv"
+
+runpy.run_module("train", run_name="__main__")

code_snapshot/train_rot_layer.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+"""Run the shared SiT trainer with the rotation-layer model implementation."""
+
+import os
+import runpy
+
+
+# train.py imports its model registry at module load time. Set both variables
+# first, and require the shared trainer to fail closed if another registry is
+# ever selected accidentally.
+os.environ["SIT_MODEL_MODULE"] = "models_rot_layer"
+os.environ["SIT_EXPECTED_MODEL_MODULE"] = "models_rot_layer"
+
+runpy.run_module("train", run_name="__main__")

experiments/bs256_lr1e-4/EXPERIMENT_RESULTS.md

@@ -0,0 +1,23 @@
+# SiT-S/2 ImageNet-256 experiment snapshot
+
+All three runs use the same training protocol unless stated otherwise: ImageNet-1K
+256px ImageFolder input, global batch size 256 on 8 GPUs, AdamW with learning
+rate `1e-4`, Linear/velocity transport, SD VAE EMA, and 5,004 training steps per
+epoch. FID values below are PyTorch-FID results using the `SII-PengZheng/discon`
+`VIRTUAL_imagenet256_labeled.npz` reference, 50,176 generated PNGs, seed 0, and
+250 Euler sampling steps. Lower is better.
+
+| Variant | Backup checkpoint | Progress | CFG=1 latest FID | CFG=4 final FID |
+| --- | --- | ---: | ---: | ---: |
+| Base | `4003200.pt` | 800 epochs | 38.871 at 4,000,000 steps | 11.042 at 4,003,200 steps |
+| Rotation-layer | `4003200.pt` | 800 epochs | 40.274 at 4,000,000 steps | 11.189 at 4,003,200 steps |
+| Conv-layer | `1950000.pt` | 389.7 epochs | 45.239 at 1,750,000 steps | not run |
+
+The convolution-layer run was intentionally stopped before its next periodic FID
+point. Its latest saved checkpoint is step 1,950,000; its 1,750,000-step FID is
+the most recent completed, same-protocol measurement.
+
+`fid_cfg1_50k_training_curves.{png,svg,csv}` and the three source TSV files are
+the authoritative training-trend records. The accompanying checkpoint manifest
+contains SHA-256 hashes and strict-load/inference validation results for every
+backup checkpoint.

experiments/bs256_lr1e-4/RESUME_GUIDE.md

@@ -0,0 +1,121 @@
+# Recreate and resume the backed-up SiT experiments
+
+This guide restores the base, rotation-layer, or convolution-layer experiment on
+a new GPU server. It assumes that the Hugging Face account has accepted the
+ImageNet-1K access conditions and can read `BlueSourceJY/SiT-Complementary`.
+
+## 1. Create the environment and obtain the code
+
+```bash
+git clone https://github.com/hu-jy-0406/SiT-Complementary.git
+cd SiT-Complementary
+conda env create -f environment.yml
+conda activate SiT
+pip install pyarrow pytorch-fid huggingface_hub
+huggingface-cli login
+```
+
+Install a CUDA/PyTorch build appropriate for the new server if the environment
+file's default build does not match its driver. Eight visible GPUs are assumed by
+the commands below.
+
+## 2. Download ImageNet-1K and convert it to ImageFolder
+
+The gated Hugging Face ImageNet export is parquet. Keep that download separate
+from the ImageFolder output (the latter is what training uses).
+
+```bash
+export DATA_ROOT=/data/imagenet-1k
+huggingface-cli download ILSVRC/imagenet-1k --repo-type dataset \
+  --local-dir "${DATA_ROOT}-hf" --include 'data/*.parquet' 'classes.py' 'README.md'
+python tools/convert_imagenet_parquet_to_imagefolder.py \
+  "${DATA_ROOT}-hf" "$DATA_ROOT" "${DATA_ROOT}-hf/classes.py"
+find "$DATA_ROOT/train" -mindepth 1 -maxdepth 1 -type d | wc -l  # must be 1000
+```
+
+The training `--data-path` is `"$DATA_ROOT/train"`. The converter also creates
+`validation/`, `class_to_idx.json`, and `conversion_summary.json` under
+`$DATA_ROOT`.
+
+## 3. Download the experiment backup and FID reference
+
+```bash
+export EXP_ROOT=/data/sit-backup
+huggingface-cli download BlueSourceJY/SiT-Complementary \
+  --local-dir "$EXP_ROOT" \
+  --include 'checkpoints/bs256_lr1e-4/*' 'experiments/bs256_lr1e-4/*'
+
+mkdir -p /data/evaluation/reference/discon-download
+huggingface-cli download SII-PengZheng/discon \
+  VIRTUAL_imagenet256_labeled.npz \
+  --local-dir /data/evaluation/reference/discon-download
+```
+
+Before resuming, compare each downloaded file's SHA-256 with
+`$EXP_ROOT/experiments/bs256_lr1e-4/checkpoint_validation.json`. The training
+checkpoint format contains `model`, `ema`, `opt`, and `args`; resume restores all
+four, including optimizer state.
+
+## 4. One-image inference check
+
+Use the model implementation that matches the checkpoint. `learn_sigma=True` is
+intentional for these checkpoints; the model returns the velocity half internally.
+
+```bash
+# Base (replace the checkpoint path for the other variants)
+torchrun --nproc_per_node=1 sample_ddp.py ODE \
+  --model SiT-S/2 --image-size 256 \
+  --ckpt "$EXP_ROOT/checkpoints/bs256_lr1e-4/base/4003200.pt" \
+  --num-fid-samples 1 --per-proc-batch-size 1 --num-sampling-steps 10 \
+  --sample-dir /data/sit-smoke-test
+
+# Rotation layer: prefix the same command with:
+# SIT_MODEL_MODULE=models_rot_layer SIT_EXPECTED_MODEL_MODULE=models_rot_layer
+# Conv layer: prefix the same command with:
+# SIT_MODEL_MODULE=models_conv SIT_EXPECTED_MODEL_MODULE=models_conv
+```
+
+Inspect the generated PNG before spending GPU time. For a repeatable FID run,
+use eight processes and the archived settings: `--num-fid-samples 50000`,
+`--per-proc-batch-size 64`, `--num-sampling-steps 250`, and the required CFG.
+
+## 5. Resume training
+
+Set W&B variables if online logging is wanted; for offline logging set
+`WANDB_MODE=offline`, then run `wandb sync <offline-run-dir>` later. Use a new
+`--results-dir` on the new server, retain the variant-specific `--run-name`, and
+pass a target total epoch count greater than the saved progress. Do not change
+the model module, global batch size, learning rate, seed, or data preprocessing.
+
+```bash
+export DATA_ROOT=/data/imagenet-1k
+export EXP_ROOT=/data/sit-backup
+export REF=/data/evaluation/reference/discon-download/VIRTUAL_imagenet256_labeled.npz
+export ENTITY='<wandb entity>' PROJECT='<wandb project>' WANDB_KEY='<wandb key>'
+
+# Base: continue 800-epoch final checkpoint to 1,000 total epochs.
+torchrun --nnodes=1 --nproc_per_node=8 train.py \
+  --model SiT-S/2 --epochs 1000 --data-path "$DATA_ROOT/train" \
+  --results-dir /data/sit-results/base --global-batch-size 256 \
+  --learning-rate 0.0001 --global-seed 0 --vae ema --num-workers 4 \
+  --log-every 100 --ckpt-every 50000 --sample-every 10000 --cfg-scale 4.0 \
+  --run-name SiT-S-2-Linear-velocity-None \
+  --ckpt "$EXP_ROOT/checkpoints/bs256_lr1e-4/base/4003200.pt" \
+  --fid-every-checkpoint --fid-every 250000 --fid-num-samples 50000 \
+  --fid-reference "$REF" --fid-per-proc-batch-size 64 \
+  --fid-inception-batch-size 128 --fid-num-workers 8 --fid-sampling-steps 250 \
+  --fid-seed 0 --fid-stop-consecutive-increases 3 \
+  --fid-stop-min-absolute-rise 0.25 --fid-stop-min-relative-rise 0.005 --wandb
+```
+
+For rotation-layer replace `train.py` with `train_rot_layer.py`, use the
+`rotation-layer/4003200.pt` checkpoint, and choose an isolated results directory.
+For conv-layer use `train_conv.py` and `conv-layer/1950000.pt`. Those wrappers
+select `models_rot_layer.py` and `models_conv.py` before importing the shared
+trainer, so never substitute plain `train.py` for them.
+
+The trainer derives the current step from the numeric checkpoint filename and
+skips consumed batches in the first resumed epoch. It will first establish or
+reuse the matching CFG=1 FID baseline, then records each periodic result in
+`fid_cfg1_50k.tsv`. It stops only after the configured sustained, material FID
+regression rule is met.

experiments/bs256_lr1e-4/base/fid_cfg1_50k.tsv

@@ -0,0 +1,49 @@
+step	checkpoint	status	fid	cfg	num_requested	num_png	seed	timestamp_utc
+1000000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1000000.pt	ok	48.29311023150069	1.0	50000	50176	0	2026-07-22T03:33:11.273379+00:00
+1050000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1050000.pt	ok	48.018564188901564	1.0	50000	50176	0	2026-07-22T06:16:58.701669+00:00
+1100000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1100000.pt	ok	47.47688964066043	1.0	50000	50176	0	2026-07-22T09:01:01.876690+00:00
+1150000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1150000.pt	ok	47.10925915220656	1.0	50000	50176	0	2026-07-22T11:44:42.044226+00:00
+1200000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1200000.pt	ok	46.718045833545546	1.0	50000	50176	0	2026-07-22T14:28:44.897893+00:00
+1250000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1250000.pt	ok	46.414511520279234	1.0	50000	50176	0	2026-07-22T17:32:00.597820+00:00
+1300000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1300000.pt	ok	45.92057227506251	1.0	50000	50176	0	2026-07-22T20:16:10.465171+00:00
+1350000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1350000.pt	ok	45.89153691429402	1.0	50000	50176	0	2026-07-22T23:00:16.427930+00:00
+1400000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1400000.pt	ok	45.54178671851474	1.0	50000	50176	0	2026-07-23T01:44:23.027126+00:00
+1450000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1450000.pt	ok	45.116327425034854	1.0	50000	50176	0	2026-07-23T04:28:29.462587+00:00
+1500000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1500000.pt	ok	44.962781651113005	1.0	50000	50176	0	2026-07-23T07:12:39.286539+00:00
+1550000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1550000.pt	ok	44.842452348401196	1.0	50000	50176	0	2026-07-23T09:56:37.374904+00:00
+1600000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1600000.pt	ok	44.640852044147266	1.0	50000	50176	0	2026-07-23T12:40:46.596453+00:00
+1650000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1650000.pt	ok	44.237907149190846	1.0	50000	50176	0	2026-07-23T15:24:58.652950+00:00
+1700000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1700000.pt	ok	44.136045077932806	1.0	50000	50176	0	2026-07-23T18:26:45.897450+00:00
+1750000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1750000.pt	ok	43.767861385718845	1.0	50000	50176	0	2026-07-23T21:10:59.723969+00:00
+1800000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1800000.pt	ok	43.64788291826136	1.0	50000	50176	0	2026-07-23T23:55:10.773687+00:00
+1850000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1850000.pt	ok	43.409102472303005	1.0	50000	50176	0	2026-07-24T02:39:16.581973+00:00
+1900000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1900000.pt	ok	43.21848172913133	1.0	50000	50176	0	2026-07-24T05:23:19.752006+00:00
+1950000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1950000.pt	ok	42.921743510806834	1.0	50000	50176	0	2026-07-24T08:07:41.369097+00:00
+2000000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2000000.pt	ok	42.80858523280904	1.0	50000	50176	0	2026-07-24T10:51:57.437664+00:00
+2050000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2050000.pt	ok	42.67307960077221	1.0	50000	50176	0	2026-07-24T13:36:12.546595+00:00
+2100000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2100000.pt	ok	42.443703960448886	1.0	50000	50176	0	2026-07-24T16:21:04.406803+00:00
+2150000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2150000.pt	ok	42.50299269928496	1.0	50000	50176	0	2026-07-24T19:22:47.091546+00:00
+2200000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2200000.pt	ok	42.20738877449185	1.0	50000	50176	0	2026-07-27T04:37:51.175075+00:00
+2250000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2250000.pt	ok	42.09083954185155	1.0	50000	50176	0	2026-07-27T07:22:06.775572+00:00
+2300000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2300000.pt	ok	41.95050261156155	1.0	50000	50176	0	2026-07-27T10:06:30.595161+00:00
+2350000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2350000.pt	ok	41.72793962332611	1.0	50000	50176	0	2026-07-27T12:51:06.430298+00:00
+2400000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2400000.pt	ok	41.68562578798634	1.0	50000	50176	0	2026-07-27T15:35:39.897651+00:00
+2450000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2450000.pt	ok	41.43696452027939	1.0	50000	50176	0	2026-07-27T18:20:03.971453+00:00
+2500000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2500000.pt	ok	41.54235319764359	1.0	50000	50176	0	2026-07-27T21:04:36.195527+00:00
+2550000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2550000.pt	ok	41.33031665291526	1.0	50000	50176	0	2026-07-28T18:27:35.792552+00:00
+2600000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2600000.pt	ok	41.09229129923432	1.0	50000	50176	0	2026-07-28T21:12:00.197229+00:00
+2650000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2650000.pt	ok	41.04462000106378	1.0	50000	50176	0	2026-07-28T23:56:30.747960+00:00
+2700000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2700000.pt	ok	40.986577043428895	1.0	50000	50176	0	2026-07-29T02:41:03.183436+00:00
+2750000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2750000.pt	ok	40.88498184813966	1.0	50000	50176	0	2026-07-29T05:25:36.916055+00:00
+2800000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2800000.pt	ok	40.794031115349355	1.0	50000	50176	0	2026-07-29T08:10:12.737171+00:00
+2850000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2850000.pt	ok	40.42894815467065	1.0	50000	50176	0	2026-07-29T10:54:46.079051+00:00
+2900000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2900000.pt	ok	40.422900985210845	1.0	50000	50176	0	2026-07-29T13:39:29.665685+00:00
+2950000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2950000.pt	ok	40.54817889332162	1.0	50000	50176	0	2026-07-29T16:46:35.078607+00:00
+3000000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3000000.pt	ok	40.46896141626422	1.0	50000	50176	0	2026-07-29T19:31:07.135111+00:00
+3050000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3050000.pt	ok	40.25837670401995	1.0	50000	50176	0	2026-07-29T22:15:37.950844+00:00
+3100000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3100000.pt	ok	40.21306397405647	1.0	50000	50176	0	2026-07-30T01:00:08.137327+00:00
+3150000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3150000.pt	ok	40.238701584494436	1.0	50000	50176	0	2026-07-30T03:44:41.736117+00:00
+3250000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3250000.pt	ok	39.798277728671735	1.0	50000	50176	0	2026-07-30T09:08:00.799118+00:00
+3500000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3500000.pt	ok	39.5760757594021	1.0	50000	50176	0	2026-07-30T22:46:26.192799+00:00
+3750000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3750000.pt	ok	39.16467116363083	1.0	50000	50176	0	2026-07-31T12:03:12.844048+00:00
+4000000	/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/4000000.pt	ok	38.87082736834208	1.0	50000	50176	0	2026-08-01T01:39:07.070873+00:00

experiments/bs256_lr1e-4/base/final_4003200_cfg4.0_pytorch_fid.txt

@@ -0,0 +1,2 @@
+
  0%|          | 0/392 [00:00<?, ?it/s]
  0%|          | 1/392 [00:01<07:19,  1.12s/it]
  1%|          | 3/392 [00:01<02:17,  2.83it/s]
  1%|▏         | 5/392 [00:01<01:21,  4.76it/s]
  2%|▏         | 7/392 [00:01<00:58,  6.54it/s]
  2%|▏         | 9/392 [00:01<00:47,  8.06it/s]
  3%|▎         | 11/392 [00:01<00:40,  9.40it/s]
  3%|▎         | 13/392 [00:02<00:36, 10.46it/s]
  4%|▍         | 15/392 [00:02<00:33, 11.28it/s]
  4%|▍         | 17/392 [00:02<00:31, 11.85it/s]
  5%|▍         | 19/392 [00:02<00:30, 12.28it/s]
  5%|▌         | 21/392 [00:02<00:29, 12.58it/s]
  6%|▌         | 23/392 [00:02<00:28, 12.78it/s]
  6%|▋         | 25/392 [00:02<00:28, 12.98it/s]
  7%|▋         | 27/392 [00:03<00:27, 13.09it/s]
  7%|▋         | 29/392 [00:03<00:27, 13.15it/s]
  8%|▊         | 31/392 [00:03<00:27, 13.23it/s]
  8%|▊         | 33/392 [00:03<00:27, 13.26it/s]
  9%|▉         | 35/392 [00:03<00:26, 13.27it/s]
  9%|▉         | 37/392 [00:03<00:26, 13.35it/s]
 10%|▉         | 39/392 [00:03<00:26, 13.38it/s]
 10%|█         | 41/392 [00:04<00:26, 13.48it/s]
 11%|█         | 43/392 [00:04<00:25, 13.51it/s]
 11%|█▏        | 45/392 [00:04<00:25, 13.49it/s]
 12%|█▏        | 47/392 [00:04<00:25, 13.43it/s]
 12%|█▎        | 49/392 [00:04<00:25, 13.42it/s]
 13%|█▎        | 51/392 [00:04<00:25, 13.37it/s]
 14%|█▎        | 53/392 [00:05<00:25, 13.37it/s]
 14%|█▍        | 55/392 [00:05<00:25, 13.40it/s]
 15%|█▍        | 57/392 [00:05<00:25, 13.40it/s]
 15%|█▌        | 59/392 [00:05<00:24, 13.41it/s]
 16%|█▌        | 61/392 [00:05<00:24, 13.44it/s]
 16%|█▌        | 63/392 [00:05<00:24, 13.40it/s]
 17%|█▋        | 65/392 [00:05<00:24, 13.32it/s]
 17%|█▋        | 67/392 [00:06<00:24, 13.34it/s]
 18%|█▊        | 69/392 [00:06<00:24, 13.35it/s]
 18%|█▊        | 71/392 [00:06<00:24, 13.32it/s]
 19%|█▊        | 73/392 [00:06<00:24, 13.27it/s]
 19%|█▉        | 75/392 [00:06<00:23, 13.32it/s]
 20%|█▉        | 77/392 [00:06<00:23, 13.34it/s]
 20%|██        | 79/392 [00:06<00:23, 13.37it/s]
 21%|██        | 81/392 [00:07<00:23, 13.34it/s]
 21%|██        | 83/392 [00:07<00:23, 13.35it/s]
 22%|██▏       | 85/392 [00:07<00:23, 13.32it/s]
 22%|██▏       | 87/392 [00:07<00:22, 13.30it/s]
 23%|██▎       | 89/392 [00:07<00:22, 13.29it/s]
 23%|██▎       | 91/392 [00:07<00:22, 13.30it/s]
 24%|██▎       | 93/392 [00:08<00:22, 13.26it/s]
 24%|██▍       | 95/392 [00:08<00:22, 13.28it/s]
 25%|██▍       | 97/392 [00:08<00:22, 13.32it/s]
 25%|██▌       | 99/392 [00:08<00:21, 13.34it/s]
 26%|██▌       | 101/392 [00:08<00:21, 13.34it/s]
 26%|██▋       | 103/392 [00:08<00:21, 13.33it/s]
 27%|██▋       | 105/392 [00:08<00:21, 13.32it/s]
 27%|██▋       | 107/392 [00:09<00:21, 13.33it/s]
 28%|██▊       | 109/392 [00:09<00:21, 13.32it/s]
 28%|██▊       | 111/392 [00:09<00:21, 13.28it/s]
 29%|██▉       | 113/392 [00:09<00:21, 13.23it/s]
 29%|██▉       | 115/392 [00:09<00:20, 13.29it/s]
 30%|██▉       | 117/392 [00:09<00:20, 13.30it/s]
 30%|███       | 119/392 [00:09<00:20, 13.30it/s]
 31%|███       | 121/392 [00:10<00:20, 13.31it/s]
 31%|███▏      | 123/392 [00:10<00:20, 13.31it/s]
 32%|███▏      | 125/392 [00:10<00:20, 13.32it/s]
 32%|███▏      | 127/392 [00:10<00:19, 13.29it/s]
 33%|███▎      | 129/392 [00:10<00:19, 13.24it/s]
 33%|███▎      | 131/392 [00:10<00:19, 13.25it/s]
 34%|███▍      | 133/392 [00:11<00:19, 13.22it/s]
 34%|███▍      | 135/392 [00:11<00:19, 13.26it/s]
 35%|███▍      | 137/392 [00:11<00:19, 13.27it/s]
 35%|███▌      | 139/392 [00:11<00:18, 13.38it/s]
 36%|███▌      | 141/392 [00:11<00:18, 13.44it/s]
 36%|███▋      | 143/392 [00:11<00:18, 13.47it/s]
 37%|███▋      | 145/392 [00:11<00:18, 13.51it/s]
 38%|███▊      | 147/392 [00:12<00:18, 13.55it/s]
 38%|███▊      | 149/392 [00:12<00:17, 13.56it/s]
 39%|███▊      | 151/392 [00:12<00:18, 13.18it/s]
 39%|███▉      | 153/392 [00:12<00:18, 12.90it/s]
 40%|███▉      | 155/392 [00:12<00:18, 12.67it/s]
 40%|████      | 157/392 [00:12<00:18, 12.42it/s]
 41%|████      | 159/392 [00:13<00:18, 12.35it/s]
 41%|████      | 161/392 [00:13<00:18, 12.45it/s]
 42%|████▏     | 163/392 [00:13<00:18, 12.42it/s]
 42%|████▏     | 165/392 [00:13<00:18, 12.25it/s]
 43%|████▎     | 167/392 [00:13<00:18, 12.26it/s]
 43%|████▎     | 169/392 [00:13<00:18, 12.27it/s]
 44%|████▎     | 171/392 [00:14<00:17, 12.35it/s]
 44%|████▍     | 173/392 [00:14<00:17, 12.34it/s]
 45%|████▍     | 175/392 [00:14<00:17, 12.43it/s]
 45%|████▌     | 177/392 [00:14<00:17, 12.13it/s]
 46%|████▌     | 179/392 [00:14<00:17, 12.07it/s]
 46%|████▌     | 181/392 [00:14<00:17, 12.07it/s]
 47%|████▋     | 183/392 [00:15<00:17, 12.09it/s]
 47%|██��█▋     | 185/392 [00:15<00:17, 12.09it/s]
 48%|████▊     | 187/392 [00:15<00:16, 12.16it/s]
 48%|████▊     | 189/392 [00:15<00:16, 12.16it/s]
 49%|████▊     | 191/392 [00:15<00:16, 11.96it/s]
 49%|████▉     | 193/392 [00:15<00:16, 12.00it/s]
 50%|████▉     | 195/392 [00:16<00:16, 12.00it/s]
 50%|█████     | 197/392 [00:16<00:16, 12.00it/s]
 51%|█████     | 199/392 [00:16<00:16, 11.91it/s]
 51%|█████▏    | 201/392 [00:16<00:15, 12.04it/s]
 52%|█████▏    | 203/392 [00:16<00:15, 12.03it/s]
 52%|█████▏    | 205/392 [00:16<00:15, 11.92it/s]
 53%|█████▎    | 207/392 [00:17<00:15, 11.97it/s]
 53%|█████▎    | 209/392 [00:17<00:15, 11.96it/s]
 54%|█████▍    | 211/392 [00:17<00:15, 11.97it/s]
 54%|█████▍    | 213/392 [00:17<00:15, 11.90it/s]
 55%|█████▍    | 215/392 [00:17<00:14, 11.95it/s]
 55%|█████▌    | 217/392 [00:17<00:14, 11.87it/s]
 56%|█████▌    | 219/392 [00:18<00:14, 11.87it/s]
 56%|█████▋    | 221/392 [00:18<00:14, 11.89it/s]
 57%|█████▋    | 223/392 [00:18<00:14, 11.84it/s]
 57%|█████▋    | 225/392 [00:18<00:14, 11.67it/s]
 58%|█████▊    | 227/392 [00:18<00:14, 11.77it/s]
 58%|█████▊    | 229/392 [00:18<00:13, 11.82it/s]
 59%|█████▉    | 231/392 [00:19<00:13, 11.75it/s]
 59%|█████▉    | 233/392 [00:19<00:13, 11.67it/s]
 60%|█████▉    | 235/392 [00:19<00:13, 11.63it/s]
 60%|██████    | 237/392 [00:19<00:13, 11.69it/s]
 61%|██████    | 239/392 [00:19<00:13, 11.69it/s]
 61%|██████▏   | 241/392 [00:19<00:12, 11.72it/s]
 62%|██████▏   | 243/392 [00:20<00:12, 11.68it/s]
 62%|██████▎   | 245/392 [00:20<00:12, 11.67it/s]
 63%|██████▎   | 247/392 [00:20<00:12, 11.69it/s]
 64%|██████▎   | 249/392 [00:20<00:12, 11.61it/s]
 64%|██████▍   | 251/392 [00:20<00:12, 11.62it/s]
 65%|██████▍   | 253/392 [00:20<00:11, 11.69it/s]
 65%|██████▌   | 255/392 [00:21<00:11, 11.75it/s]
 66%|██████▌   | 257/392 [00:21<00:11, 11.61it/s]
 66%|██████▌   | 259/392 [00:21<00:11, 11.58it/s]
 67%|██████▋   | 261/392 [00:21<00:11, 11.65it/s]
 67%|██████▋   | 263/392 [00:21<00:11, 11.69it/s]
 68%|██████▊   | 265/392 [00:21<00:10, 11.62it/s]
 68%|██████▊   | 267/392 [00:22<00:10, 11.67it/s]
 69%|██████▊   | 269/392 [00:22<00:10, 11.78it/s]
 69%|██████▉   | 271/392 [00:22<00:10, 11.71it/s]
 70%|██████▉   | 273/392 [00:22<00:10, 11.59it/s]
 70%|███████   | 275/392 [00:22<00:10, 11.63it/s]
 71%|███████   | 277/392 [00:22<00:09, 11.70it/s]
 71%|███████   | 279/392 [00:23<00:09, 11.67it/s]
 72%|███████▏  | 281/392 [00:23<00:09, 11.67it/s]
 72%|███████▏  | 283/392 [00:23<00:09, 11.74it/s]
 73%|███████▎  | 285/392 [00:23<00:09, 11.75it/s]
 73%|███████▎  | 287/392 [00:23<00:09, 11.59it/s]
 74%|███████▎  | 289/392 [00:24<00:08, 11.47it/s]
 74%|███████▍  | 291/392 [00:24<00:08, 11.46it/s]
 75%|███████▍  | 293/392 [00:24<00:08, 11.56it/s]
 75%|███████▌  | 295/392 [00:24<00:08, 11.61it/s]
 76%|███████▌  | 297/392 [00:24<00:08, 11.66it/s]
 76%|███████▋  | 299/392 [00:24<00:08, 11.53it/s]
 77%|███████▋  | 301/392 [00:25<00:07, 11.61it/s]
 77%|███████▋  | 303/392 [00:25<00:07, 11.54it/s]
 78%|███████▊  | 305/392 [00:25<00:07, 11.57it/s]
 78%|███████▊  | 307/392 [00:25<00:07, 11.55it/s]
 79%|███████▉  | 309/392 [00:25<00:07, 11.65it/s]
 79%|███████▉  | 311/392 [00:25<00:06, 11.68it/s]
 80%|███████▉  | 313/392 [00:26<00:06, 11.65it/s]
 80%|████████  | 315/392 [00:26<00:06, 11.57it/s]
 81%|████████  | 317/392 [00:26<00:06, 11.59it/s]
 81%|████████▏ | 319/392 [00:26<00:06, 11.65it/s]
 82%|████████▏ | 321/392 [00:26<00:06, 11.63it/s]
 82%|████████▏ | 323/392 [00:26<00:05, 11.64it/s]
 83%|████████▎ | 325/392 [00:27<00:05, 11.69it/s]
 83%|████████▎ | 327/392 [00:27<00:05, 11.76it/s]
 84%|████████▍ | 329/392 [00:27<00:05, 11.68it/s]
 84%|████████▍ | 331/392 [00:27<00:05, 11.65it/s]
 85%|████████▍ | 333/392 [00:27<00:05, 11.60it/s]
 85%|████████▌ | 335/392 [00:27<00:04, 11.65it/s]
 86%|████████▌ | 337/392 [00:28<00:04, 11.65it/s]
 86%|████████▋ | 339/392 [00:28<00:04, 11.65it/s]
 87%|████████▋ | 341/392 [00:28<00:04, 11.69it/s]
 88%|████████▊ | 343/392 [00:28<00:04, 11.81it/s]
 88%|████████▊ | 345/392 [00:28<00:03, 11.81it/s]
 89%|████████▊ | 347/392 [00:28<00:03, 11.65it/s]
 89%|████████▉ | 349/392 [00:29<00:03, 11.60it/s]
 90%|████████▉ | 351/392 [00:29<00:03, 11.69it/s]
 90%|█████████ | 353/392 [00:29<00:03, 11.73it/s]
 91%|█████████ | 355/392 [00:29<00:03, 11.73it/s]
 91%|█████████ | 357/392 [00:29<00:02, 11.68it/s]
 92%|█████████▏| 359/392 [00:30<00:02, 11.71it/s]
 92%|█████████▏| 361/392 [00:30<00:02, 11.66it/s]
 93%|█████████▎| 363/392 [00:30<00:02, 11.71it/s]
 93%|█████████▎| 365/392 [00:30<00:02, 11.64it/s]
 94%|█████████▎| 367/392 [00:30<00:02, 11.67it/s]
 94%|█████████▍| 369/392 [00:30<00:01, 11.69it/s]
 95%|█████████▍| 371/392 [00:31<00:01, 11.66it/s]
 95%|█████████▌| 373/392 [00:31<00:01, 11.59it/s]
 96%|█████████▌| 375/392 [00:31<00:01, 11.64it/s]
 96%|█████████▌| 377/392 [00:31<00:01, 11.66it/s]
 97%|█████████▋| 379/392 [00:31<00:01, 11.63it/s]
 97%|█████████▋| 381/392 [00:31<00:00, 11.61it/s]
 98%|█████████▊| 383/392 [00:32<00:00, 11.63it/s]
 98%|█████████▊| 385/392 [00:32<00:00, 11.54it/s]
 99%|█████████▊| 387/392 [00:32<00:00, 11.52it/s]
 99%|█████████▉| 389/392 [00:32<00:00, 11.55it/s]
+FID:  11.042092769873989

experiments/bs256_lr1e-4/conv-layer/fid_cfg1_50k.tsv

@@ -0,0 +1,8 @@
+step	checkpoint	status	fid	cfg	num_requested	num_png	seed	timestamp_utc
+250000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/0250000.pt	ok	68.02801770968483	1.0	50000	50176	0	2026-08-12T02:49:22.920355+00:00
+500000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/0500000.pt	ok	56.895089833391694	1.0	50000	50176	0	2026-08-12T16:07:49.756361+00:00
+750000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/0750000.pt	ok	51.979633242087516	1.0	50000	50176	0	2026-08-13T05:38:42.872765+00:00
+1000000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1000000.pt	ok	49.241666747978115	1.0	50000	50176	0	2026-08-13T19:10:04.964165+00:00
+1250000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1250000.pt	ok	47.5775979137519	1.0	50000	50176	0	2026-08-14T08:27:28.159212+00:00
+1500000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1500000.pt	ok	46.296574099817235	1.0	50000	50176	0	2026-08-14T21:59:13.996458+00:00
+1750000	/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1750000.pt	ok	45.23863928744123	1.0	50000	50176	0	2026-08-15T11:18:22.287188+00:00

experiments/bs256_lr1e-4/fid_cfg1_50k_training_curves.csv

@@ -0,0 +1,73 @@
+model,step,epoch,fid,cfg,num_png,checkpoint,timestamp_utc,source
+Base,1000000,199.84012789768187,48.29311023150069,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1000000.pt,2026-07-22T03:33:11.273379+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1050000,209.83213429256594,48.018564188901564,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1050000.pt,2026-07-22T06:16:58.701669+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1100000,219.82414068745004,47.47688964066043,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1100000.pt,2026-07-22T09:01:01.876690+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1150000,229.81614708233414,47.10925915220656,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1150000.pt,2026-07-22T11:44:42.044226+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1200000,239.8081534772182,46.718045833545546,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1200000.pt,2026-07-22T14:28:44.897893+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1250000,249.8001598721023,46.414511520279234,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1250000.pt,2026-07-22T17:32:00.597820+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1300000,259.7921662669864,45.92057227506251,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1300000.pt,2026-07-22T20:16:10.465171+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1350000,269.7841726618705,45.89153691429402,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1350000.pt,2026-07-22T23:00:16.427930+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1400000,279.7761790567546,45.54178671851474,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1400000.pt,2026-07-23T01:44:23.027126+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1450000,289.7681854516387,45.116327425034854,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1450000.pt,2026-07-23T04:28:29.462587+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1500000,299.7601918465228,44.962781651113005,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1500000.pt,2026-07-23T07:12:39.286539+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1550000,309.75219824140686,44.842452348401196,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1550000.pt,2026-07-23T09:56:37.374904+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1600000,319.744204636291,44.640852044147266,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1600000.pt,2026-07-23T12:40:46.596453+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1650000,329.73621103117506,44.237907149190846,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1650000.pt,2026-07-23T15:24:58.652950+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1700000,339.72821742605913,44.136045077932806,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1700000.pt,2026-07-23T18:26:45.897450+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1750000,349.72022382094326,43.767861385718845,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1750000.pt,2026-07-23T21:10:59.723969+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1800000,359.71223021582733,43.64788291826136,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1800000.pt,2026-07-23T23:55:10.773687+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1850000,369.7042366107114,43.409102472303005,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1850000.pt,2026-07-24T02:39:16.581973+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1900000,379.69624300559553,43.21848172913133,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1900000.pt,2026-07-24T05:23:19.752006+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,1950000,389.6882494004796,42.921743510806834,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/1950000.pt,2026-07-24T08:07:41.369097+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2000000,399.68025579536373,42.80858523280904,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2000000.pt,2026-07-24T10:51:57.437664+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2050000,409.6722621902478,42.67307960077221,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2050000.pt,2026-07-24T13:36:12.546595+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2100000,419.6642685851319,42.443703960448886,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2100000.pt,2026-07-24T16:21:04.406803+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2150000,429.656274980016,42.50299269928496,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2150000.pt,2026-07-24T19:22:47.091546+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2200000,439.6482813749001,42.20738877449185,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2200000.pt,2026-07-27T04:37:51.175075+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2250000,449.64028776978415,42.09083954185155,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2250000.pt,2026-07-27T07:22:06.775572+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2300000,459.6322941646683,41.95050261156155,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2300000.pt,2026-07-27T10:06:30.595161+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2350000,469.62430055955235,41.72793962332611,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2350000.pt,2026-07-27T12:51:06.430298+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2400000,479.6163069544364,41.68562578798634,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2400000.pt,2026-07-27T15:35:39.897651+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2450000,489.60831334932055,41.43696452027939,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2450000.pt,2026-07-27T18:20:03.971453+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2500000,499.6003197442046,41.54235319764359,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2500000.pt,2026-07-27T21:04:36.195527+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2550000,509.59232613908875,41.33031665291526,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2550000.pt,2026-07-28T18:27:35.792552+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2600000,519.5843325339728,41.09229129923432,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2600000.pt,2026-07-28T21:12:00.197229+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2650000,529.5763389288569,41.04462000106378,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2650000.pt,2026-07-28T23:56:30.747960+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2700000,539.568345323741,40.986577043428895,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2700000.pt,2026-07-29T02:41:03.183436+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2750000,549.5603517186252,40.88498184813966,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2750000.pt,2026-07-29T05:25:36.916055+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2800000,559.5523581135092,40.794031115349355,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2800000.pt,2026-07-29T08:10:12.737171+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2850000,569.5443645083933,40.42894815467065,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2850000.pt,2026-07-29T10:54:46.079051+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2900000,579.5363709032774,40.422900985210845,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2900000.pt,2026-07-29T13:39:29.665685+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,2950000,589.5283772981614,40.54817889332162,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/2950000.pt,2026-07-29T16:46:35.078607+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3000000,599.5203836930456,40.46896141626422,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3000000.pt,2026-07-29T19:31:07.135111+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3050000,609.5123900879297,40.25837670401995,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3050000.pt,2026-07-29T22:15:37.950844+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3100000,619.5043964828137,40.21306397405647,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3100000.pt,2026-07-30T01:00:08.137327+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3150000,629.4964028776978,40.238701584494436,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3150000.pt,2026-07-30T03:44:41.736117+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3250000,649.480415667466,39.798277728671735,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3250000.pt,2026-07-30T09:08:00.799118+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3500000,699.4404476418865,39.5760757594021,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3500000.pt,2026-07-30T22:46:26.192799+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,3750000,749.4004796163069,39.16467116363083,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/3750000.pt,2026-07-31T12:03:12.844048+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Base,4000000,799.3605115907275,38.87082736834208,1.0,50176,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/4000000.pt,2026-08-01T01:39:07.070873+00:00,/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/fid_cfg1_50k.tsv
+Rotation-layer,250000,49.96003197442047,67.91903261534316,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/0250000.pt,2026-08-01T21:02:29.476441+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,500000,99.92006394884093,57.01010638417159,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/0500000.pt,2026-08-02T10:11:28.935049+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,750000,149.8800959232614,52.441390102274056,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/0750000.pt,2026-08-02T23:36:06.565998+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,1000000,199.84012789768187,49.64487908424047,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000000.pt,2026-08-03T12:45:14.796403+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,1000800,200.0,49.63772618622812,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000800.pt,2026-08-04T15:15:35.360973+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,1250000,249.8001598721023,47.591701349064124,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1250000.pt,2026-08-05T04:47:11.319539+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,1500000,299.7601918465228,46.44435903091937,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1500000.pt,2026-08-05T18:19:19.835092+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,1750000,349.72022382094326,45.36629934762789,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1750000.pt,2026-08-06T07:30:17.398008+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,2000000,399.68025579536373,44.46750925270834,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2000000.pt,2026-08-06T21:00:42.968130+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,2250000,449.64028776978415,43.70016640676096,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2250000.pt,2026-08-07T10:10:54.646029+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,2500000,499.6003197442046,43.01911391830271,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2500000.pt,2026-08-07T23:39:35.548958+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,2750000,549.5603517186252,42.307006340956036,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2750000.pt,2026-08-08T12:50:10.743320+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,3000000,599.5203836930456,41.88882400358102,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3000000.pt,2026-08-09T02:17:50.216261+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,3250000,649.480415667466,41.4511384135609,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3250000.pt,2026-08-09T15:28:36.464899+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,3500000,699.4404476418865,41.06018444188339,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3500000.pt,2026-08-10T04:56:51.985194+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,3750000,749.4004796163069,40.635836038709044,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3750000.pt,2026-08-10T18:25:12.497226+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Rotation-layer,4000000,799.3605115907275,40.273841521884094,1.0,50176,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/4000000.pt,2026-08-11T07:36:22.477469+00:00,/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/fid_cfg1_50k.tsv
+Conv-layer,250000,49.96003197442047,68.02801770968483,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/0250000.pt,2026-08-12T02:49:22.920355+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv
+Conv-layer,500000,99.92006394884093,56.895089833391694,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/0500000.pt,2026-08-12T16:07:49.756361+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv
+Conv-layer,750000,149.8800959232614,51.979633242087516,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/0750000.pt,2026-08-13T05:38:42.872765+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv
+Conv-layer,1000000,199.84012789768187,49.241666747978115,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1000000.pt,2026-08-13T19:10:04.964165+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv
+Conv-layer,1250000,249.8001598721023,47.5775979137519,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1250000.pt,2026-08-14T08:27:28.159212+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv
+Conv-layer,1500000,299.7601918465228,46.296574099817235,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1500000.pt,2026-08-14T21:59:13.996458+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv
+Conv-layer,1750000,349.72022382094326,45.23863928744123,1.0,50176,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1750000.pt,2026-08-15T11:18:22.287188+00:00,/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/fid_cfg1_50k.tsv

experiments/bs256_lr1e-4/rotation-layer/fid_cfg1_50k.tsv

@@ -0,0 +1,18 @@
+step	checkpoint	status	fid	cfg	num_requested	num_png	seed	timestamp_utc
+250000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/0250000.pt	ok	67.91903261534316	1.0	50000	50176	0	2026-08-01T21:02:29.476441+00:00
+500000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/0500000.pt	ok	57.01010638417159	1.0	50000	50176	0	2026-08-02T10:11:28.935049+00:00
+750000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/0750000.pt	ok	52.441390102274056	1.0	50000	50176	0	2026-08-02T23:36:06.565998+00:00
+1000000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000000.pt	ok	49.64487908424047	1.0	50000	50176	0	2026-08-03T12:45:14.796403+00:00
+1000800	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1000800.pt	ok	49.63772618622812	1.0	50000	50176	0	2026-08-04T15:15:35.360973+00:00
+1250000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1250000.pt	ok	47.591701349064124	1.0	50000	50176	0	2026-08-05T04:47:11.319539+00:00
+1500000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1500000.pt	ok	46.44435903091937	1.0	50000	50176	0	2026-08-05T18:19:19.835092+00:00
+1750000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/1750000.pt	ok	45.36629934762789	1.0	50000	50176	0	2026-08-06T07:30:17.398008+00:00
+2000000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2000000.pt	ok	44.46750925270834	1.0	50000	50176	0	2026-08-06T21:00:42.968130+00:00
+2250000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2250000.pt	ok	43.70016640676096	1.0	50000	50176	0	2026-08-07T10:10:54.646029+00:00
+2500000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2500000.pt	ok	43.01911391830271	1.0	50000	50176	0	2026-08-07T23:39:35.548958+00:00
+2750000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/2750000.pt	ok	42.307006340956036	1.0	50000	50176	0	2026-08-08T12:50:10.743320+00:00
+3000000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3000000.pt	ok	41.88882400358102	1.0	50000	50176	0	2026-08-09T02:17:50.216261+00:00
+3250000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3250000.pt	ok	41.4511384135609	1.0	50000	50176	0	2026-08-09T15:28:36.464899+00:00
+3500000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3500000.pt	ok	41.06018444188339	1.0	50000	50176	0	2026-08-10T04:56:51.985194+00:00
+3750000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/3750000.pt	ok	40.635836038709044	1.0	50000	50176	0	2026-08-10T18:25:12.497226+00:00
+4000000	/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/4000000.pt	ok	40.273841521884094	1.0	50000	50176	0	2026-08-11T07:36:22.477469+00:00

experiments/bs256_lr1e-4/rotation-layer/final_4003200_cfg4.0_pytorch_fid.txt

@@ -0,0 +1 @@
+FID:  11.188732265256363

experiments/bs256_lr1e-4/validation/checkpoint_validation.json

@@ -0,0 +1,101 @@
+{
+  "base": {
+    "source_path": "/home/nvidia/SiT-official-bs256/results-official-bs256-400k/000-SiT-S-2-Linear-velocity-None/checkpoints/4003200.pt",
+    "filename": "4003200.pt",
+    "bytes": 526844432,
+    "sha256": "81c9246463c94a3200e7d1861a95bd923e3f3cac92e4966d861a843348142f4e",
+    "metadata_fields": [
+      "args",
+      "ema",
+      "model",
+      "opt"
+    ],
+    "checkpoint_steps": null,
+    "configured_epochs": 800,
+    "global_batch_size": 256,
+    "learning_rate": null,
+    "model_module": "models",
+    "state_dict": "ema",
+    "learn_sigma": true,
+    "strict_load": true,
+    "forward_shape": [
+      1,
+      4,
+      32,
+      32
+    ],
+    "decoded_shape": [
+      1,
+      3,
+      256,
+      256
+    ],
+    "preview": "base_ema_euler10_seed20260815.png"
+  },
+  "rot-layer": {
+    "source_path": "/home/nvidia/SiT-rot-layer-bs256/results-200ep/SiT-S-2-RotLayer-bs256-lr1e-4-200ep/checkpoints/4003200.pt",
+    "filename": "4003200.pt",
+    "bytes": 526845456,
+    "sha256": "0d0a191eda63f82d6262f8b41644a3f9ceb0f3cfe66cc0accc711c8b488bbc32",
+    "metadata_fields": [
+      "args",
+      "ema",
+      "model",
+      "opt"
+    ],
+    "checkpoint_steps": null,
+    "configured_epochs": 800,
+    "global_batch_size": 256,
+    "learning_rate": 0.0001,
+    "model_module": "models_rot_layer",
+    "state_dict": "ema",
+    "learn_sigma": true,
+    "strict_load": true,
+    "forward_shape": [
+      1,
+      4,
+      32,
+      32
+    ],
+    "decoded_shape": [
+      1,
+      3,
+      256,
+      256
+    ],
+    "preview": "rot-layer_ema_euler10_seed20260815.png"
+  },
+  "conv-layer": {
+    "source_path": "/data/nvidia/SiT-conv-layer-bs256/results-800ep/SiT-S-2-ConvLayer-bs256-lr1e-4-800ep/checkpoints/1950000.pt",
+    "filename": "1950000.pt",
+    "bytes": 526985776,
+    "sha256": "f3724fa2651c6fbb3a624664f057c1dd56c658d68fb356c52cf51a7684ae7548",
+    "metadata_fields": [
+      "args",
+      "ema",
+      "model",
+      "opt"
+    ],
+    "checkpoint_steps": null,
+    "configured_epochs": 800,
+    "global_batch_size": 256,
+    "learning_rate": 0.0001,
+    "model_module": "models_conv",
+    "state_dict": "ema",
+    "learn_sigma": true,
+    "strict_load": true,
+    "forward_shape": [
+      1,
+      4,
+      32,
+      32
+    ],
+    "decoded_shape": [
+      1,
+      3,
+      256,
+      256
+    ],
+    "preview": "conv-layer_ema_euler10_seed20260815.png"
+  }
+}