modeling_tinyimagegen.py

import math
from dataclasses import dataclass
from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import ModelOutput
from safetensors.torch import load_file
import os

try:
    from .configuration_tinyimagegen import TinyImageGenConfig
except Exception:
    from configuration_tinyimagegen import TinyImageGenConfig

@torch.no_grad()
def get_hadamard_matrix(d: int, dtype=torch.float32) -> torch.Tensor:
    p2 = 1 << (d - 1).bit_length()
    eye = torch.eye(p2, dtype=dtype)
    h = 1
    out = eye.clone()
    while h < p2:
        out = out.view(-1, 2, h)
        u = out[:, 0, :]
        v = out[:, 1, :]
        out = torch.cat((u + v, u - v), dim=-2)
        out = out.view(p2, p2)
        h *= 2
    out = out * (1.0 / math.sqrt(p2))
    return out[:d, :d].contiguous()

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x * norm * self.weight

class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size: int, frequency_embedding_size: int = 128):
        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: float = 10000.0) -> torch.Tensor:
        half = dim // 2
        freqs = torch.exp(
            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half
        )
        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:
        t_freq = self.timestep_embedding(t * 1000.0, self.frequency_embedding_size)
        return self.mlp(t_freq)

class RotaryEmbedding2D(nn.Module):
    def __init__(self, head_dim: int, base: float = 10000.0):
        super().__init__()
        self.head_dim = head_dim
        self.dim_h = 2 * (head_dim // 4)
        self.dim_w = head_dim - self.dim_h
        self.base = base

        inv_freq_h = 1.0 / (self.base ** (torch.arange(0, self.dim_h, 2, dtype=torch.float32) / self.dim_h))
        inv_freq_w = 1.0 / (self.base ** (torch.arange(0, self.dim_w, 2, dtype=torch.float32) / self.dim_w))
        self.register_buffer("inv_freq_h", inv_freq_h, persistent=False)
        self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)

    def forward(self, grid_h: int, grid_w: int, device: torch.device, dtype: torch.dtype = torch.float32):
        if self.inv_freq_h is None or self.inv_freq_h.device.type == "meta" or (self.inv_freq_h == 0).all():
            self.inv_freq_h = 1.0 / (self.base ** (torch.arange(0, self.dim_h, 2, dtype=torch.float32, device=device) / self.dim_h))
            self.inv_freq_w = 1.0 / (self.base ** (torch.arange(0, self.dim_w, 2, dtype=torch.float32, device=device) / self.dim_w))

        t_h = torch.arange(grid_h, device=device, dtype=torch.float32)
        t_w = torch.arange(grid_w, device=device, dtype=torch.float32)

        freqs_h = torch.outer(t_h, self.inv_freq_h.to(device=device, dtype=torch.float32))
        freqs_w = torch.outer(t_w, self.inv_freq_w.to(device=device, dtype=torch.float32))

        emb_h = torch.cat((freqs_h, freqs_h), dim=-1)
        emb_w = torch.cat((freqs_w, freqs_w), dim=-1)

        emb_h_grid = emb_h[:, None, :].expand(-1, grid_w, -1).reshape(grid_h * grid_w, self.dim_h)
        emb_w_grid = emb_w[None, :, :].expand(grid_h, -1, -1).reshape(grid_h * grid_w, self.dim_w)

        cos_h = emb_h_grid.cos().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
        sin_h = emb_h_grid.sin().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
        cos_w = emb_w_grid.cos().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
        sin_w = emb_w_grid.sin().to(dtype=dtype).unsqueeze(0).unsqueeze(0)
        return cos_h, sin_h, cos_w, sin_w

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb_2d(q: torch.Tensor, k: torch.Tensor, cos_h: torch.Tensor, sin_h: torch.Tensor, cos_w: torch.Tensor, sin_w: torch.Tensor):
    d_h = cos_h.shape[-1]
    qh, qw = q[..., :d_h], q[..., d_h:]
    kh, kw = k[..., :d_h], k[..., d_h:]

    qh_rot = (qh * cos_h) + (rotate_half(qh) * sin_h)
    qw_rot = (qw * cos_w) + (rotate_half(qw) * sin_w)
    kh_rot = (kh * cos_h) + (rotate_half(kh) * sin_h)
    kw_rot = (kw * cos_w) + (rotate_half(kw) * sin_w)

    return torch.cat([qh_rot, qw_rot], dim=-1), torch.cat([kh_rot, kw_rot], dim=-1)

