Delete jit_diffusers

jit_diffusers/__init__.py

@@ -1,11 +0,0 @@
-from .modeling_jit_transformer_2d import JiTTransformer2DModel, JiTDiffusersModel
-from .pipeline_jit import JiTPipeline, JiTPipelineOutput
-from .scheduling_jit import JiTScheduler
-
-__all__ = [
-    "JiTTransformer2DModel",
-    "JiTDiffusersModel",
-    "JiTPipeline",
-    "JiTPipelineOutput",
-    "JiTScheduler",
-]

jit_diffusers/modeling_jit_backbone.py

@@ -1,321 +0,0 @@
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .modeling_jit_utils import VisionRotaryEmbeddingFast, get_2d_sincos_pos_embed, RMSNorm
-
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-class BottleneckPatchEmbed(nn.Module):
-    def __init__(self, img_size=224, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True):
-        super().__init__()
-        img_size = (img_size, img_size)
-        patch_size = (patch_size, patch_size)
-        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
-        self.img_size = img_size
-        self.patch_size = patch_size
-        self.num_patches = num_patches
-
-        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
-        self.proj2 = nn.Conv2d(pca_dim, embed_dim, kernel_size=1, stride=1, bias=bias)
-
-    def forward(self, x):
-        _, _, height, width = x.shape
-        assert height == self.img_size[0] and width == self.img_size[1], (
-            f"Input image size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
-        )
-        x = self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
-        return x
-
-
-class TimestepEmbedder(nn.Module):
-    def __init__(self, hidden_size, frequency_embedding_size=256):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=True),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        half = dim // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
-            device=t.device
-        )
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_freq = t_freq.to(dtype=self.mlp[0].weight.dtype)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class LabelEmbedder(nn.Module):
-    def __init__(self, num_classes, hidden_size):
-        super().__init__()
-        self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)
-        self.num_classes = num_classes
-
-    def forward(self, labels):
-        embeddings = self.embedding_table(labels)
-        return embeddings
-
-
-def scaled_dot_product_attention(query, key, value, dropout_p=0.0) -> torch.Tensor:
-    query_len, key_len = query.size(-2), key.size(-2)
-    scale_factor = 1 / math.sqrt(query.size(-1))
-    attn_bias = torch.zeros(query.size(0), 1, query_len, key_len, dtype=query.dtype, device=query.device)
-
-    with torch.amp.autocast("cuda", enabled=False):
-        attn_weight = query.float() @ key.float().transpose(-2, -1) * scale_factor
-    attn_weight += attn_bias
-    attn_weight = torch.softmax(attn_weight, dim=-1)
-    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
-    out = attn_weight @ value.float()
-    return out.to(query.dtype)
-
-
-class Attention(nn.Module):
-    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_norm=True, attn_drop=0.0, proj_drop=0.0):
-        super().__init__()
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-
-        self.q_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
-        self.k_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.attn_drop = nn.Dropout(attn_drop)
-        self.proj = nn.Linear(dim, dim)
-        self.proj_drop = nn.Dropout(proj_drop)
-
-    def forward(self, x, rope):
-        batch_size, num_tokens, channels = x.shape
-        qkv = self.qkv(x).reshape(batch_size, num_tokens, 3, self.num_heads, channels // self.num_heads).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]
-
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-        q = rope(q)
-        k = rope(k)
-
-        x = scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0)
-        x = x.transpose(1, 2).reshape(batch_size, num_tokens, channels)
-        x = x.to(self.proj.weight.dtype)
-        x = self.proj(x)
-        x = self.proj_drop(x)
-        return x
-
-
-class SwiGLUFFN(nn.Module):
-    def __init__(self, dim: int, hidden_dim: int, drop=0.0, bias=True) -> None:
-        super().__init__()
-        hidden_dim = int(hidden_dim * 2 / 3)
-        self.w12 = nn.Linear(dim, 2 * hidden_dim, bias=bias)
-        self.w3 = nn.Linear(hidden_dim, dim, bias=bias)
-        self.ffn_dropout = nn.Dropout(drop)
-
-    def forward(self, x):
-        x12 = self.w12(x)
-        x1, x2 = x12.chunk(2, dim=-1)
-        hidden = F.silu(x1) * x2
-        return self.w3(self.ffn_dropout(hidden))
-
-
-class FinalLayer(nn.Module):
-    def __init__(self, hidden_size, patch_size, out_channels):
-        super().__init__()
-        self.norm_final = RMSNorm(hidden_size)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
-
-    def forward(self, x, c):
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-
-class JiTBlock(nn.Module):
-    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, attn_drop=0.0, proj_drop=0.0):
-        super().__init__()
-        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
-        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, attn_drop=attn_drop, proj_drop=proj_drop)
-        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
-        mlp_hidden_dim = int(hidden_size * mlp_ratio)
-        self.mlp = SwiGLUFFN(hidden_size, mlp_hidden_dim, drop=proj_drop)
-        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
-
-    def forward(self, x, c, feat_rope=None):
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
-        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), rope=feat_rope)
-        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
-        return x
-
-
-class JiT(nn.Module):
-    def __init__(
-        self,
-        input_size=256,
-        patch_size=16,
-        in_channels=3,
-        hidden_size=1024,
-        depth=24,
-        num_heads=16,
-        mlp_ratio=4.0,
-        attn_drop=0.0,
-        proj_drop=0.0,
-        num_classes=1000,
-        bottleneck_dim=128,
-        in_context_len=32,
-        in_context_start=8,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.patch_size = patch_size
-        self.num_heads = num_heads
-        self.hidden_size = hidden_size
-        self.input_size = input_size
-        self.in_context_len = in_context_len
-        self.in_context_start = in_context_start
-        self.num_classes = num_classes
-
-        self.t_embedder = TimestepEmbedder(hidden_size)
-        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
-        self.x_embedder = BottleneckPatchEmbed(input_size, patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
-
-        num_patches = self.x_embedder.num_patches
-        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
-
-        if self.in_context_len > 0:
-            self.in_context_posemb = nn.Parameter(torch.zeros(1, self.in_context_len, hidden_size), requires_grad=True)
-            torch.nn.init.normal_(self.in_context_posemb, std=0.02)
-
-        half_head_dim = hidden_size // num_heads // 2
-        hw_seq_len = input_size // patch_size
-        self.feat_rope = VisionRotaryEmbeddingFast(dim=half_head_dim, pt_seq_len=hw_seq_len, num_cls_token=0)
-        self.feat_rope_incontext = VisionRotaryEmbeddingFast(dim=half_head_dim, pt_seq_len=hw_seq_len, num_cls_token=self.in_context_len)
-
-        self.blocks = nn.ModuleList(
-            [
-                JiTBlock(
-                    hidden_size,
-                    num_heads,
-                    mlp_ratio=mlp_ratio,
-                    attn_drop=attn_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
-                    proj_drop=proj_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
-                )
-                for i in range(depth)
-            ]
-        )
-
-        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
-        self.initialize_weights()
-
-    def initialize_weights(self):
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-
-        self.apply(_basic_init)
-
-        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))
-
-        w1 = self.x_embedder.proj1.weight.data
-        nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
-        w2 = self.x_embedder.proj2.weight.data
-        nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
-        nn.init.constant_(self.x_embedder.proj2.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, patch_size):
-        channels = self.out_channels
-        height = width = int(x.shape[1] ** 0.5)
-        assert height * width == x.shape[1]
-
-        x = x.reshape(shape=(x.shape[0], height, width, patch_size, patch_size, channels))
-        x = torch.einsum("nhwpqc->nchpwq", x)
-        images = x.reshape(shape=(x.shape[0], channels, height * patch_size, height * patch_size))
-        return images
-
-    def forward(self, x, t, y):
-        t_emb = self.t_embedder(t)
-        y_emb = self.y_embedder(y)
-        c = t_emb + y_emb
-
-        x = self.x_embedder(x)
-        x += self.pos_embed
-
-        for i, block in enumerate(self.blocks):
-            if self.in_context_len > 0 and i == self.in_context_start:
-                in_context_tokens = y_emb.unsqueeze(1).repeat(1, self.in_context_len, 1)
-                in_context_tokens += self.in_context_posemb
-                x = torch.cat([in_context_tokens, x], dim=1)
-            x = block(x, c, self.feat_rope if i < self.in_context_start else self.feat_rope_incontext)
-
-        x = x[:, self.in_context_len :]
-        x = self.final_layer(x, c)
-        output = self.unpatchify(x, self.patch_size)
-        return output
-
-
-def JiT_B_16(**kwargs):
-    return JiT(depth=12, hidden_size=768, num_heads=12, bottleneck_dim=128, in_context_len=32, in_context_start=4, patch_size=16, **kwargs)
-
-
-def JiT_B_32(**kwargs):
-    return JiT(depth=12, hidden_size=768, num_heads=12, bottleneck_dim=128, in_context_len=32, in_context_start=4, patch_size=32, **kwargs)
-
-
-def JiT_L_16(**kwargs):
-    return JiT(depth=24, hidden_size=1024, num_heads=16, bottleneck_dim=128, in_context_len=32, in_context_start=8, patch_size=16, **kwargs)
-
-
-def JiT_L_32(**kwargs):
-    return JiT(depth=24, hidden_size=1024, num_heads=16, bottleneck_dim=128, in_context_len=32, in_context_start=8, patch_size=32, **kwargs)
-
-
-def JiT_H_16(**kwargs):
-    return JiT(depth=32, hidden_size=1280, num_heads=16, bottleneck_dim=256, in_context_len=32, in_context_start=10, patch_size=16, **kwargs)
-
-
-def JiT_H_32(**kwargs):
-    return JiT(depth=32, hidden_size=1280, num_heads=16, bottleneck_dim=256, in_context_len=32, in_context_start=10, patch_size=32, **kwargs)
-
-
-JiT_models = {
-    "JiT-B/16": JiT_B_16,
-    "JiT-B/32": JiT_B_32,
-    "JiT-L/16": JiT_L_16,
-    "JiT-L/32": JiT_L_32,
-    "JiT-H/16": JiT_H_16,
-    "JiT-H/32": JiT_H_32,
-}

