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)