class HadamardMLP(nn.Module):
    def __init__(self, config: TinyImageGenConfig):
        super().__init__()
        self.dim = config.hidden_size
        self.scale1 = nn.Parameter(torch.ones(self.dim))
        self.scale2 = nn.Parameter(torch.ones(self.dim))
        self.gate = nn.Parameter(torch.ones(self.dim))
        self.bias = nn.Parameter(torch.zeros(self.dim))

        hadamard_mat = get_hadamard_matrix(self.dim)
        self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not hasattr(self, "hadamard_mat") or self.hadamard_mat is None or self.hadamard_mat.device.type == "meta" or (self.hadamard_mat == 0).all() or self.hadamard_mat.abs().max() > 10.0:
            hadamard_mat = get_hadamard_matrix(self.dim, dtype=torch.float32).to(x.device)
            self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)
        mat = self.hadamard_mat.type_as(x)
        h = (x * self.scale1) @ mat
        g = F.silu(x * self.gate)
        out = ((h * g) @ mat) * self.scale2 + self.bias
        return out

class SwiGLUMLP(nn.Module):
    def __init__(self, config: TinyImageGenConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

class XSAGQAttention(nn.Module):
    def __init__(self, config: TinyImageGenConfig):
        super().__init__()
        self.dim = config.hidden_size
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.n_heads // self.n_kv_heads
        self.use_xsa = config.use_xsa
        self.use_per_head_gating = config.use_per_head_gating

        self.wq = nn.Linear(self.dim, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False)

        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

        if self.use_per_head_gating:
            self.head_gate = nn.Linear(self.dim, self.n_heads, bias=True)
            nn.init.constant_(self.head_gate.bias, 1.0)
            nn.init.zeros_(self.head_gate.weight)

    def forward(self, x: torch.Tensor, cos_h: torch.Tensor, sin_h: torch.Tensor, cos_w: torch.Tensor, sin_w: torch.Tensor) -> torch.Tensor:
        bsz, seqlen, _ = x.shape

        xq = self.wq(x).view(bsz, seqlen, self.n_heads, self.head_dim).transpose(1, 2)
        xk = self.wk(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2)
        xv = self.wv(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2)

        xq = self.q_norm(xq)
        xk = self.k_norm(xk)

        xq, xk = apply_rotary_pos_emb_2d(xq, xk, cos_h, sin_h, cos_w, sin_w)

        if self.num_kv_groups > 1:
            xk = xk.repeat_interleave(self.num_kv_groups, dim=1)
            xv_expanded = xv.repeat_interleave(self.num_kv_groups, dim=1)
        else:
            xv_expanded = xv

        attn_out = F.scaled_dot_product_attention(xq, xk, xv_expanded, is_causal=False)

        if self.use_xsa:
            vn = F.normalize(xv_expanded, p=2, dim=-1, eps=1e-6)
            proj = (attn_out * vn).sum(dim=-1, keepdim=True)
            attn_out = attn_out - proj * vn

        if self.use_per_head_gating:
            gate = torch.sigmoid(self.head_gate(x)).transpose(1, 2).unsqueeze(-1)
            attn_out = attn_out * gate

        out = attn_out.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(out)

def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class MultiLaneBlock(nn.Module):
    def __init__(self, config: TinyImageGenConfig, layer_idx: int):
        super().__init__()
        self.num_lanes = config.num_lanes
        self.dim = config.hidden_size
        self.layer_idx = layer_idx

        self.attn_norm = RMSNorm(self.dim, eps=config.rms_norm_eps)
        self.attn = XSAGQAttention(config)

        self.mlp_norm = RMSNorm(self.dim, eps=config.rms_norm_eps)
        if config.swiglu_interval == 0:
            self.use_swiglu = False
        elif config.swiglu_interval == 1:
            self.use_swiglu = True
        else:
            self.use_swiglu = ((layer_idx + 1) % config.swiglu_interval == 0)

        if self.use_swiglu:
            self.mlp = SwiGLUMLP(config)
        else:
            self.mlp = HadamardMLP(config)

        self.lane_mix_attn = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes))
        self.lane_mix_mlp = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes))

        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(config.hidden_size, 6 * config.hidden_size, bias=True)
        )
        nn.init.zeros_(self.adaLN_modulation[-1].weight)
        nn.init.zeros_(self.adaLN_modulation[-1].bias)

    def forward(self, lanes: torch.Tensor, t_emb: torch.Tensor, cos_h: torch.Tensor, sin_h: torch.Tensor, cos_w: torch.Tensor, sin_w: torch.Tensor) -> torch.Tensor:
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t_emb).chunk(6, dim=-1)

        primary = lanes[0]
        normed_primary = modulate(self.attn_norm(primary), shift_msa, scale_msa)
        attn_update = self.attn(normed_primary, cos_h, sin_h, cos_w, sin_w) * gate_msa.unsqueeze(1)

        mixed = torch.matmul(self.lane_mix_attn, lanes.view(self.num_lanes, -1)).view_as(lanes)
        lanes = torch.cat([(mixed[0] + attn_update).unsqueeze(0), mixed[1:]], dim=0)

        normed_primary = modulate(self.mlp_norm(lanes[0]), shift_mlp, scale_mlp)
        mlp_update = self.mlp(normed_primary) * gate_mlp.unsqueeze(1)

        mixed = torch.matmul(self.lane_mix_mlp, lanes.view(self.num_lanes, -1)).view_as(lanes)
        lanes = torch.cat([(mixed[0] + mlp_update).unsqueeze(0), mixed[1:]], dim=0)
        return lanes