jit_diffusers/modeling_jit_transformer_2d.py

@@ -1,200 +0,0 @@
-from __future__ import annotations
-
-import argparse
-from collections.abc import Mapping
-from dataclasses import dataclass
-from typing import Any, Dict, Literal, Tuple
-
-import torch
-from diffusers import ConfigMixin, ModelMixin
-from diffusers.configuration_utils import register_to_config
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-
-from .modeling_jit_backbone import JiT_models
-
-
-def _extract_module_state_dict(
-    state_dict: Dict[str, torch.Tensor], prefixes: Tuple[str, ...] = ("transformer.", "net.")
-) -> Dict[str, torch.Tensor]:
-    """Extract module state by stripping the first fully-matching prefix.
-
-    Prefix precedence is left-to-right; `"transformer."` is preferred over legacy `"net."`.
-    """
-    for prefix in prefixes:
-        if all(key.startswith(prefix) for key in state_dict.keys()):
-            return {k[len(prefix):]: v for k, v in state_dict.items()}
-    return state_dict
-
-
-def _build_jit_kwargs(
-    image_size: int,
-    num_classes: int,
-    attn_dropout: float,
-    proj_dropout: float,
-    model_name: str | None = None,
-) -> Dict[str, object]:
-    # Keep model_name for backward-compatible internal call signatures.
-    _ = model_name
-    return {
-        "input_size": image_size,
-        "in_channels": 3,
-        "num_classes": num_classes,
-        "attn_drop": attn_dropout,
-        "proj_drop": proj_dropout,
-    }
-
-
-@dataclass
-class JiTCheckpointConfig:
-    model_name: str
-    image_size: int
-    num_classes: int
-    attn_dropout: float
-    proj_dropout: float
-
-
-def _config_from_checkpoint(ckpt_args: argparse.Namespace | Mapping[str, Any]) -> JiTCheckpointConfig:
-    if isinstance(ckpt_args, argparse.Namespace):
-        args_dict = vars(ckpt_args)
-    elif isinstance(ckpt_args, Mapping):
-        args_dict = ckpt_args
-    else:
-        raise TypeError(f"Unsupported checkpoint args type: {type(ckpt_args)}")
-
-    def _get_first_available(*keys: str, default=None):
-        for key in keys:
-            if key in args_dict and args_dict[key] is not None:
-                return args_dict[key]
-        return default
-
-    model_name = _get_first_available("model", "model_name", "model_type")
-    image_size = _get_first_available("img_size", "image_size", "sample_size")
-    num_classes = _get_first_available("class_num", "num_classes", "num_class_embeds")
-    if model_name is None or image_size is None or num_classes is None:
-        raise ValueError("Checkpoint args are missing model/image_size/num_classes information.")
-
-    return JiTCheckpointConfig(
-        model_name=str(model_name),
-        image_size=int(image_size),
-        num_classes=int(num_classes),
-        attn_dropout=float(_get_first_available("attn_dropout", "attention_dropout", default=0.0)),
-        proj_dropout=float(_get_first_available("proj_dropout", "dropout", default=0.0)),
-    )
-
-
-class JiTTransformer2DModel(ModelMixin, ConfigMixin):
-    @register_to_config
-    def __init__(
-        self,
-        model_type: str = "JiT-B/16",
-        sample_size: int = 256,
-        num_class_embeds: int = 1000,
-        attention_dropout: float = 0.0,
-        dropout: float = 0.0,
-        model_name: str | None = None,
-        image_size: int | None = None,
-        num_classes: int | None = None,
-        attn_dropout: float | None = None,
-        proj_dropout: float | None = None,
-    ):
-        super().__init__()
-        resolved_model_type = model_type if model_name is None else model_name
-        resolved_sample_size = sample_size if image_size is None else image_size
-        resolved_num_class_embeds = num_class_embeds if num_classes is None else num_classes
-        resolved_attention_dropout = attention_dropout if attn_dropout is None else attn_dropout
-        resolved_dropout = dropout if proj_dropout is None else proj_dropout
-
-        if resolved_model_type not in JiT_models:
-            raise ValueError(f"Unknown model '{resolved_model_type}'. Available: {list(JiT_models.keys())}")
-
-        self.transformer = JiT_models[resolved_model_type](
-            **_build_jit_kwargs(
-                image_size=resolved_sample_size,
-                num_classes=resolved_num_class_embeds,
-                attn_dropout=resolved_attention_dropout,
-                proj_dropout=resolved_dropout,
-                model_name=resolved_model_type,
-            )
-        )
-
-    def forward(
-        self,
-        sample: torch.Tensor,
-        timestep: torch.Tensor,
-        class_labels: torch.Tensor,
-        return_dict: bool = True,
-    ):
-        timestep = torch.as_tensor(timestep, device=sample.device)
-        if timestep.ndim == 0:
-            timestep = timestep.repeat(sample.shape[0])
-        else:
-            timestep = timestep.reshape(-1)
-            if timestep.shape[0] == 1 and sample.shape[0] > 1:
-                timestep = timestep.repeat(sample.shape[0])
-
-        denoised = self.transformer(sample, timestep, class_labels)
-        if not return_dict:
-            return (denoised,)
-        return Transformer2DModelOutput(sample=denoised)
-
-    @classmethod
-    def from_jit_checkpoint(
-        cls,
-        checkpoint_path: str,
-        weights: Literal["model", "ema1", "ema2"] = "ema1",
-        map_location: str = "cpu",
-        strict: bool = True,
-    ) -> Tuple["JiTTransformer2DModel", Dict[str, object]]:
-        checkpoint = torch.load(checkpoint_path, map_location=map_location)
-        if "args" not in checkpoint:
-            raise ValueError("Checkpoint is missing 'args', cannot infer JiT architecture config.")
-
-        config = _config_from_checkpoint(checkpoint["args"])
-        model = cls(
-            model_type=config.model_name,
-            sample_size=config.image_size,
-            num_class_embeds=config.num_classes,
-            attention_dropout=config.attn_dropout,
-            dropout=config.proj_dropout,
-        )
-
-        key = "model" if weights == "model" else f"model_{weights}"
-        if key not in checkpoint:
-            raise ValueError(f"Checkpoint key '{key}' not found. Available keys: {list(checkpoint.keys())}")
-
-        model_state = _extract_module_state_dict(checkpoint[key])
-        model.transformer.load_state_dict(model_state, strict=strict)
-
-        metadata = {
-            "checkpoint_path": checkpoint_path,
-            "weights": weights,
-            "epoch": checkpoint.get("epoch"),
-            "source_args": checkpoint.get("args"),
-        }
-        return model, metadata
-
-    def to_jit_checkpoint(
-        self,
-        ema_mode: Literal["none", "copy_to_both"] = "copy_to_both",
-        prefix: str = "net.",
-    ) -> Dict[str, object]:
-        base_state = {f"{prefix}{k}": v.detach().cpu() for k, v in self.transformer.state_dict().items()}
-        checkpoint = {"model": base_state}
-        if ema_mode == "copy_to_both":
-            checkpoint["model_ema1"] = {k: v.clone() for k, v in base_state.items()}
-            checkpoint["model_ema2"] = {k: v.clone() for k, v in base_state.items()}
-        elif ema_mode != "none":
-            raise ValueError(f"Unsupported ema_mode='{ema_mode}'.")
-        return checkpoint
-
-    @property
-    def net(self):
-        return self.transformer
-
-    @net.setter
-    def net(self, module):
-        self.transformer = module
-
-
-# Backward-compatible alias.
-JiTDiffusersModel = JiTTransformer2DModel

