Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torch.utils.checkpoint import checkpoint | |
| from typing import Tuple | |
| from .pixeldit_c2i import PatchTokenEmbedder, PixelTokenEmbedder, PiTBlock | |
| from .modules import ( | |
| FinalLayer, | |
| FeedForward, | |
| PixelDiTJointAttnProcessor, | |
| RMSNorm, | |
| TimestepConditioner, | |
| apply_adaln, | |
| apply_rotary_emb, | |
| precompute_freqs_cis_2d, | |
| ) | |
| class MMDiTJointAttention(nn.Module): | |
| def __init__( | |
| self, | |
| dim: int, | |
| num_heads: int = 8, | |
| qkv_bias: bool = False, | |
| attn_drop: float = 0., | |
| proj_drop: float = 0., | |
| ) -> None: | |
| super().__init__() | |
| assert dim % num_heads == 0, 'dim should be divisible by num_heads' | |
| self.dim = dim | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.qkv_x = nn.Linear(dim, dim * 3, bias=qkv_bias) | |
| self.qkv_y = nn.Linear(dim, dim * 3, bias=qkv_bias) | |
| self.q_norm_x = RMSNorm(self.head_dim) | |
| self.k_norm_x = RMSNorm(self.head_dim) | |
| self.q_norm_y = RMSNorm(self.head_dim) | |
| self.k_norm_y = RMSNorm(self.head_dim) | |
| self.proj_x = nn.Linear(dim, dim) | |
| self.proj_y = nn.Linear(dim, dim) | |
| self.attn_drop = nn.Dropout(attn_drop) | |
| self.proj_drop_x = nn.Dropout(proj_drop) | |
| self.proj_drop_y = nn.Dropout(proj_drop) | |
| self.processor = PixelDiTJointAttnProcessor() | |
| def set_processor(self, processor) -> None: | |
| self.processor = processor | |
| def get_processor(self): | |
| return self.processor | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| y: torch.Tensor, | |
| pos_img: torch.Tensor, | |
| pos_txt: torch.Tensor = None, | |
| attn_mask: torch.Tensor = None, | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| B, Nx, C = x.shape | |
| By, Ny, Cy = y.shape | |
| assert B == By and C == Cy, "x and y must share batch and channel dims" | |
| qkv_x = self.qkv_x(x).reshape(B, Nx, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) | |
| qx, kx, vx = qkv_x[0], qkv_x[1], qkv_x[2] | |
| qx = self.q_norm_x(qx) | |
| kx = self.k_norm_x(kx) | |
| qkv_y = self.qkv_y(y).reshape(B, Ny, 3, self.num_heads, C // self.num_heads).permute(2, 0, 1, 3, 4) | |
| qy, ky, vy = qkv_y[0], qkv_y[1], qkv_y[2] | |
| qy = self.q_norm_y(qy) | |
| ky = self.k_norm_y(ky) | |
| qx, kx = apply_rotary_emb(qx, kx, freqs_cis=pos_img) | |
| if pos_txt is not None: | |
| qy, ky = apply_rotary_emb(qy, ky, freqs_cis=pos_txt) | |
| qx = qx.transpose(1, 2) | |
| kx = kx.transpose(1, 2) | |
| vx = vx.transpose(1, 2) | |
| qy = qy.transpose(1, 2) | |
| ky = ky.transpose(1, 2) | |
| vy = vy.transpose(1, 2) | |
| q_joint = torch.cat([qy, qx], dim=2) | |
| k_joint = torch.cat([ky, kx], dim=2) | |
| v_joint = torch.cat([vy, vx], dim=2) | |
| out_joint = self.processor(self, q_joint, k_joint, v_joint, attn_mask=attn_mask) | |
| out_y = out_joint[:, :, :Ny, :] | |
| out_x = out_joint[:, :, Ny:, :] | |
| out_y = out_y.transpose(1, 2).reshape(B, Ny, C) | |
| out_x = out_x.transpose(1, 2).reshape(B, Nx, C) | |
| out_x = self.proj_drop_x(self.proj_x(out_x)) | |
| out_y = self.proj_drop_y(self.proj_y(out_y)) | |
| return out_x, out_y | |
| class MMDiTBlockT2I(nn.Module): | |
| def __init__(self, hidden_size, groups, mlp_ratio=4.0, adaLN_modulation_img=None, adaLN_modulation_txt=None): | |
| super().__init__() | |
| self.hidden_size = hidden_size | |
| self.groups = groups | |
| self.head_dim = hidden_size // groups | |
| self.norm_x1 = RMSNorm(hidden_size, eps=1e-6) | |
| self.norm_y1 = RMSNorm(hidden_size, eps=1e-6) | |
| self.attn = MMDiTJointAttention(hidden_size, num_heads=groups, qkv_bias=False) | |
| self.norm_x2 = RMSNorm(hidden_size, eps=1e-6) | |
| self.norm_y2 = RMSNorm(hidden_size, eps=1e-6) | |
| mlp_hidden_dim = int(hidden_size * mlp_ratio) | |
| self.mlp_x = FeedForward(hidden_size, mlp_hidden_dim) | |
| self.mlp_y = FeedForward(hidden_size, mlp_hidden_dim) | |
| self.adaLN_modulation_img = adaLN_modulation_img if adaLN_modulation_img is not None else nn.Sequential(nn.Linear(hidden_size, 6 * hidden_size, bias=True)) | |
| self.adaLN_modulation_txt = adaLN_modulation_txt if adaLN_modulation_txt is not None else nn.Sequential(nn.Linear(hidden_size, 6 * hidden_size, bias=True)) | |
| def forward(self, x, y, c, pos_img, pos_txt=None, attn_mask=None): | |
| shift_msa_x, scale_msa_x, gate_msa_x, shift_mlp_x, scale_mlp_x, gate_mlp_x = self.adaLN_modulation_img(c).chunk(6, dim=-1) | |
| shift_msa_y, scale_msa_y, gate_msa_y, shift_mlp_y, scale_mlp_y, gate_mlp_y = self.adaLN_modulation_txt(c).chunk(6, dim=-1) | |
| x_norm = apply_adaln(self.norm_x1(x), shift_msa_x, scale_msa_x) | |
| y_norm = apply_adaln(self.norm_y1(y), shift_msa_y, scale_msa_y) | |
| attn_x, attn_y = self.attn(x_norm, y_norm, pos_img, pos_txt, attn_mask) | |
| x = x + gate_msa_x * attn_x | |
| y = y + gate_msa_y * attn_y | |
| x = x + gate_mlp_x * self.mlp_x(apply_adaln(self.norm_x2(x), shift_mlp_x, scale_mlp_x)) | |
| y = y + gate_mlp_y * self.mlp_y(apply_adaln(self.norm_y2(y), shift_mlp_y, scale_mlp_y)) | |
| return x, y | |
| class PixDiT_T2I(nn.Module): | |
| def __init__( | |
| self, | |
| in_channels=3, | |
| num_groups=16, | |
| hidden_size=1152, | |
| pixel_hidden_size=64, | |
| pixel_attn_hidden_size=None, | |
| pixel_num_groups=None, | |
| patch_depth=26, | |
| pixel_depth=2, | |
| num_text_blocks=4, | |
| patch_size=16, | |
| txt_embed_dim=4096, | |
| txt_max_length=1024, | |
| use_text_rope: bool = True, | |
| text_rope_theta: float = 10000.0, | |
| repa_encoder_index: int = -1, | |
| use_pixel_abs_pos: bool = True, | |
| ): | |
| super().__init__() | |
| self.in_channels = int(in_channels) | |
| self.out_channels = int(in_channels) | |
| self.hidden_size = int(hidden_size) | |
| self.num_groups = int(num_groups) | |
| self.patch_depth = int(patch_depth) | |
| self.pixel_depth = int(pixel_depth) | |
| self.num_text_blocks = int(num_text_blocks) | |
| self.patch_size = int(patch_size) | |
| self.pixel_hidden_size = int(pixel_hidden_size) | |
| self.txt_embed_dim = int(txt_embed_dim) | |
| self.txt_max_length = int(txt_max_length) | |
| self.use_text_rope = bool(use_text_rope) | |
| self.text_rope_theta = float(text_rope_theta) | |
| self.repa_encoder_index = int(repa_encoder_index) | |
| self.use_pixel_abs_pos = bool(use_pixel_abs_pos) | |
| if self.pixel_depth <= 0: | |
| raise ValueError("PixDiT_T2I expects pixel_depth > 0 to retain the pixel pathway") | |
| self.pixel_embedder = PixelTokenEmbedder(in_channels, self.pixel_hidden_size, use_pixel_abs_pos=self.use_pixel_abs_pos) | |
| self.s_embedder = PatchTokenEmbedder(in_channels * patch_size ** 2, hidden_size, bias=True) | |
| self.t_embedder = TimestepConditioner(hidden_size) | |
| self.y_embedder = PatchTokenEmbedder(self.txt_embed_dim, hidden_size, bias=True, norm_layer=RMSNorm) | |
| self.y_pos_embedding = nn.Parameter(torch.randn(1, self.txt_max_length, hidden_size)) | |
| self._shared_cond_adaln = None | |
| self._shared_cond_adaln_img = None | |
| self._shared_cond_adaln_txt = None | |
| self.patch_blocks = nn.ModuleList([ | |
| MMDiTBlockT2I( | |
| self.hidden_size, | |
| self.num_groups, | |
| adaLN_modulation_img=self._shared_cond_adaln_img, | |
| adaLN_modulation_txt=self._shared_cond_adaln_txt, | |
| ) | |
| for _ in range(self.patch_depth) | |
| ]) | |
| self.text_refine_blocks = None | |
| self.pixel_attn_hidden_size = ( | |
| int(pixel_attn_hidden_size) if pixel_attn_hidden_size is not None else self.hidden_size | |
| ) | |
| self.pixel_num_groups = int(pixel_num_groups) if pixel_num_groups is not None else self.num_groups | |
| self.pixel_blocks = nn.ModuleList( | |
| [ | |
| PiTBlock( | |
| self.pixel_hidden_size, | |
| self.hidden_size, | |
| patch_size=self.patch_size, | |
| num_heads=self.num_groups, | |
| mlp_ratio=4.0, | |
| attn_hidden_size=self.pixel_attn_hidden_size, | |
| attn_num_heads=self.pixel_num_groups, | |
| rope_fn=precompute_freqs_cis_2d, | |
| ) | |
| for _ in range(self.pixel_depth) | |
| ] | |
| ) | |
| self.final_layer = FinalLayer(self.pixel_hidden_size, self.out_channels) | |
| self.precompute_pos = dict() | |
| self.precompute_pos_txt = dict() | |
| self.last_repa_tokens = None | |
| self.gradient_checkpointing = False | |
| self.initialize_weights() | |
| def fetch_pos(self, height, width, device): | |
| if (height, width) in self.precompute_pos: | |
| return self.precompute_pos[(height, width)].to(device) | |
| else: | |
| pos = precompute_freqs_cis_2d(self.hidden_size // self.num_groups, height, width).to(device) | |
| self.precompute_pos[(height, width)] = pos | |
| return pos | |
| def fetch_pos_text(self, length, device): | |
| if length in self.precompute_pos_txt: | |
| return self.precompute_pos_txt[length].to(device) | |
| head_dim = self.hidden_size // self.num_groups | |
| freqs = 1.0 / (self.text_rope_theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) | |
| positions = torch.arange(0, length, device=device).float().unsqueeze(1) | |
| angles = positions * freqs.unsqueeze(0) | |
| freqs_cis = torch.polar(torch.ones_like(angles), angles) | |
| self.precompute_pos_txt[length] = freqs_cis | |
| return freqs_cis | |
| def initialize_weights(self): | |
| w = self.s_embedder.proj.weight.data | |
| nn.init.xavier_uniform_(w.view([w.shape[0], -1])) | |
| nn.init.constant_(self.s_embedder.proj.bias, 0) | |
| nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) | |
| nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) | |
| nn.init.zeros_(self.final_layer.linear.weight) | |
| nn.init.zeros_(self.final_layer.linear.bias) | |
| def forward(self, x, t, y, s=None, mask=None): | |
| B, _, H, W = x.shape | |
| Hs = H // self.patch_size | |
| Ws = W // self.patch_size | |
| L = Hs * Ws | |
| pos = self.fetch_pos(Hs, Ws, x.device) | |
| x_patches = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size).transpose(1, 2) | |
| t_emb = self.t_embedder(t.view(-1)).view(B, -1, self.hidden_size) | |
| if y.dim() != 3: | |
| raise ValueError("Text embedding y must be [B, L, D]") | |
| Ltxt = min(y.shape[1], self.txt_max_length) | |
| y = y[:, :Ltxt, :] | |
| y_emb = self.y_embedder(y).view(B, Ltxt, self.hidden_size) | |
| y_emb = y_emb + self.y_pos_embedding[:, :Ltxt, :].to(y_emb.dtype) | |
| condition = torch.nn.functional.silu(t_emb) | |
| if s is None: | |
| s0 = self.s_embedder(x_patches) | |
| pos_txt = self.fetch_pos_text(Ltxt, x.device) if self.use_text_rope else None | |
| attn_mask_joint = None | |
| if mask is not None and isinstance(mask, torch.Tensor): | |
| m = mask | |
| while m.dim() > 2 and m.size(1) == 1: | |
| m = m.squeeze(1) | |
| if m.dim() == 3 and m.size(1) == 1: | |
| m = m.squeeze(1) | |
| if m.dim() == 2: | |
| pad = (m == 0) | |
| pad_img = torch.zeros((B, L), dtype=torch.bool, device=x.device) | |
| attn_mask_joint = torch.cat([pad[:, :Ltxt], pad_img], dim=1).view(B, 1, 1, Ltxt + L) | |
| self.last_repa_tokens = None | |
| s = s0 | |
| for i in range(self.patch_depth): | |
| if self.gradient_checkpointing and self.training: | |
| s, y_emb = checkpoint( | |
| self.patch_blocks[i], | |
| s, y_emb, condition, pos, pos_txt, attn_mask_joint, | |
| use_reentrant=False, | |
| ) | |
| else: | |
| s, y_emb = self.patch_blocks[i](s, y_emb, condition, pos, pos_txt, attn_mask_joint) | |
| if 0 < self.repa_encoder_index == (i + 1): | |
| self.last_repa_tokens = s | |
| s = torch.nn.functional.silu(t_emb + s) | |
| if not (0 < self.repa_encoder_index <= self.patch_depth): | |
| self.last_repa_tokens = s | |
| batch_size, length, _ = s.shape | |
| if length != L: | |
| if length > L: | |
| s = s[:, :L, :] | |
| else: | |
| pad_len = L - length | |
| s = torch.cat([s, s.new_zeros(B, pad_len, s.shape[2])], dim=1) | |
| length = L | |
| s_cond = s.view(B * L, self.hidden_size) | |
| x_pixels = self.pixel_embedder(x, img_height=H, img_width=W, patch_size=self.patch_size) | |
| for blk in self.pixel_blocks: | |
| x_pixels = blk(x_pixels, s_cond, H, W, self.patch_size, mask) | |
| x_pixels = self.final_layer(x_pixels) | |
| C_out = self.out_channels | |
| P2 = self.patch_size * self.patch_size | |
| x_pixels = x_pixels.view(B, L, P2, C_out).permute(0, 3, 2, 1).contiguous() | |
| x_pixels = x_pixels.view(B, C_out * P2, L) | |
| x_img = torch.nn.functional.fold(x_pixels, (H, W), kernel_size=self.patch_size, stride=self.patch_size) | |
| return x_img | |