Add files using upload-large-folder tool

README.md

@@ -0,0 +1,132 @@
+---
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - sit
+  - image-generation
+  - class-conditional
+  - imagenet
+license: mit
+inference: true
+---
+
+# SiT-diffusers
+
+Diffusers-ready checkpoints for **Scalable Interpolant Transformers (SiT)**, converted for local/offline use.
+
+This root folder is a model collection that contains:
+
+- `SiT-S-2-256-diffusers`
+- `SiT-B-2-256-diffusers`
+- `SiT-L-2-256-diffusers`
+- `SiT-XL-2-256-diffusers`
+- `SiT-XL-2-512-diffusers`
+
+Each subfolder is a self-contained Diffusers model repo with:
+
+- `pipeline.py`
+- `transformer/transformer_sit.py`
+- `scheduler/scheduling_flow_match_sit.py`
+- `transformer/diffusion_pytorch_model.safetensors`
+- `vae/diffusion_pytorch_model.safetensors`
+
+## Model Paths
+
+Use paths relative to this root README:
+
+| Model | Resolution | Local path |
+|---|---:|---|
+| SiT-S/2 | 256x256 | `./SiT-S-2-256-diffusers` |
+| SiT-B/2 | 256x256 | `./SiT-B-2-256-diffusers` |
+| SiT-L/2 | 256x256 | `./SiT-L-2-256-diffusers` |
+| SiT-XL/2 | 256x256 | `./SiT-XL-2-256-diffusers` |
+| SiT-XL/2 | 512x512 | `./SiT-XL-2-512-diffusers` |
+
+## Inference Demo (Diffusers)
+
+### 1) Load a local subfolder checkpoint
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+model_path = "./SiT-XL-2-512-diffusers"  # change to any path in the table above
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = DiffusionPipeline.from_pretrained(
+    model_path,
+    trust_remote_code=True,
+).to(device)
+
+generator = torch.Generator(device=device).manual_seed(0)
+
+# ImageNet class example: 207 = golden retriever
+result = pipe(
+    class_labels=207,
+    height=512,
+    width=512,
+    num_inference_steps=250,  # official SiT comparisons commonly use 250 steps
+    guidance_scale=4.0,
+    generator=generator,
+)
+
+image = result.images[0]
+image.save("sit_xl_512_demo.png")
+```
+
+### 2) Quick variant switch (256 models)
+
+```python
+model_path = "./SiT-S-2-256-diffusers"
+# model_path = "./SiT-B-2-256-diffusers"
+# model_path = "./SiT-L-2-256-diffusers"
+# model_path = "./SiT-XL-2-256-diffusers"
+
+pipe = DiffusionPipeline.from_pretrained(model_path, trust_remote_code=True).to(device)
+image = pipe(
+    class_labels=207,
+    height=256,
+    width=256,
+    num_inference_steps=250,
+    guidance_scale=4.0,
+    generator=generator,
+).images[0]
+image.save("sit_256_demo.png")
+```
+
+## FID Reference (from Official SiT Results)
+
+The table below summarizes widely cited SiT numbers from the official project materials for class-conditional ImageNet generation.
+
+| Model / setting | Resolution | FID-50K (lower is better) |
+|---|---:|---:|
+| SiT-S (400K steps) | 256x256 | 57.6 |
+| SiT-B (400K steps) | 256x256 | 33.5 |
+| SiT-L (400K steps) | 256x256 | 17.2 |
+| SiT-XL (400K steps) | 256x256 | 8.6 |
+| SiT-XL (cfg=1.5, ODE) | 256x256 | 2.15 |
+| SiT-XL (cfg=1.5, SDE, `w(t)=sigma_t`) | 256x256 | 2.06 |
+| SiT-XL (sample showcase) | 512x512 | Not reported in the same benchmark table |
+
+> Note: FID depends on training recipe, sampler choice (ODE/SDE), guidance scale, and evaluation protocol. Treat this table as a reference to official SiT reports, not as guaranteed reproducibility for every conversion/export.
+
+## Source and Paper
+
+- Official SiT code: [willisma/SiT](https://github.com/willisma/SiT)
+- Project page: [scalable-interpolant.github.io](https://scalable-interpolant.github.io/)
+- Paper (arXiv): [SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers](https://arxiv.org/abs/2401.08740)
+
+## Citation
+
+If you use SiT in your work, please cite:
+
+```bibtex
+@inproceedings{ma2024sit,
+  title={SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers},
+  author={Ma, Nanye and Goldstein, Mark and Albergo, Michael S. and Boffi, Nicholas M. and Vanden-Eijnden, Eric and Xie, Saining},
+  booktitle={European Conference on Computer Vision (ECCV)},
+  year={2024},
+  note={Accepted to ECCV 2024}
+}
+```

SiT-B-2-256-diffusers/README.md

@@ -0,0 +1,42 @@
+---
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - sit
+  - image-generation
+  - class-conditional
+inference: true
+---
+
+# SiT-B-2-256-diffusers
+
+Self-contained Diffusers checkpoint repo for SiT.
+
+## Usage
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("./").to("cuda" if torch.cuda.is_available() else "cpu")
+generator = torch.Generator(device=pipe.device).manual_seed(0)
+
+image = pipe(
+    class_labels=207,
+    height=256,
+    width=256,
+    num_inference_steps=250,
+    guidance_scale=4.0,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+## Components
+
+- `pipeline.py`
+- `transformer/transformer_sit.py`
+- `scheduler/scheduling_flow_match_sit.py`
+- `transformer/diffusion_pytorch_model.safetensors`
+- `vae/diffusion_pytorch_model.safetensors`

SiT-B-2-256-diffusers/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "SiTPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
+  "scheduler": [
+    "scheduling_flow_match_sit",
+    "SiTFlowMatchScheduler"
+  ],
+  "transformer": [
+    "transformer_sit",
+    "SiTTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

SiT-B-2-256-diffusers/pipeline.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+import torch
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class SiTPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, scheduler, vae):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
+        self.vae_scale_factor = 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        class_labels: Union[int, List[int]] = 207,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 250,
+        guidance_scale: float = 4.0,
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        device = self._execution_device
+        if isinstance(class_labels, int):
+            class_labels = [class_labels]
+        batch_size = len(class_labels)
+
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        latents = randn_tensor(
+            (batch_size, self.transformer.config.in_channels, latent_h, latent_w),
+            generator=generator,
+            device=device,
+            dtype=self.transformer.dtype,
+        )
+
+        labels = torch.tensor(class_labels, device=device, dtype=torch.long)
+        do_cfg = guidance_scale is not None and guidance_scale > 1.0
+        if do_cfg:
+            null_label = torch.full((batch_size,), self.transformer.config.num_classes, device=device, dtype=torch.long)
+            labels = torch.cat([labels, null_label], dim=0)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        for t in self.progress_bar(timesteps):
+            t_batch = torch.full((batch_size,), t, device=device, dtype=latents.dtype)
+            model_input = latents
+            if do_cfg:
+                model_input = torch.cat([latents, latents], dim=0)
+                t_batch = torch.cat([t_batch, t_batch], dim=0)
+
+            model_pred = self.transformer(
+                hidden_states=model_input,
+                timestep=t_batch,
+                class_labels=labels,
+            ).sample
+
+            if do_cfg:
+                cond, uncond = model_pred.chunk(2, dim=0)
+                model_pred = uncond + guidance_scale * (cond - uncond)
+
+            latents = self.scheduler.step(model_pred, t, latents, generator=generator).prev_sample
+
+        image = self.vae.decode(latents / 0.18215).sample
+        # Keep PyTorch outputs in raw VAE range [-1, 1] to match original SiT scripts.
+        if output_type == "pt":
+            image = image
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

SiT-B-2-256-diffusers/scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "SiTFlowMatchScheduler",
+  "_diffusers_version": "0.36.0",
+  "diffusion_form": "sigma",
+  "diffusion_norm": 1.0,
+  "mode": "ode",
+  "num_train_timesteps": 1000,
+  "shift": 1.0
+}

SiT-B-2-256-diffusers/scheduler/scheduling_flow_match_sit.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SiTFlowMatchSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class SiTFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        mode: str = "ode",
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        diffusion_form: str = "sigma",
+        diffusion_norm: float = 1.0,
+    ):
+        self.timesteps = None
+        self.sigmas = None
+        self._step_index = None
+
+    def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None):
+        # Flow matching integrates from noise (t=0) to data (t=1).
+        ts = torch.linspace(0.0, 1.0, num_inference_steps + 1, device=device, dtype=torch.float32)
+        self.timesteps = ts[:-1]
+        self.sigmas = 1.0 - self.timesteps
+        self._step_index = 0
+        return self.timesteps
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def _diffusion(self, t: torch.Tensor) -> torch.Tensor:
+        form = self.config.diffusion_form
+        norm = self.config.diffusion_norm
+        if form == "constant":
+            return torch.full_like(t, norm)
+        if form == "sigma":
+            return norm * (1.0 - t)
+        if form == "linear":
+            return norm * (1.0 - t)
+        if form == "decreasing":
+            return 0.25 * (norm * torch.cos(torch.pi * t) + 1) ** 2
+        if form == "increasing-decreasing":
+            return norm * torch.sin(torch.pi * t) ** 2
+        # "SBDM" approximated with sigma-based schedule for compatibility.
+        return norm * (1.0 - t)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SiTFlowMatchSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None:
+            raise ValueError("Call `set_timesteps` before `step`.")
+        if self._step_index is None:
+            self._step_index = 0
+
+        step_index = min(self._step_index, len(self.timesteps) - 1)
+        t = self.timesteps[step_index].to(sample.device)
+        next_t = 1.0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1].to(sample.device)
+        dt = next_t - t
+
+        prev_sample = sample + model_output * dt
+        if self.config.mode.lower() == "sde":
+            diffusion = self._diffusion(torch.full((sample.shape[0],), t, device=sample.device, dtype=sample.dtype))
+            while diffusion.dim() < sample.dim():
+                diffusion = diffusion.unsqueeze(-1)
+            noise = torch.randn(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = prev_sample + torch.sqrt(torch.clamp(2.0 * diffusion * torch.abs(dt), min=0.0)) * noise
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return SiTFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        sigma = (1.0 - timesteps).view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1 - sigma) * original_samples + sigma * noise

SiT-B-2-256-diffusers/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "SiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "class_dropout_prob": 0.1,
+  "depth": 12,
+  "hidden_size": 768,
+  "in_channels": 4,
+  "input_size": 32,
+  "learn_sigma": true,
+  "mlp_ratio": 4.0,
+  "num_classes": 1000,
+  "num_heads": 12,
+  "patch_size": 2
+}

SiT-B-2-256-diffusers/transformer/diffusion_pytorch_model.safetensors

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

SiT-B-2-256-diffusers/transformer/transformer_sit.py

@@ -0,0 +1,224 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import Attention, Mlp, PatchEmbed
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+@dataclass
+class SiTTransformer2DModelOutput(BaseOutput):
+    sample: torch.Tensor
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 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: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        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: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size))
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
+        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: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        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: torch.Tensor,
+        train: bool,
+        force_drop_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        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)
+        return self.embedding_table(labels)
+
+
+class SiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        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):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        return self.linear(x)
+
+
+class SiTTransformer2DModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        input_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        hidden_size: int = 1152,
+        depth: int = 28,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        class_dropout_prob: float = 0.1,
+        num_classes: int = 1000,
+        learn_sigma: bool = 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_classes = num_classes
+
+        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
+        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) -> None:
+        def _basic_init(module: nn.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)
+        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))
+
+        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)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+        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: torch.Tensor) -> torch.Tensor:
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        force_drop_ids: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> SiTTransformer2DModelOutput:
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(class_labels, self.training, force_drop_ids=force_drop_ids)
+        c = t + y
+        for block in self.blocks:
+            x = block(x, c)
+        x = self.final_layer(x, c)
+        x = self.unpatchify(x)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        if not return_dict:
+            return (x,)
+        return SiTTransformer2DModelOutput(sample=x)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0):
+    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)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray):
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)

SiT-B-2-256-diffusers/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

SiT-B-2-256-diffusers/vae/diffusion_pytorch_model.safetensors

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

SiT-L-2-256-diffusers/README.md

@@ -0,0 +1,42 @@
+---
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - sit
+  - image-generation
+  - class-conditional
+inference: true
+---
+
+# SiT-L-2-256-diffusers
+
+Self-contained Diffusers checkpoint repo for SiT.
+
+## Usage
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("./").to("cuda" if torch.cuda.is_available() else "cpu")
+generator = torch.Generator(device=pipe.device).manual_seed(0)
+
+image = pipe(
+    class_labels=207,
+    height=256,
+    width=256,
+    num_inference_steps=250,
+    guidance_scale=4.0,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+## Components
+
+- `pipeline.py`
+- `transformer/transformer_sit.py`
+- `scheduler/scheduling_flow_match_sit.py`
+- `transformer/diffusion_pytorch_model.safetensors`
+- `vae/diffusion_pytorch_model.safetensors`

SiT-L-2-256-diffusers/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "SiTPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
+  "scheduler": [
+    "scheduling_flow_match_sit",
+    "SiTFlowMatchScheduler"
+  ],
+  "transformer": [
+    "transformer_sit",
+    "SiTTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

SiT-L-2-256-diffusers/pipeline.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+import torch
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class SiTPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, scheduler, vae):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
+        self.vae_scale_factor = 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        class_labels: Union[int, List[int]] = 207,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 250,
+        guidance_scale: float = 4.0,
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        device = self._execution_device
+        if isinstance(class_labels, int):
+            class_labels = [class_labels]
+        batch_size = len(class_labels)
+
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        latents = randn_tensor(
+            (batch_size, self.transformer.config.in_channels, latent_h, latent_w),
+            generator=generator,
+            device=device,
+            dtype=self.transformer.dtype,
+        )
+
+        labels = torch.tensor(class_labels, device=device, dtype=torch.long)
+        do_cfg = guidance_scale is not None and guidance_scale > 1.0
+        if do_cfg:
+            null_label = torch.full((batch_size,), self.transformer.config.num_classes, device=device, dtype=torch.long)
+            labels = torch.cat([labels, null_label], dim=0)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        for t in self.progress_bar(timesteps):
+            t_batch = torch.full((batch_size,), t, device=device, dtype=latents.dtype)
+            model_input = latents
+            if do_cfg:
+                model_input = torch.cat([latents, latents], dim=0)
+                t_batch = torch.cat([t_batch, t_batch], dim=0)
+
+            model_pred = self.transformer(
+                hidden_states=model_input,
+                timestep=t_batch,
+                class_labels=labels,
+            ).sample
+
+            if do_cfg:
+                cond, uncond = model_pred.chunk(2, dim=0)
+                model_pred = uncond + guidance_scale * (cond - uncond)
+
+            latents = self.scheduler.step(model_pred, t, latents, generator=generator).prev_sample
+
+        image = self.vae.decode(latents / 0.18215).sample
+        # Keep PyTorch outputs in raw VAE range [-1, 1] to match original SiT scripts.
+        if output_type == "pt":
+            image = image
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

SiT-L-2-256-diffusers/scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "SiTFlowMatchScheduler",
+  "_diffusers_version": "0.36.0",
+  "diffusion_form": "sigma",
+  "diffusion_norm": 1.0,
+  "mode": "ode",
+  "num_train_timesteps": 1000,
+  "shift": 1.0
+}

SiT-L-2-256-diffusers/scheduler/scheduling_flow_match_sit.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SiTFlowMatchSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class SiTFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        mode: str = "ode",
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        diffusion_form: str = "sigma",
+        diffusion_norm: float = 1.0,
+    ):
+        self.timesteps = None
+        self.sigmas = None
+        self._step_index = None
+
+    def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None):
+        # Flow matching integrates from noise (t=0) to data (t=1).
+        ts = torch.linspace(0.0, 1.0, num_inference_steps + 1, device=device, dtype=torch.float32)
+        self.timesteps = ts[:-1]
+        self.sigmas = 1.0 - self.timesteps
+        self._step_index = 0
+        return self.timesteps
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def _diffusion(self, t: torch.Tensor) -> torch.Tensor:
+        form = self.config.diffusion_form
+        norm = self.config.diffusion_norm
+        if form == "constant":
+            return torch.full_like(t, norm)
+        if form == "sigma":
+            return norm * (1.0 - t)
+        if form == "linear":
+            return norm * (1.0 - t)
+        if form == "decreasing":
+            return 0.25 * (norm * torch.cos(torch.pi * t) + 1) ** 2
+        if form == "increasing-decreasing":
+            return norm * torch.sin(torch.pi * t) ** 2
+        # "SBDM" approximated with sigma-based schedule for compatibility.
+        return norm * (1.0 - t)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SiTFlowMatchSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None:
+            raise ValueError("Call `set_timesteps` before `step`.")
+        if self._step_index is None:
+            self._step_index = 0
+
+        step_index = min(self._step_index, len(self.timesteps) - 1)
+        t = self.timesteps[step_index].to(sample.device)
+        next_t = 1.0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1].to(sample.device)
+        dt = next_t - t
+
+        prev_sample = sample + model_output * dt
+        if self.config.mode.lower() == "sde":
+            diffusion = self._diffusion(torch.full((sample.shape[0],), t, device=sample.device, dtype=sample.dtype))
+            while diffusion.dim() < sample.dim():
+                diffusion = diffusion.unsqueeze(-1)
+            noise = torch.randn(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = prev_sample + torch.sqrt(torch.clamp(2.0 * diffusion * torch.abs(dt), min=0.0)) * noise
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return SiTFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        sigma = (1.0 - timesteps).view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1 - sigma) * original_samples + sigma * noise

SiT-L-2-256-diffusers/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "SiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "class_dropout_prob": 0.1,
+  "depth": 24,
+  "hidden_size": 1024,
+  "in_channels": 4,
+  "input_size": 32,
+  "learn_sigma": true,
+  "mlp_ratio": 4.0,
+  "num_classes": 1000,
+  "num_heads": 16,
+  "patch_size": 2
+}

SiT-L-2-256-diffusers/transformer/transformer_sit.py

@@ -0,0 +1,224 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import Attention, Mlp, PatchEmbed
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+@dataclass
+class SiTTransformer2DModelOutput(BaseOutput):
+    sample: torch.Tensor
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 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: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        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: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size))
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
+        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: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        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: torch.Tensor,
+        train: bool,
+        force_drop_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        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)
+        return self.embedding_table(labels)
+
+
+class SiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        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):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        return self.linear(x)
+
+
+class SiTTransformer2DModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        input_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        hidden_size: int = 1152,
+        depth: int = 28,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        class_dropout_prob: float = 0.1,
+        num_classes: int = 1000,
+        learn_sigma: bool = 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_classes = num_classes
+
+        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
+        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) -> None:
+        def _basic_init(module: nn.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)
+        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))
+
+        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)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+        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: torch.Tensor) -> torch.Tensor:
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        force_drop_ids: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> SiTTransformer2DModelOutput:
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(class_labels, self.training, force_drop_ids=force_drop_ids)
+        c = t + y
+        for block in self.blocks:
+            x = block(x, c)
+        x = self.final_layer(x, c)
+        x = self.unpatchify(x)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        if not return_dict:
+            return (x,)
+        return SiTTransformer2DModelOutput(sample=x)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0):
+    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)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray):
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)

SiT-L-2-256-diffusers/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

SiT-L-2-256-diffusers/vae/diffusion_pytorch_model.safetensors

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

SiT-S-2-256-diffusers/README.md

@@ -0,0 +1,42 @@
+---
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - sit
+  - image-generation
+  - class-conditional
+inference: true
+---
+
+# SiT-S-2-256-diffusers
+
+Self-contained Diffusers checkpoint repo for SiT.
+
+## Usage
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("./").to("cuda" if torch.cuda.is_available() else "cpu")
+generator = torch.Generator(device=pipe.device).manual_seed(0)
+
+image = pipe(
+    class_labels=207,
+    height=256,
+    width=256,
+    num_inference_steps=250,
+    guidance_scale=4.0,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+## Components
+
+- `pipeline.py`
+- `transformer/transformer_sit.py`
+- `scheduler/scheduling_flow_match_sit.py`
+- `transformer/diffusion_pytorch_model.safetensors`
+- `vae/diffusion_pytorch_model.safetensors`

SiT-S-2-256-diffusers/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "SiTPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
+  "scheduler": [
+    "scheduling_flow_match_sit",
+    "SiTFlowMatchScheduler"
+  ],
+  "transformer": [
+    "transformer_sit",
+    "SiTTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

SiT-S-2-256-diffusers/pipeline.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+import torch
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class SiTPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, scheduler, vae):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
+        self.vae_scale_factor = 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        class_labels: Union[int, List[int]] = 207,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 250,
+        guidance_scale: float = 4.0,
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        device = self._execution_device
+        if isinstance(class_labels, int):
+            class_labels = [class_labels]
+        batch_size = len(class_labels)
+
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        latents = randn_tensor(
+            (batch_size, self.transformer.config.in_channels, latent_h, latent_w),
+            generator=generator,
+            device=device,
+            dtype=self.transformer.dtype,
+        )
+
+        labels = torch.tensor(class_labels, device=device, dtype=torch.long)
+        do_cfg = guidance_scale is not None and guidance_scale > 1.0
+        if do_cfg:
+            null_label = torch.full((batch_size,), self.transformer.config.num_classes, device=device, dtype=torch.long)
+            labels = torch.cat([labels, null_label], dim=0)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        for t in self.progress_bar(timesteps):
+            t_batch = torch.full((batch_size,), t, device=device, dtype=latents.dtype)
+            model_input = latents
+            if do_cfg:
+                model_input = torch.cat([latents, latents], dim=0)
+                t_batch = torch.cat([t_batch, t_batch], dim=0)
+
+            model_pred = self.transformer(
+                hidden_states=model_input,
+                timestep=t_batch,
+                class_labels=labels,
+            ).sample
+
+            if do_cfg:
+                cond, uncond = model_pred.chunk(2, dim=0)
+                model_pred = uncond + guidance_scale * (cond - uncond)
+
+            latents = self.scheduler.step(model_pred, t, latents, generator=generator).prev_sample
+
+        image = self.vae.decode(latents / 0.18215).sample
+        # Keep PyTorch outputs in raw VAE range [-1, 1] to match original SiT scripts.
+        if output_type == "pt":
+            image = image
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

SiT-S-2-256-diffusers/scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "SiTFlowMatchScheduler",
+  "_diffusers_version": "0.36.0",
+  "diffusion_form": "sigma",
+  "diffusion_norm": 1.0,
+  "mode": "ode",
+  "num_train_timesteps": 1000,
+  "shift": 1.0
+}

SiT-S-2-256-diffusers/scheduler/scheduling_flow_match_sit.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SiTFlowMatchSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class SiTFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        mode: str = "ode",
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        diffusion_form: str = "sigma",
+        diffusion_norm: float = 1.0,
+    ):
+        self.timesteps = None
+        self.sigmas = None
+        self._step_index = None
+
+    def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None):
+        # Flow matching integrates from noise (t=0) to data (t=1).
+        ts = torch.linspace(0.0, 1.0, num_inference_steps + 1, device=device, dtype=torch.float32)
+        self.timesteps = ts[:-1]
+        self.sigmas = 1.0 - self.timesteps
+        self._step_index = 0
+        return self.timesteps
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def _diffusion(self, t: torch.Tensor) -> torch.Tensor:
+        form = self.config.diffusion_form
+        norm = self.config.diffusion_norm
+        if form == "constant":
+            return torch.full_like(t, norm)
+        if form == "sigma":
+            return norm * (1.0 - t)
+        if form == "linear":
+            return norm * (1.0 - t)
+        if form == "decreasing":
+            return 0.25 * (norm * torch.cos(torch.pi * t) + 1) ** 2
+        if form == "increasing-decreasing":
+            return norm * torch.sin(torch.pi * t) ** 2
+        # "SBDM" approximated with sigma-based schedule for compatibility.
+        return norm * (1.0 - t)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SiTFlowMatchSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None:
+            raise ValueError("Call `set_timesteps` before `step`.")
+        if self._step_index is None:
+            self._step_index = 0
+
+        step_index = min(self._step_index, len(self.timesteps) - 1)
+        t = self.timesteps[step_index].to(sample.device)
+        next_t = 1.0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1].to(sample.device)
+        dt = next_t - t
+
+        prev_sample = sample + model_output * dt
+        if self.config.mode.lower() == "sde":
+            diffusion = self._diffusion(torch.full((sample.shape[0],), t, device=sample.device, dtype=sample.dtype))
+            while diffusion.dim() < sample.dim():
+                diffusion = diffusion.unsqueeze(-1)
+            noise = torch.randn(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = prev_sample + torch.sqrt(torch.clamp(2.0 * diffusion * torch.abs(dt), min=0.0)) * noise
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return SiTFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        sigma = (1.0 - timesteps).view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1 - sigma) * original_samples + sigma * noise

SiT-S-2-256-diffusers/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "SiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "class_dropout_prob": 0.1,
+  "depth": 12,
+  "hidden_size": 384,
+  "in_channels": 4,
+  "input_size": 32,
+  "learn_sigma": true,
+  "mlp_ratio": 4.0,
+  "num_classes": 1000,
+  "num_heads": 6,
+  "patch_size": 2
+}

SiT-S-2-256-diffusers/transformer/diffusion_pytorch_model.safetensors

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

SiT-S-2-256-diffusers/transformer/transformer_sit.py

@@ -0,0 +1,224 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import Attention, Mlp, PatchEmbed
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+@dataclass
+class SiTTransformer2DModelOutput(BaseOutput):
+    sample: torch.Tensor
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 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: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        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: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size))
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
+        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: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        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: torch.Tensor,
+        train: bool,
+        force_drop_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        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)
+        return self.embedding_table(labels)
+
+
+class SiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        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):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        return self.linear(x)
+
+
+class SiTTransformer2DModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        input_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        hidden_size: int = 1152,
+        depth: int = 28,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        class_dropout_prob: float = 0.1,
+        num_classes: int = 1000,
+        learn_sigma: bool = 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_classes = num_classes
+
+        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
+        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) -> None:
+        def _basic_init(module: nn.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)
+        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))
+
+        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)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+        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: torch.Tensor) -> torch.Tensor:
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        force_drop_ids: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> SiTTransformer2DModelOutput:
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(class_labels, self.training, force_drop_ids=force_drop_ids)
+        c = t + y
+        for block in self.blocks:
+            x = block(x, c)
+        x = self.final_layer(x, c)
+        x = self.unpatchify(x)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        if not return_dict:
+            return (x,)
+        return SiTTransformer2DModelOutput(sample=x)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0):
+    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)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray):
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)

SiT-S-2-256-diffusers/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

SiT-S-2-256-diffusers/vae/diffusion_pytorch_model.safetensors

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

SiT-XL-2-256-diffusers/README.md

@@ -0,0 +1,42 @@
+---
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - sit
+  - image-generation
+  - class-conditional
+inference: true
+---
+
+# SiT-XL-2-256-diffusers
+
+Self-contained Diffusers checkpoint repo for SiT.
+
+## Usage
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("./").to("cuda" if torch.cuda.is_available() else "cpu")
+generator = torch.Generator(device=pipe.device).manual_seed(0)
+
+image = pipe(
+    class_labels=207,
+    height=256,
+    width=256,
+    num_inference_steps=250,
+    guidance_scale=4.0,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+## Components
+
+- `pipeline.py`
+- `transformer/transformer_sit.py`
+- `scheduler/scheduling_flow_match_sit.py`
+- `transformer/diffusion_pytorch_model.safetensors`
+- `vae/diffusion_pytorch_model.safetensors`

SiT-XL-2-256-diffusers/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "SiTPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
+  "scheduler": [
+    "scheduling_flow_match_sit",
+    "SiTFlowMatchScheduler"
+  ],
+  "transformer": [
+    "transformer_sit",
+    "SiTTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

SiT-XL-2-256-diffusers/pipeline.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+import torch
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class SiTPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, scheduler, vae):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
+        self.vae_scale_factor = 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        class_labels: Union[int, List[int]] = 207,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 250,
+        guidance_scale: float = 4.0,
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        device = self._execution_device
+        if isinstance(class_labels, int):
+            class_labels = [class_labels]
+        batch_size = len(class_labels)
+
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        latents = randn_tensor(
+            (batch_size, self.transformer.config.in_channels, latent_h, latent_w),
+            generator=generator,
+            device=device,
+            dtype=self.transformer.dtype,
+        )
+
+        labels = torch.tensor(class_labels, device=device, dtype=torch.long)
+        do_cfg = guidance_scale is not None and guidance_scale > 1.0
+        if do_cfg:
+            null_label = torch.full((batch_size,), self.transformer.config.num_classes, device=device, dtype=torch.long)
+            labels = torch.cat([labels, null_label], dim=0)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        for t in self.progress_bar(timesteps):
+            t_batch = torch.full((batch_size,), t, device=device, dtype=latents.dtype)
+            model_input = latents
+            if do_cfg:
+                model_input = torch.cat([latents, latents], dim=0)
+                t_batch = torch.cat([t_batch, t_batch], dim=0)
+
+            model_pred = self.transformer(
+                hidden_states=model_input,
+                timestep=t_batch,
+                class_labels=labels,
+            ).sample
+
+            if do_cfg:
+                cond, uncond = model_pred.chunk(2, dim=0)
+                model_pred = uncond + guidance_scale * (cond - uncond)
+
+            latents = self.scheduler.step(model_pred, t, latents, generator=generator).prev_sample
+
+        image = self.vae.decode(latents / 0.18215).sample
+        # Keep PyTorch outputs in raw VAE range [-1, 1] to match original SiT scripts.
+        if output_type == "pt":
+            image = image
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

SiT-XL-2-256-diffusers/scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "SiTFlowMatchScheduler",
+  "_diffusers_version": "0.36.0",
+  "diffusion_form": "sigma",
+  "diffusion_norm": 1.0,
+  "mode": "ode",
+  "num_train_timesteps": 1000,
+  "shift": 1.0
+}

SiT-XL-2-256-diffusers/scheduler/scheduling_flow_match_sit.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SiTFlowMatchSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class SiTFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        mode: str = "ode",
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        diffusion_form: str = "sigma",
+        diffusion_norm: float = 1.0,
+    ):
+        self.timesteps = None
+        self.sigmas = None
+        self._step_index = None
+
+    def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None):
+        # Flow matching integrates from noise (t=0) to data (t=1).
+        ts = torch.linspace(0.0, 1.0, num_inference_steps + 1, device=device, dtype=torch.float32)
+        self.timesteps = ts[:-1]
+        self.sigmas = 1.0 - self.timesteps
+        self._step_index = 0
+        return self.timesteps
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def _diffusion(self, t: torch.Tensor) -> torch.Tensor:
+        form = self.config.diffusion_form
+        norm = self.config.diffusion_norm
+        if form == "constant":
+            return torch.full_like(t, norm)
+        if form == "sigma":
+            return norm * (1.0 - t)
+        if form == "linear":
+            return norm * (1.0 - t)
+        if form == "decreasing":
+            return 0.25 * (norm * torch.cos(torch.pi * t) + 1) ** 2
+        if form == "increasing-decreasing":
+            return norm * torch.sin(torch.pi * t) ** 2
+        # "SBDM" approximated with sigma-based schedule for compatibility.
+        return norm * (1.0 - t)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SiTFlowMatchSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None:
+            raise ValueError("Call `set_timesteps` before `step`.")
+        if self._step_index is None:
+            self._step_index = 0
+
+        step_index = min(self._step_index, len(self.timesteps) - 1)
+        t = self.timesteps[step_index].to(sample.device)
+        next_t = 1.0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1].to(sample.device)
+        dt = next_t - t
+
+        prev_sample = sample + model_output * dt
+        if self.config.mode.lower() == "sde":
+            diffusion = self._diffusion(torch.full((sample.shape[0],), t, device=sample.device, dtype=sample.dtype))
+            while diffusion.dim() < sample.dim():
+                diffusion = diffusion.unsqueeze(-1)
+            noise = torch.randn(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = prev_sample + torch.sqrt(torch.clamp(2.0 * diffusion * torch.abs(dt), min=0.0)) * noise
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return SiTFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        sigma = (1.0 - timesteps).view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1 - sigma) * original_samples + sigma * noise

SiT-XL-2-256-diffusers/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "SiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "class_dropout_prob": 0.1,
+  "depth": 28,
+  "hidden_size": 1152,
+  "in_channels": 4,
+  "input_size": 32,
+  "learn_sigma": true,
+  "mlp_ratio": 4.0,
+  "num_classes": 1000,
+  "num_heads": 16,
+  "patch_size": 2
+}

SiT-XL-2-256-diffusers/transformer/diffusion_pytorch_model.safetensors

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

SiT-XL-2-256-diffusers/transformer/transformer_sit.py

@@ -0,0 +1,224 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import Attention, Mlp, PatchEmbed
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+@dataclass
+class SiTTransformer2DModelOutput(BaseOutput):
+    sample: torch.Tensor
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 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: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        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: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size))
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
+        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: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        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: torch.Tensor,
+        train: bool,
+        force_drop_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        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)
+        return self.embedding_table(labels)
+
+
+class SiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        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):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        return self.linear(x)
+
+
+class SiTTransformer2DModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        input_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        hidden_size: int = 1152,
+        depth: int = 28,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        class_dropout_prob: float = 0.1,
+        num_classes: int = 1000,
+        learn_sigma: bool = 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_classes = num_classes
+
+        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
+        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) -> None:
+        def _basic_init(module: nn.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)
+        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))
+
+        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)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+        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: torch.Tensor) -> torch.Tensor:
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        force_drop_ids: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> SiTTransformer2DModelOutput:
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(class_labels, self.training, force_drop_ids=force_drop_ids)
+        c = t + y
+        for block in self.blocks:
+            x = block(x, c)
+        x = self.final_layer(x, c)
+        x = self.unpatchify(x)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        if not return_dict:
+            return (x,)
+        return SiTTransformer2DModelOutput(sample=x)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0):
+    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)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray):
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)

SiT-XL-2-256-diffusers/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

SiT-XL-2-256-diffusers/vae/diffusion_pytorch_model.safetensors

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

SiT-XL-2-512-diffusers/README.md

@@ -0,0 +1,42 @@
+---
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - sit
+  - image-generation
+  - class-conditional
+inference: true
+---
+
+# SiT-XL-2-512-diffusers
+
+Self-contained Diffusers checkpoint repo for SiT.
+
+## Usage
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+pipe = DiffusionPipeline.from_pretrained("./").to("cuda" if torch.cuda.is_available() else "cpu")
+generator = torch.Generator(device=pipe.device).manual_seed(0)
+
+image = pipe(
+    class_labels=207,
+    height=512,
+    width=512,
+    num_inference_steps=250,
+    guidance_scale=4.0,
+    generator=generator,
+).images[0]
+image.save("demo.png")
+```
+
+## Components
+
+- `pipeline.py`
+- `transformer/transformer_sit.py`
+- `scheduler/scheduling_flow_match_sit.py`
+- `transformer/diffusion_pytorch_model.safetensors`
+- `vae/diffusion_pytorch_model.safetensors`