jit_diffusers/modeling_jit_utils.py

@@ -1,129 +0,0 @@
-from math import pi
-
-import numpy as np
-import torch
-from einops import rearrange, repeat
-from torch import nn
-
-
-def broadcat(tensors, dim=-1):
-    num_tensors = len(tensors)
-    shape_lens = set(list(map(lambda tensor: len(tensor.shape), tensors)))
-    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
-    shape_len = list(shape_lens)[0]
-    dim = (dim + shape_len) if dim < 0 else dim
-    dims = list(zip(*map(lambda tensor: list(tensor.shape), tensors)))
-    expandable_dims = [(index, val) for index, val in enumerate(dims) if index != dim]
-    assert all([*map(lambda tensor: len(set(tensor[1])) <= 2, expandable_dims)]), "invalid dimensions for broadcastable concatenation"
-    max_dims = list(map(lambda tensor: (tensor[0], max(tensor[1])), expandable_dims))
-    expanded_dims = list(map(lambda tensor: (tensor[0], (tensor[1],) * num_tensors), max_dims))
-    expanded_dims.insert(dim, (dim, dims[dim]))
-    expandable_shapes = list(zip(*map(lambda tensor: tensor[1], expanded_dims)))
-    tensors = list(map(lambda tensor: tensor[0].expand(*tensor[1]), zip(tensors, expandable_shapes)))
-    return torch.cat(tensors, dim=dim)
-
-
-def rotate_half(x):
-    x = rearrange(x, "... (d r) -> ... d r", r=2)
-    x1, x2 = x.unbind(dim=-1)
-    x = torch.stack((-x2, x1), dim=-1)
-    return rearrange(x, "... d r -> ... (d r)")
-
-
-class VisionRotaryEmbeddingFast(nn.Module):
-    def __init__(
-        self,
-        dim,
-        pt_seq_len=16,
-        ft_seq_len=None,
-        custom_freqs=None,
-        freqs_for="lang",
-        theta=10000,
-        max_freq=10,
-        num_freqs=1,
-        num_cls_token=0,
-    ):
-        super().__init__()
-        if custom_freqs:
-            freqs = custom_freqs
-        elif freqs_for == "lang":
-            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
-        elif freqs_for == "pixel":
-            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
-        elif freqs_for == "constant":
-            freqs = torch.ones(num_freqs).float()
-        else:
-            raise ValueError(f"unknown modality {freqs_for}")
-
-        if ft_seq_len is None:
-            ft_seq_len = pt_seq_len
-        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
-
-        freqs = torch.einsum("..., f -> ... f", t, freqs)
-        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
-        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
-
-        if num_cls_token > 0:
-            freqs_flat = freqs.view(-1, freqs.shape[-1])
-            cos_img = freqs_flat.cos()
-            sin_img = freqs_flat.sin()
-            _, dim_freq = cos_img.shape
-            cos_pad = torch.ones(num_cls_token, dim_freq, dtype=cos_img.dtype, device=cos_img.device)
-            sin_pad = torch.zeros(num_cls_token, dim_freq, dtype=sin_img.dtype, device=sin_img.device)
-            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
-            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
-        else:
-            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
-            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
-
-    def forward(self, tensor):
-        freqs_cos = self.freqs_cos.to(device=tensor.device, dtype=tensor.dtype)
-        freqs_sin = self.freqs_sin.to(device=tensor.device, dtype=tensor.dtype)
-        return tensor * freqs_cos + rotate_half(tensor) * freqs_sin
-
-
-class RMSNorm(nn.Module):
-    def __init__(self, hidden_size, eps=1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        variance = hidden_states.pow(2).mean(-1, keepdim=True)
-        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        return (self.weight * hidden_states).to(input_dtype)
-
-
-def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=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)
-    grid = grid.reshape([2, 1, grid_size, grid_size])
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    assert embed_dim % 2 == 0
-    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])
-    emb = np.concatenate([emb_h, emb_w], axis=1)
-    return emb
-
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    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)
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)
-    return emb