@dataclass
class DiffusionOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    v_pred: Optional[torch.FloatTensor] = None

class TinyImageGenPreTrainedModel(PreTrainedModel):
    config_class = TinyImageGenConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MultiLaneBlock"]

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=std)
            if hasattr(module, "bias") and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, RMSNorm):
            module.weight.data.fill_(1.0)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        kwargs.pop("trust_remote_code", None)
        torch_dtype = kwargs.pop("torch_dtype", None)
        kwargs.pop("device_map", None)
        kwargs.pop("low_cpu_mem_usage", None)

        if config is None:
            config = TinyImageGenConfig.from_pretrained(pretrained_model_name_or_path)

        model = cls(config, *model_args)

        st_file = os.path.join(pretrained_model_name_or_path, "model.safetensors")
        bin_file = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")

        if os.path.exists(st_file):
            state_dict = load_file(st_file)
        elif os.path.exists(bin_file):
            state_dict = torch.load(bin_file, map_location="cpu")
        else:
            model = super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
            for module in model.modules():
                if type(module).__name__ == "RotaryEmbedding2D":
                    module.inv_freq_h = 1.0 / (module.base ** (torch.arange(0, module.dim_h, 2, dtype=torch.float32) / module.dim_h))
                    module.inv_freq_w = 1.0 / (module.base ** (torch.arange(0, module.dim_w, 2, dtype=torch.float32) / module.dim_w))
                elif type(module).__name__ == "HadamardMLP":
                    module.register_buffer("hadamard_mat", get_hadamard_matrix(module.dim, dtype=torch.float32), persistent=False)
            return model

        model_keys = set(model.state_dict().keys())
        st_keys = set(state_dict.keys())
        if not model_keys.intersection(st_keys):
            if any(k.startswith("model.") for k in st_keys):
                state_dict = {k[6:] if k.startswith("model.") else k: v for k, v in state_dict.items()}
            elif any(k.startswith("model.") for k in model_keys):
                state_dict = {f"model.{k}": v for k, v in state_dict.items()}

        model.load_state_dict(state_dict, strict=True)

        for module in model.modules():
            if type(module).__name__ == "RotaryEmbedding2D":
                module.inv_freq_h = 1.0 / (module.base ** (torch.arange(0, module.dim_h, 2, dtype=torch.float32) / module.dim_h))
                module.inv_freq_w = 1.0 / (module.base ** (torch.arange(0, module.dim_w, 2, dtype=torch.float32) / module.dim_w))
            elif type(module).__name__ == "HadamardMLP":
                module.register_buffer("hadamard_mat", get_hadamard_matrix(module.dim, dtype=torch.float32), persistent=False)

        if torch_dtype is not None:
            model.to(dtype=torch_dtype)

        return model