SiT-XL-2-512-diffusers/model_index.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": [
+    "pipeline",
+    "SiTPipeline"
+  ],
+  "_diffusers_version": "0.36.0",
+  "scheduler": [
+    "scheduling_flow_match_sit",
+    "SiTFlowMatchScheduler"
+  ],
+  "transformer": [
+    "transformer_sit",
+    "SiTTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

SiT-XL-2-512-diffusers/pipeline.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+import torch
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+
+class SiTPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer->vae"
+
+    def __init__(self, transformer, scheduler, vae):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
+        self.vae_scale_factor = 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        class_labels: Union[int, List[int]] = 207,
+        height: int = 256,
+        width: int = 256,
+        num_inference_steps: int = 250,
+        guidance_scale: float = 4.0,
+        generator: Optional[torch.Generator] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ):
+        device = self._execution_device
+        if isinstance(class_labels, int):
+            class_labels = [class_labels]
+        batch_size = len(class_labels)
+
+        latent_h = height // self.vae_scale_factor
+        latent_w = width // self.vae_scale_factor
+        latents = randn_tensor(
+            (batch_size, self.transformer.config.in_channels, latent_h, latent_w),
+            generator=generator,
+            device=device,
+            dtype=self.transformer.dtype,
+        )
+
+        labels = torch.tensor(class_labels, device=device, dtype=torch.long)
+        do_cfg = guidance_scale is not None and guidance_scale > 1.0
+        if do_cfg:
+            null_label = torch.full((batch_size,), self.transformer.config.num_classes, device=device, dtype=torch.long)
+            labels = torch.cat([labels, null_label], dim=0)
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        for t in self.progress_bar(timesteps):
+            t_batch = torch.full((batch_size,), t, device=device, dtype=latents.dtype)
+            model_input = latents
+            if do_cfg:
+                model_input = torch.cat([latents, latents], dim=0)
+                t_batch = torch.cat([t_batch, t_batch], dim=0)
+
+            model_pred = self.transformer(
+                hidden_states=model_input,
+                timestep=t_batch,
+                class_labels=labels,
+            ).sample
+
+            if do_cfg:
+                cond, uncond = model_pred.chunk(2, dim=0)
+                model_pred = uncond + guidance_scale * (cond - uncond)
+
+            latents = self.scheduler.step(model_pred, t, latents, generator=generator).prev_sample
+
+        image = self.vae.decode(latents / 0.18215).sample
+        # Keep PyTorch outputs in raw VAE range [-1, 1] to match original SiT scripts.
+        if output_type == "pt":
+            image = image
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        if not return_dict:
+            return (image,)
+        return ImagePipelineOutput(images=image)