jit_diffusers/pipeline_jit.py

@@ -1,179 +0,0 @@
-from dataclasses import dataclass
-from pathlib import Path
-from typing import List, Tuple
-
-import numpy as np
-import torch
-from diffusers import DiffusionPipeline
-from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
-from diffusers.utils import BaseOutput
-from diffusers.utils.torch_utils import randn_tensor
-
-from .modeling_jit_transformer_2d import JiTTransformer2DModel
-from .scheduling_jit import JiTScheduler
-
-
-RECOMMENDED_CFG_BY_MODEL = {
-    "JiT-B/16": 3.0,
-    "JiT-L/16": 2.4,
-    "JiT-H/16": 2.2,
-    "JiT-B/32": 3.0,
-    "JiT-L/32": 2.5,
-    "JiT-H/32": 2.3,
-}
-
-RECOMMENDED_NOISE_BY_RESOLUTION = {
-    256: 1.0,
-    512: 2.0,
-}
-
-
-@dataclass
-class JiTPipelineOutput(BaseOutput):
-    images: List["PIL.Image.Image"] | np.ndarray | torch.Tensor
-
-
-class JiTPipeline(DiffusionPipeline):
-    model_cpu_offload_seq = "transformer"
-
-    def __init__(self, transformer: JiTTransformer2DModel, scheduler: JiTScheduler | None = None):
-        super().__init__()
-        self.register_modules(transformer=transformer, scheduler=scheduler or JiTScheduler())
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
-        model_kwargs = dict(kwargs)
-        transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
-        scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
-        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
-        if transformer_subfolder is not None:
-            transformer_path = str(Path(pretrained_model_name_or_path) / transformer_subfolder)
-        else:
-            transformer_path = pretrained_model_name_or_path
-        transformer = JiTTransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
-        try:
-            scheduler = JiTScheduler.from_pretrained(
-                pretrained_model_name_or_path,
-                subfolder=scheduler_subfolder,
-                **scheduler_kwargs,
-            )
-        except Exception:
-            scheduler = JiTScheduler(**scheduler_kwargs)
-        return cls(transformer=transformer, scheduler=scheduler)
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        class_labels: int | List[int] | torch.Tensor,
-        num_inference_steps: int = 50,
-        guidance_scale: float | None = None,
-        guidance_interval_min: float = 0.1,
-        guidance_interval_max: float = 1.0,
-        noise_scale: float | None = None,
-        t_eps: float = 5e-2,
-        sampling_method: str | None = None,
-        generator: torch.Generator | List[torch.Generator] | None = None,
-        output_type: str = "pil",
-        return_dict: bool = True,
-    ) -> JiTPipelineOutput | ImagePipelineOutput | Tuple:
-        if output_type not in {"pil", "np", "pt"}:
-            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
-        if sampling_method is not None and sampling_method not in {"heun", "euler"}:
-            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
-        if num_inference_steps < 2:
-            raise ValueError("num_inference_steps must be >= 2.")
-        if sampling_method is not None and sampling_method != self.scheduler.config.solver:
-            self.scheduler = JiTScheduler.from_config(self.scheduler.config, solver=sampling_method)
-
-        if isinstance(class_labels, int):
-            class_labels = [class_labels]
-        if isinstance(class_labels, list):
-            class_labels = torch.tensor(class_labels, device=self._execution_device, dtype=torch.long)
-        else:
-            class_labels = class_labels.to(self._execution_device, dtype=torch.long).reshape(-1)
-
-        batch_size = class_labels.shape[0]
-        latent_size = int(self.transformer.config.sample_size)
-        latent_channels = int(getattr(self.transformer.config, "in_channels", 3))
-        num_classes = int(self.transformer.config.num_class_embeds)
-        model_type = str(getattr(self.transformer.config, "model_type", ""))
-
-        if guidance_scale is None:
-            guidance_scale = RECOMMENDED_CFG_BY_MODEL.get(model_type, 2.9)
-        if noise_scale is None:
-            noise_scale = RECOMMENDED_NOISE_BY_RESOLUTION.get(latent_size, 1.0)
-
-        class_labels = class_labels.clamp(0, num_classes - 1)
-        class_null = torch.full_like(class_labels, num_classes)
-
-        latents = randn_tensor(
-            shape=(batch_size, latent_channels, latent_size, latent_size),
-            generator=generator,
-            device=self._execution_device,
-            dtype=self.transformer.dtype,
-        ) * noise_scale
-        self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self._execution_device)
-        timesteps = self.scheduler.timesteps.to(device=self._execution_device, dtype=latents.dtype)
-
-        def forward_cfg(z_value: torch.Tensor, t: torch.Tensor | float) -> torch.Tensor:
-            t = torch.as_tensor(t, device=self._execution_device, dtype=latents.dtype)
-            x_cond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_labels).sample
-            v_cond = (x_cond - z_value) / (1.0 - t).clamp_min(t_eps)
-
-            x_uncond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_null).sample
-            v_uncond = (x_uncond - z_value) / (1.0 - t).clamp_min(t_eps)
-
-            interval_mask = t < guidance_interval_max
-            if guidance_interval_min != 0.0:
-                interval_mask = interval_mask & (t > guidance_interval_min)
-            scale = torch.where(
-                interval_mask,
-                torch.tensor(guidance_scale, device=self._execution_device, dtype=latents.dtype),
-                torch.tensor(1.0, device=self._execution_device, dtype=latents.dtype),
-            )
-            return v_uncond + scale * (v_cond - v_uncond)
-
-        for i in self.progress_bar(range(num_inference_steps - 1)):
-            t, t_next = timesteps[i], timesteps[i + 1]
-            model_output = forward_cfg(latents, t)
-            if self.scheduler.config.solver == "heun":
-                latents = self.scheduler.step(
-                    model_output=model_output,
-                    timestep=t,
-                    next_timestep=t_next,
-                    sample=latents,
-                    model_fn=forward_cfg,
-                ).prev_sample
-            else:
-                latents = self.scheduler.step(
-                    model_output=model_output,
-                    timestep=t,
-                    next_timestep=t_next,
-                    sample=latents,
-                ).prev_sample
-
-        # Match the original JiT implementation: always use Euler for the final step.
-        t, t_next = timesteps[-2], timesteps[-1]
-        model_output = forward_cfg(latents, t)
-        latents = self.scheduler.euler_step(
-            model_output=model_output,
-            timestep=t,
-            next_timestep=t_next,
-            sample=latents,
-        ).prev_sample
-
-        images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
-        if output_type == "pt":
-            images = images_pt
-        else:
-            images_np = images_pt.permute(0, 2, 3, 1).numpy()
-            if output_type == "np":
-                images = images_np
-            else:
-                images = self.numpy_to_pil(images_np)
-
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (images,)
-        return JiTPipelineOutput(images=images)

