Delete SiT-XL-2-256-diffusers

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

@@ -1,42 +0,0 @@
----
-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

@@ -1,19 +0,0 @@
-{
-  "_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

@@ -1,82 +0,0 @@
-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

@@ -1,9 +0,0 @@
-{
-  "_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

@@ -1,98 +0,0 @@
-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

@@ -1,14 +0,0 @@
-{
-  "_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

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

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

@@ -1,224 +0,0 @@
-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

@@ -1,38 +0,0 @@
-{
-  "_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

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