SiT-XL-2-512-diffusers/scheduler/scheduler_config.json

@@ -0,0 +1,9 @@
+{
+  "_class_name": "SiTFlowMatchScheduler",
+  "_diffusers_version": "0.36.0",
+  "diffusion_form": "sigma",
+  "diffusion_norm": 1.0,
+  "mode": "ode",
+  "num_train_timesteps": 1000,
+  "shift": 1.0
+}

SiT-XL-2-512-diffusers/scheduler/scheduling_flow_match_sit.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class SiTFlowMatchSchedulerOutput(BaseOutput):
+    prev_sample: torch.Tensor
+
+
+class SiTFlowMatchScheduler(SchedulerMixin, ConfigMixin):
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        mode: str = "ode",
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        diffusion_form: str = "sigma",
+        diffusion_norm: float = 1.0,
+    ):
+        self.timesteps = None
+        self.sigmas = None
+        self._step_index = None
+
+    def set_timesteps(self, num_inference_steps: int, device: Optional[Union[str, torch.device]] = None):
+        # Flow matching integrates from noise (t=0) to data (t=1).
+        ts = torch.linspace(0.0, 1.0, num_inference_steps + 1, device=device, dtype=torch.float32)
+        self.timesteps = ts[:-1]
+        self.sigmas = 1.0 - self.timesteps
+        self._step_index = 0
+        return self.timesteps
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return sample
+
+    def _diffusion(self, t: torch.Tensor) -> torch.Tensor:
+        form = self.config.diffusion_form
+        norm = self.config.diffusion_norm
+        if form == "constant":
+            return torch.full_like(t, norm)
+        if form == "sigma":
+            return norm * (1.0 - t)
+        if form == "linear":
+            return norm * (1.0 - t)
+        if form == "decreasing":
+            return 0.25 * (norm * torch.cos(torch.pi * t) + 1) ** 2
+        if form == "increasing-decreasing":
+            return norm * torch.sin(torch.pi * t) ** 2
+        # "SBDM" approximated with sigma-based schedule for compatibility.
+        return norm * (1.0 - t)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: Union[float, torch.Tensor],
+        sample: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        return_dict: bool = True,
+    ) -> Union[SiTFlowMatchSchedulerOutput, Tuple[torch.Tensor]]:
+        if self.timesteps is None:
+            raise ValueError("Call `set_timesteps` before `step`.")
+        if self._step_index is None:
+            self._step_index = 0
+
+        step_index = min(self._step_index, len(self.timesteps) - 1)
+        t = self.timesteps[step_index].to(sample.device)
+        next_t = 1.0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1].to(sample.device)
+        dt = next_t - t
+
+        prev_sample = sample + model_output * dt
+        if self.config.mode.lower() == "sde":
+            diffusion = self._diffusion(torch.full((sample.shape[0],), t, device=sample.device, dtype=sample.dtype))
+            while diffusion.dim() < sample.dim():
+                diffusion = diffusion.unsqueeze(-1)
+            noise = torch.randn(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = prev_sample + torch.sqrt(torch.clamp(2.0 * diffusion * torch.abs(dt), min=0.0)) * noise
+
+        self._step_index += 1
+        if not return_dict:
+            return (prev_sample,)
+        return SiTFlowMatchSchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.Tensor,
+    ) -> torch.Tensor:
+        sigma = (1.0 - timesteps).view(-1, *([1] * (original_samples.ndim - 1)))
+        return (1 - sigma) * original_samples + sigma * noise