jit_diffusers/scheduling_jit.py

@@ -1,71 +0,0 @@
-from __future__ import annotations
-
-from typing import Callable
-
-import torch
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.schedulers.scheduling_utils import SchedulerMixin, SchedulerOutput
-
-
-class JiTScheduler(SchedulerMixin, ConfigMixin):
-    order = 1
-
-    @register_to_config
-    def __init__(
-        self,
-        solver: str = "heun",
-        timestep_start: float = 0.0,
-        timestep_end: float = 1.0,
-    ):
-        if solver not in {"heun", "euler"}:
-            raise ValueError("solver must be one of: 'heun', 'euler'.")
-        if timestep_end <= timestep_start:
-            raise ValueError("timestep_end must be greater than timestep_start.")
-        self.timesteps = torch.tensor([])
-
-    def set_timesteps(self, num_inference_steps: int, device: str | torch.device | None = None):
-        if num_inference_steps < 2:
-            raise ValueError("num_inference_steps must be >= 2.")
-        self.timesteps = torch.linspace(
-            self.config.timestep_start,
-            self.config.timestep_end,
-            num_inference_steps + 1,
-            device=device,
-            dtype=torch.float32,
-        )
-
-    def euler_step(
-        self,
-        model_output: torch.Tensor,
-        timestep: torch.Tensor,
-        next_timestep: torch.Tensor,
-        sample: torch.Tensor,
-        return_dict: bool = True,
-    ) -> SchedulerOutput | tuple[torch.Tensor]:
-        prev_sample = sample + (next_timestep - timestep) * model_output
-        if not return_dict:
-            return (prev_sample,)
-        return SchedulerOutput(prev_sample=prev_sample)
-
-    def step(
-        self,
-        model_output: torch.Tensor,
-        timestep: torch.Tensor,
-        next_timestep: torch.Tensor,
-        sample: torch.Tensor,
-        model_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
-        return_dict: bool = True,
-    ) -> SchedulerOutput | tuple[torch.Tensor]:
-        if self.config.solver == "euler":
-            return self.euler_step(model_output, timestep, next_timestep, sample, return_dict=return_dict)
-
-        if model_fn is None:
-            raise ValueError("model_fn is required when solver='heun'.")
-
-        sample_euler = sample + (next_timestep - timestep) * model_output
-        model_output_next = model_fn(sample_euler, next_timestep)
-        prev_sample = sample + (next_timestep - timestep) * 0.5 * (model_output + model_output_next)
-
-        if not return_dict:
-            return (prev_sample,)
-        return SchedulerOutput(prev_sample=prev_sample)