class TinyImageGenModel(TinyImageGenPreTrainedModel):
    def __init__(self, config: TinyImageGenConfig, *args, **kwargs):
        super().__init__(config)
        self.config = config
        self.num_lanes = config.num_lanes
        self.gradient_checkpointing = False

        self.x_embedder = nn.Linear(config.patch_dim, config.hidden_size, bias=True)
        self.t_embedder = TimestepEmbedder(config.hidden_size)
        self.rotary_emb = RotaryEmbedding2D(config.head_dim, base=config.rope_theta)

        self.layers = nn.ModuleList([
            MultiLaneBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)
        ])

        self.lane_pool_weights = nn.Parameter(torch.tensor([1.0] + [0.1] * (config.num_lanes - 1)))
        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.final_adaLN = nn.Sequential(
            nn.SiLU(),
            nn.Linear(config.hidden_size, 2 * config.hidden_size, bias=True)
        )
        self.final_proj = nn.Linear(config.hidden_size, config.patch_dim, bias=True)

        nn.init.zeros_(self.final_adaLN[-1].weight)
        nn.init.zeros_(self.final_adaLN[-1].bias)
        nn.init.zeros_(self.final_proj.weight)
        nn.init.zeros_(self.final_proj.bias)

        self.post_init()

    def patchify(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        p = self.config.patch_size
        h_patches, w_patches = H // p, W // p
        x = x.view(B, C, h_patches, p, w_patches, p)
        x = torch.einsum("bchpwq->bhwpcq", x)
        x = x.reshape(B, h_patches * w_patches, p * p * C)
        return x

    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
        B, N, _ = x.shape
        p = self.config.patch_size
        h_patches = self.config.num_patches_side
        w_patches = self.config.num_patches_side
        c = self.config.in_channels
        x = x.reshape(B, h_patches, w_patches, p, p, c)
        x = torch.einsum("bhwpqc->bchpwq", x)
        x = x.reshape(B, c, h_patches * p, w_patches * p)
        return x

    def forward(self, x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        bsz = x_t.shape[0]

        h0 = self.x_embedder(self.patchify(x_t))
        t_emb = self.t_embedder(t)

        lanes = h0.unsqueeze(0).repeat(self.num_lanes, 1, 1, 1)

        cos_h, sin_h, cos_w, sin_w = self.rotary_emb(self.config.num_patches_side, self.config.num_patches_side, device=x_t.device, dtype=x_t.dtype)

        for layer in self.layers:
            if self.gradient_checkpointing and self.training:
                lanes = cp.checkpoint(layer, lanes, t_emb, cos_h, sin_h, cos_w, sin_w, use_reentrant=False)
            else:
                lanes = layer(lanes, t_emb, cos_h, sin_h, cos_w, sin_w)

        pool_weights = F.softmax(self.lane_pool_weights, dim=0).view(self.num_lanes, 1, 1, 1)
        pooled = (lanes * pool_weights).sum(dim=0)

        shift, scale = self.final_adaLN(t_emb).chunk(2, dim=-1)
        out = modulate(self.final_norm(pooled), shift, scale)
        out = self.final_proj(out)
        return self.unpatchify(out)

    @torch.no_grad()
    def sample(self, num_samples: int, device: torch.device, num_steps: int = 25) -> torch.Tensor:
        was_training = self.training
        self.eval()
        x = torch.randn((num_samples, self.config.in_channels, self.config.image_size, self.config.image_size), device=device)
        dt = 1.0 / num_steps

        for step in range(num_steps):
            t_val = step / num_steps
            t = torch.full((num_samples,), t_val, device=device, dtype=torch.float32)
            v = self(x, t)
            x = x + v * dt

        if was_training:
            self.train()
        return x.clamp(-1.0, 1.0)

class TinyImageGenModelForImageDiffusion(TinyImageGenPreTrainedModel):
    def __init__(self, config: TinyImageGenConfig, *args, **kwargs):
        super().__init__(config)
        self.model = TinyImageGenModel(config)
        self.post_init()

    def forward(

        self,

        pixel_values: Optional[torch.Tensor] = None,

        x_t: Optional[torch.Tensor] = None,

        t: Optional[torch.Tensor] = None,

        return_dict: Optional[bool] = None,

    ) -> DiffusionOutput:
        return_dict = return_dict if return_dict is not None else getattr(self.config, "return_dict", True)

        loss = None
        v_pred = None

        if pixel_values is not None:
            x_1 = pixel_values
            bsz = x_1.shape[0]
            x_0 = torch.randn_like(x_1)
            t_rand = torch.rand(bsz, device=x_1.device)
            t_exp = t_rand.view(bsz, 1, 1, 1)

            x_t_flow = (1.0 - t_exp) * x_0 + t_exp * x_1
            v_target = x_1 - x_0

            v_pred = self.model(x_t_flow, t_rand)
            loss = F.mse_loss(v_pred, v_target)
        elif x_t is not None and t is not None:
            v_pred = self.model(x_t, t)
        else:
            raise ValueError("You must pass either 'pixel_values' for training or ('x_t', 't') for inference.")

        if not return_dict:
            return (loss, v_pred) if loss is not None else (v_pred,)

        return DiffusionOutput(loss=loss, v_pred=v_pred)

    @torch.no_grad()
    def sample(self, num_samples: int, device: torch.device, num_steps: int = 25) -> torch.Tensor:
        return self.model.sample(num_samples=num_samples, device=device, num_steps=num_steps)