SiT-XL-2-512-diffusers/transformer/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "SiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "class_dropout_prob": 0.1,
+  "depth": 28,
+  "hidden_size": 1152,
+  "in_channels": 4,
+  "input_size": 64,
+  "learn_sigma": false,
+  "mlp_ratio": 4.0,
+  "num_classes": 1000,
+  "num_heads": 16,
+  "patch_size": 2
+}

SiT-XL-2-512-diffusers/transformer/diffusion_pytorch_model.safetensors

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

SiT-XL-2-512-diffusers/transformer/transformer_sit.py

@@ -0,0 +1,224 @@
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import Attention, Mlp, PatchEmbed
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+@dataclass
+class SiTTransformer2DModelOutput(BaseOutput):
+    sample: torch.Tensor
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 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: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
+        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: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size))
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes: int, hidden_size: int, dropout_prob: float):
+        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: torch.Tensor, force_drop_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+        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: torch.Tensor,
+        train: bool,
+        force_drop_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        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)
+        return self.embedding_table(labels)
+
+
+class SiTBlock(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        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):
+    def __init__(self, hidden_size: int, patch_size: int, out_channels: int):
+        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: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        return self.linear(x)
+
+
+class SiTTransformer2DModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        input_size: int = 32,
+        patch_size: int = 2,
+        in_channels: int = 4,
+        hidden_size: int = 1152,
+        depth: int = 28,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        class_dropout_prob: float = 0.1,
+        num_classes: int = 1000,
+        learn_sigma: bool = 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_classes = num_classes
+
+        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
+        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) -> None:
+        def _basic_init(module: nn.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)
+        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))
+
+        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)
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+        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: torch.Tensor) -> torch.Tensor:
+        c = self.out_channels
+        p = self.x_embedder.patch_size[0]
+        h = w = int(x.shape[1] ** 0.5)
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        force_drop_ids: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> SiTTransformer2DModelOutput:
+        x = self.x_embedder(hidden_states) + self.pos_embed
+        t = self.t_embedder(timestep)
+        y = self.y_embedder(class_labels, self.training, force_drop_ids=force_drop_ids)
+        c = t + y
+        for block in self.blocks:
+            x = block(x, c)
+        x = self.final_layer(x, c)
+        x = self.unpatchify(x)
+        if self.learn_sigma:
+            x, _ = x.chunk(2, dim=1)
+        if not return_dict:
+            return (x,)
+        return SiTTransformer2DModelOutput(sample=x)
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False, extra_tokens: int = 0):
+    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)
+    grid = np.stack(grid, axis=0).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: int, grid: np.ndarray):
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    return np.concatenate([emb_h, emb_w], axis=1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)

SiT-XL-2-512-diffusers/vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "stabilityai/sd-vae-ft-mse",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}