Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.utils.checkpoint | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from diffusers.configuration_utils import ConfigMixin, register_to_config | |
| from diffusers.loaders import PeftAdapterMixin | |
| from diffusers.models.attention import FeedForward | |
| from diffusers.models.attention_processor import Attention, AttentionProcessor | |
| from diffusers.models.embeddings import ( | |
| GaussianFourierProjection, | |
| TimestepEmbedding, | |
| Timesteps, | |
| ) | |
| import numpy as np | |
| from diffusers.models.modeling_utils import ModelMixin | |
| from diffusers.models.normalization import ( | |
| AdaLayerNormContinuous, | |
| FP32LayerNorm, | |
| LayerNorm, | |
| ) | |
| from torchvision.ops import generalized_box_iou_loss | |
| from diffusers.utils import ( | |
| USE_PEFT_BACKEND, | |
| is_torch_version, | |
| logging, | |
| scale_lora_layers, | |
| unscale_lora_layers, | |
| ) | |
| from diffusers.utils.torch_utils import maybe_allow_in_graph | |
| from typing import * | |
| from ..attention_processor import FusedTripoSGAttnProcessor2_0, TripoSGAttnProcessor2_0, PartCrafterAttnProcessor | |
| from .modeling_outputs import Transformer1DModelOutput | |
| logger = logging.get_logger(__name__) # pylint: disable=invalid-name | |
| def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Helper function for AdaLN. | |
| x: [N, T, D] | |
| shift, scale: [N, D] | |
| """ | |
| # We unsqueeze(1) to make shift/scale broadcastable with x | |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) | |
| class PointPromptEncoder(nn.Module): | |
| """ | |
| Encodes 2D point prompts (x, y) into a high-dimensional embedding. | |
| Uses GaussianFourierProjection for positional encoding. | |
| """ | |
| def __init__(self, output_dim: int, fourier_dim: int = 128, fourier_scale: float = 4.0): | |
| super().__init__() | |
| self.output_dim = output_dim | |
| self.point_gff_encoder = GaussianFourierProjection(embedding_size=fourier_dim, scale=fourier_scale,log=False) | |
| prompt_input_dim = fourier_dim * 2 * 2 | |
| self.prompt_mlp = nn.Sequential( | |
| nn.Linear(prompt_input_dim, self.output_dim), | |
| nn.ReLU(), | |
| nn.Linear(self.output_dim, self.output_dim) | |
| ) | |
| def forward(self, point_prompt_xy: torch.Tensor) -> torch.Tensor: | |
| x_prompt = point_prompt_xy[:, 0:1] # [N, 1] | |
| y_prompt = point_prompt_xy[:, 1:2] # [N, 1] | |
| x_embed = self.point_gff_encoder(x_prompt.squeeze(-1)) # [N, fourier_dim * 2] | |
| y_embed = self.point_gff_encoder(y_prompt.squeeze(-1)) # [N, fourier_dim * 2] | |
| prompt_fourier_features = torch.cat([x_embed, y_embed], dim=-1) | |
| point_prompt_condition = self.prompt_mlp(prompt_fourier_features) | |
| return point_prompt_condition | |
| class BoxPromptEncoder(nn.Module): | |
| """ | |
| Encodes 2D bounding boxes (Box) into a high-dimensional embedding. | |
| (Inspired by the SAM architecture) | |
| 1. Decompose (x1, y1, x2, y2) into two corner points: top-left and bottom-right. | |
| 2. Encode these two (x, y) coordinates using Gaussian Fourier Features (GFF). | |
| 3. Add learnable "corner type" embeddings for "tl" (top-left) and "br" (bottom-right) respectively. | |
| 4. Pass both through an MLP. | |
| 5. Sum the final embeddings of the two corners to produce the box_embedding. | |
| """ | |
| def __init__(self, output_dim: int, fourier_dim: int = 128, fourier_scale: float = 4.0): | |
| super().__init__() | |
| self.output_dim = output_dim | |
| self.gff_encoder = GaussianFourierProjection( | |
| embedding_size=fourier_dim, scale=fourier_scale, log=False | |
| ) | |
| gff_output_dim = fourier_dim * 2 * 2 | |
| self.corner_type_embedding = nn.Embedding(2, gff_output_dim) | |
| self.corner_mlp = nn.Sequential( | |
| nn.Linear(gff_output_dim, self.output_dim), | |
| nn.ReLU(), | |
| nn.Linear(self.output_dim, self.output_dim) | |
| ) | |
| def _encode_corner(self, xy_coords: torch.Tensor, corner_type: int) -> torch.Tensor: | |
| x_c = xy_coords[:, 0:1] | |
| y_c = xy_coords[:, 1:2] | |
| x_embed = self.gff_encoder(x_c.squeeze(-1)) # [B, fourier_dim * 2] | |
| y_embed = self.gff_encoder(y_c.squeeze(-1)) # [B, fourier_dim * 2] | |
| features = torch.cat([x_embed, y_embed], dim=-1) | |
| type_emb = self.corner_type_embedding( | |
| torch.tensor(corner_type, device=features.device) | |
| ) | |
| return self.corner_mlp(features + type_emb) | |
| def forward(self, box_xyxy: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Args: | |
| box_xyxy (torch.Tensor): [B, 4] 归一化坐标 (x1, y1, x2, y2) | |
| """ | |
| # [B, 2] | |
| top_left = box_xyxy[:, :2] | |
| bottom_right = box_xyxy[:, 2:] | |
| # 编码两个角 | |
| emb_tl = self._encode_corner(top_left, corner_type=0) # [B, output_dim] | |
| emb_br = self._encode_corner(bottom_right, corner_type=1) # [B, output_dim] | |
| # 将两个嵌入相加，得到最终的 Box 嵌入 | |
| box_embedding = emb_tl + emb_br | |
| return box_embedding | |
| class MaskPromptEncoder(nn.Module): | |
| def __init__(self, output_dim=256): | |
| super().__init__() | |
| base_dim = 16 | |
| self.conv_net = nn.Sequential( | |
| nn.Conv2d(1, base_dim, kernel_size=3, stride=2, padding=1), | |
| nn.BatchNorm2d(base_dim), | |
| nn.ReLU(), | |
| nn.Conv2d(base_dim, base_dim*2, kernel_size=3, stride=2, padding=1), | |
| nn.BatchNorm2d(base_dim*2), | |
| nn.ReLU(), | |
| nn.Conv2d(base_dim*2, base_dim*4, kernel_size=3, stride=2, padding=1), | |
| nn.BatchNorm2d(base_dim*4), | |
| nn.ReLU(), | |
| ) | |
| self.gap = nn.AdaptiveAvgPool2d((1, 1)) | |
| self.proj = nn.Linear(base_dim * 4, output_dim) | |
| def forward(self, masks: torch.Tensor) -> torch.Tensor: | |
| if masks.dim() == 3: | |
| x = masks.unsqueeze(1) | |
| else: | |
| x = masks | |
| B, C, H, W = x.shape | |
| x = F.interpolate(x, size=(512, 512), mode='nearest') | |
| feat = self.conv_net(x) | |
| feat = self.gap(feat).flatten(1) | |
| feat = self.proj(feat) | |
| return feat | |
| class Box3DPromptEncoder(nn.Module): | |
| def __init__(self, output_dim: int, fourier_dim: int = 128, fourier_scale: float = 4.0): | |
| super().__init__() | |
| self.output_dim = output_dim | |
| self.gff_encoder = GaussianFourierProjection( | |
| embedding_size=fourier_dim, scale=fourier_scale,log=False | |
| ) | |
| gff_output_dim = fourier_dim * 2 * 3 | |
| self.corner_type_embedding = nn.Embedding(2, gff_output_dim) | |
| self.corner_mlp = nn.Sequential( | |
| nn.Linear(gff_output_dim, self.output_dim), | |
| nn.LayerNorm(self.output_dim), | |
| nn.GELU(), | |
| nn.Linear(self.output_dim, self.output_dim) | |
| ) | |
| def _encode_corner(self, xyz_coords: torch.Tensor, corner_type: int) -> torch.Tensor: | |
| x = xyz_coords[:, 0] | |
| y = xyz_coords[:, 1] | |
| z = xyz_coords[:, 2] | |
| x_emb = self.gff_encoder(x) | |
| y_emb = self.gff_encoder(y) | |
| z_emb = self.gff_encoder(z) | |
| features = torch.cat([x_emb, y_emb, z_emb], dim=-1) | |
| type_emb = self.corner_type_embedding( | |
| torch.tensor(corner_type, device=features.device) | |
| ) | |
| return self.corner_mlp(features + type_emb) | |
| def forward(self, box_3d: torch.Tensor) -> torch.Tensor: | |
| box_3d = box_3d | |
| # 分解为 Min 角和 Max 角 | |
| min_corner = box_3d[:, 0, :] # [B, 3] | |
| max_corner = box_3d[:, 1, :] # [B, 3] | |
| # 分别编码 | |
| emb_min = self._encode_corner(min_corner, corner_type=0) | |
| emb_max = self._encode_corner(max_corner, corner_type=1) | |
| box_embedding = emb_min + emb_max | |
| return box_embedding | |
| class SinkhornWithDustbin(nn.Module): | |
| """ | |
| Sinkhorn in log-space with a single learnable dustbin score. | |
| Supports inputs of shape [N, M] or [B, N, M] (square or rectangular). | |
| Assumes logits already scaled by external logit_scale. | |
| """ | |
| def __init__(self, n_iters: int = 50, learnable=True, init_value: float = 1.0): | |
| super().__init__() | |
| self.n_iters = n_iters | |
| if learnable: | |
| self.bin_score = nn.Parameter(torch.tensor(init_value)) | |
| else: | |
| self.register_buffer("bin_score", torch.tensor(init_value)) | |
| def forward(self, logits, mask_row=None, mask_col=None): | |
| """ | |
| logits: [N, M] or [B, N, M] (not yet augmented) | |
| mask_row: optional bool mask for valid rows shape [N] or [B, N] | |
| mask_col: optional bool mask for valid cols shape [M] or [B, M] | |
| returns dict with augmented P and components: | |
| "P_aug": [B, N+1, M+1] or [N+1, M+1] | |
| "P_match": [B, N, M] or [N, M] | |
| "P_row_bin": [B, N] or [N] | |
| "P_col_bin": [B, M] or [M] | |
| "P_bin": scalar or [B] | |
| """ | |
| single = logits.dim() == 2 | |
| if single: | |
| logits = logits.unsqueeze(0) # [1, N, M] | |
| B, N, M = logits.shape | |
| device = logits.device | |
| b = self.bin_score | |
| # Build augmented logit matrix L of shape [B, N+1, M+1] | |
| L = logits.new_full((B, N + 1, M + 1), fill_value=-1e4) # large negative for unused by default | |
| # top-left: original logits | |
| L[:, :N, :M] = logits | |
| # top-right: row -> bin (N x 1) fill with b | |
| L[:, :N, M] = b | |
| # bottom-left: bin <- col (1 x M) fill with b | |
| L[:, N, :M] = b | |
| # bottom-right: bin-bin = 0 | |
| L[:, N, M] = 0.0 | |
| # Optional: mask invalid rows/cols by setting to -inf so they get zero prob. | |
| # mask_row: if provided, shape [B, N] or [N] | |
| if mask_row is not None: | |
| mr = mask_row | |
| if mr.dim() == 1: | |
| mr = mr.unsqueeze(0).expand(B, -1) # [B, N] | |
| # where invalid rows -> set entire that row (except bin) to large negative | |
| invalid = (~mr).to(torch.bool) | |
| if invalid.any(): | |
| L[invalid, :M] = -1e9 # sets top-left entries in invalid rows | |
| L[invalid, M] = b # keep the row->bin entry as b (or set to -inf if you want forbid) | |
| if mask_col is not None: | |
| mc = mask_col | |
| if mc.dim() == 1: | |
| mc = mc.unsqueeze(0).expand(B, -1) # [B, M] | |
| invalidc = (~mc).to(torch.bool) | |
| if invalidc.any(): | |
| L[invalidc, :, :M] = -1e9 # careful broadcasting, you may want explicit loops | |
| L[:, N, invalidc] = b # keep col->bin | |
| # Sinkhorn in log-space | |
| log_alpha = L # already logits scaled outside | |
| for _ in range(self.n_iters): | |
| # row normalize: logsumexp over columns | |
| log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=2, keepdim=True) | |
| # col normalize: logsumexp over rows | |
| log_alpha = log_alpha - torch.logsumexp(log_alpha, dim=1, keepdim=True) | |
| P = torch.exp(log_alpha) # [B, N+1, M+1] | |
| P_match = P[:, :N, :M] | |
| P_row_bin = P[:, :N, M] | |
| P_col_bin = P[:, N, :M] | |
| P_bin = P[:, N, M] | |
| if single: | |
| # squeeze batch dim | |
| return { | |
| "P_aug": P[0], | |
| "P_match": P_match[0], | |
| "P_row_bin": P_row_bin[0], | |
| "P_col_bin": P_col_bin[0], | |
| "P_bin": P_bin[0] | |
| } | |
| else: | |
| return { | |
| "P_aug": P, | |
| "P_match": P_match, | |
| "P_row_bin": P_row_bin, | |
| "P_col_bin": P_col_bin, | |
| "P_bin": P_bin | |
| } | |
| class PartVisualBridge(nn.Module): | |
| def __init__(self, part_dim=256, image_dim=768, num_heads=8): | |
| super().__init__() | |
| self.cross_attn = nn.MultiheadAttention(embed_dim=part_dim, kdim=image_dim, vdim=image_dim, num_heads=num_heads, batch_first=True) | |
| self.norm1 = nn.LayerNorm(part_dim) | |
| self.norm2 = nn.LayerNorm(part_dim) | |
| self.ffn = nn.Sequential( | |
| nn.Linear(part_dim, part_dim * 4), | |
| nn.GELU(), | |
| nn.Linear(part_dim * 4, part_dim) | |
| ) | |
| def forward(self, part_emb, image_emb): | |
| # 1. Cross Attention | |
| # query=part, key=value=image | |
| part_emb = part_emb.unsqueeze(1) | |
| attn_out, _ = self.cross_attn(query=self.norm1(part_emb), key=image_emb, value=image_emb) | |
| # Residual Connection | |
| x = part_emb + attn_out | |
| # 2. FFN | |
| x = x + self.ffn(self.norm2(x)) | |
| return x.squeeze(1) | |
| class DiTBlock(nn.Module): | |
| def __init__( | |
| self, | |
| dim: int, | |
| num_attention_heads: int, | |
| adaln_input_dim: Optional[int] = None,# <--- 必需参数 | |
| use_self_attention: bool = True, | |
| self_attention_norm_type: Optional[str] = None, | |
| use_cross_attention: bool = True, | |
| cross_attention_dim: Optional[int] = None, | |
| cross_attention_norm_type: Optional[str] = "fp32_layer_norm", | |
| dropout=0.0, | |
| activation_fn: str = "gelu", | |
| norm_type: str = "fp32_layer_norm", | |
| norm_eps: float = 1e-5, | |
| final_dropout: bool = False, | |
| ff_inner_dim: Optional[int] = None, | |
| ff_bias: bool = True, | |
| skip: bool = False, | |
| skip_concat_front: bool = False, | |
| skip_norm_last: bool = False, | |
| qk_norm: bool = True, | |
| qkv_bias: bool = True, | |
| ): | |
| super().__init__() | |
| self.use_self_attention = use_self_attention | |
| self.use_cross_attention = use_cross_attention | |
| self.skip_concat_front = skip_concat_front | |
| self.skip_norm_last = skip_norm_last | |
| self.adaln_input_dim = adaln_input_dim | |
| self.dim = dim | |
| current_norm_affine = True | |
| if use_self_attention: | |
| if ( | |
| self_attention_norm_type == "fp32_layer_norm" | |
| or self_attention_norm_type is None | |
| ): | |
| self.norm1 = FP32LayerNorm(dim, norm_eps, current_norm_affine) | |
| else: | |
| raise NotImplementedError | |
| self.attn1 = Attention( | |
| query_dim=dim, | |
| cross_attention_dim=None, | |
| dim_head=dim // num_attention_heads, | |
| heads=num_attention_heads, | |
| qk_norm="rms_norm" if qk_norm else None, | |
| eps=1e-6, | |
| bias=qkv_bias, | |
| processor=TripoSGAttnProcessor2_0(), | |
| ) | |
| else: | |
| self.norm1 = None | |
| self.attn1 = None | |
| if use_cross_attention: | |
| assert cross_attention_dim is not None | |
| self.norm2 = FP32LayerNorm(dim, norm_eps, current_norm_affine) | |
| self.attn2 = Attention( | |
| query_dim=dim, | |
| cross_attention_dim=cross_attention_dim, | |
| dim_head=dim // num_attention_heads, | |
| heads=num_attention_heads, | |
| qk_norm="rms_norm" if qk_norm else None, | |
| cross_attention_norm=cross_attention_norm_type, | |
| eps=1e-6, | |
| bias=qkv_bias, | |
| processor=TripoSGAttnProcessor2_0(), | |
| ) | |
| else: | |
| self.norm2 = None | |
| self.attn2 = None | |
| self.norm3 = FP32LayerNorm(dim, norm_eps, current_norm_affine) | |
| self.ff = FeedForward( | |
| dim, | |
| dropout=dropout, | |
| activation_fn=activation_fn, | |
| final_dropout=final_dropout, | |
| inner_dim=ff_inner_dim, | |
| bias=ff_bias, | |
| ) | |
| self.adaLN_modulation = None | |
| if skip: | |
| self.skip_norm = FP32LayerNorm(dim, norm_eps, elementwise_affine=True) | |
| self.skip_linear = nn.Linear(2 * dim, dim) | |
| else: | |
| self.skip_linear = None | |
| self._chunk_size = None | |
| self._chunk_dim = 0 | |
| def set_topk(self, topk): | |
| self.flash_processor.topk = topk | |
| def set_flash_processor(self, flash_processor): | |
| self.flash_processor = flash_processor | |
| self.attn2.processor = self.flash_processor | |
| # Copied from diffusers.models.attention.BasicTransformerBlock.set_chunk_feed_forward | |
| def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): | |
| # Sets chunk feed-forward | |
| self._chunk_size = chunk_size | |
| self._chunk_dim = dim | |
| def add_adaln_modules(self): | |
| num_params = 0 | |
| if self.use_self_attention: | |
| num_params += self.dim * 3 | |
| if self.use_cross_attention: | |
| num_params += self.dim * 3 | |
| num_params += self.dim * 3 # for FeedForward | |
| self.adaLN_modulation = nn.Sequential( | |
| nn.SiLU(), | |
| nn.Linear(self.adaln_input_dim, num_params) | |
| ) | |
| def forward( | |
| self, | |
| hidden_states: torch.Tensor, | |
| encoder_hidden_states: Optional[torch.Tensor] = None, | |
| temb: Optional[torch.Tensor] = None, | |
| part_condition: Optional[torch.Tensor] = None, # <--- AdaLN 条件 | |
| image_rotary_emb: Optional[torch.Tensor] = None, | |
| skip: Optional[torch.Tensor] = None, | |
| attention_kwargs: Optional[Dict[str, Any]] = None, | |
| ) -> torch.Tensor: | |
| attention_kwargs = attention_kwargs or {} | |
| assert part_condition is not None, "Gated AdaLN modulation requires part_condition" | |
| all_params = self.adaLN_modulation(part_condition) | |
| param_list = list(all_params.chunk(all_params.shape[1] // hidden_states.shape[-1], dim=1)) | |
| if self.skip_linear is not None: | |
| cat = torch.cat( | |
| ( | |
| [skip, hidden_states] | |
| if self.skip_concat_front | |
| else [hidden_states, skip] | |
| ), | |
| dim=-1, | |
| ) | |
| if self.skip_norm_last: | |
| hidden_states = self.skip_linear(cat) | |
| hidden_states = self.skip_norm(hidden_states) | |
| else: | |
| cat = self.skip_norm(cat) | |
| hidden_states = self.skip_linear(cat) | |
| if self.use_self_attention: | |
| shift_attn1 = param_list.pop(0) | |
| scale_attn1 = param_list.pop(0) | |
| gate_attn1 = param_list.pop(0) | |
| norm_x = self.norm1(hidden_states) | |
| modulated_x = modulate(norm_x, shift_attn1, scale_attn1) | |
| attn_output = self.attn1( | |
| modulated_x, | |
| image_rotary_emb=image_rotary_emb, | |
| **attention_kwargs, | |
| ) | |
| hidden_states = hidden_states + gate_attn1.unsqueeze(1) * attn_output | |
| if self.use_cross_attention: | |
| shift_attn2 = param_list.pop(0) | |
| scale_attn2 = param_list.pop(0) | |
| gate_attn2 = param_list.pop(0) | |
| norm_x = self.norm2(hidden_states) | |
| modulated_x = modulate(norm_x, shift_attn2, scale_attn2) | |
| attn_output = self.attn2( | |
| modulated_x, | |
| encoder_hidden_states=encoder_hidden_states, | |
| image_rotary_emb=image_rotary_emb, | |
| **attention_kwargs, | |
| ) | |
| hidden_states = hidden_states + gate_attn2.unsqueeze(1) * attn_output | |
| shift_ff = param_list.pop(0) | |
| scale_ff = param_list.pop(0) | |
| gate_ff = param_list.pop(0) | |
| norm_x = self.norm3(hidden_states) | |
| modulated_x = modulate(norm_x, shift_ff, scale_ff) | |
| ff_output = self.ff(modulated_x) | |
| hidden_states = hidden_states + gate_ff.unsqueeze(1) * ff_output | |
| return hidden_states | |
| class FlexPartDiTModel(ModelMixin, ConfigMixin, PeftAdapterMixin): | |
| """ | |
| TripoSG: Diffusion model with a Transformer backbone. | |
| Inherit ModelMixin and ConfigMixin to be compatible with the sampler StableDiffusionPipeline of diffusers. | |
| Parameters: | |
| num_attention_heads (`int`, *optional*, defaults to 16): | |
| The number of heads to use for multi-head attention. | |
| attention_head_dim (`int`, *optional*, defaults to 88): | |
| The number of channels in each head. | |
| in_channels (`int`, *optional*): | |
| The number of channels in the input and output (specify if the input is **continuous**). | |
| patch_size (`int`, *optional*): | |
| The size of the patch to use for the input. | |
| activation_fn (`str`, *optional*, defaults to `"geglu"`): | |
| Activation function to use in feed-forward. | |
| sample_size (`int`, *optional*): | |
| The width of the latent images. This is fixed during training since it is used to learn a number of | |
| position embeddings. | |
| dropout (`float`, *optional*, defaults to 0.0): | |
| The dropout probability to use. | |
| cross_attention_dim (`int`, *optional*): | |
| The number of dimension in the clip text embedding. | |
| hidden_size (`int`, *optional*): | |
| The size of hidden layer in the conditioning embedding layers. | |
| num_layers (`int`, *optional*, defaults to 1): | |
| The number of layers of Transformer blocks to use. | |
| mlp_ratio (`float`, *optional*, defaults to 4.0): | |
| The ratio of the hidden layer size to the input size. | |
| learn_sigma (`bool`, *optional*, defaults to `True`): | |
| Whether to predict variance. | |
| cross_attention_dim_t5 (`int`, *optional*): | |
| The number dimensions in t5 text embedding. | |
| pooled_projection_dim (`int`, *optional*): | |
| The size of the pooled projection. | |
| text_len (`int`, *optional*): | |
| The length of the clip text embedding. | |
| text_len_t5 (`int`, *optional*): | |
| The length of the T5 text embedding. | |
| use_style_cond_and_image_meta_size (`bool`, *optional*): | |
| Whether or not to use style condition and image meta size. True for version <=1.1, False for version >= 1.2 | |
| """ | |
| _supports_gradient_checkpointing = True | |
| def __init__( | |
| self, | |
| num_attention_heads: int = 16, | |
| width: int = 2048, | |
| in_channels: int = 64, | |
| num_layers: int = 21, | |
| cross_attention_dim: int = 1024, | |
| max_num_parts: int = 32, | |
| enable_part_embedding=True, | |
| enable_local_cross_attn: bool = True, | |
| enable_global_cross_attn: bool = True, | |
| global_attn_block_ids: Optional[List[int]] = None, | |
| global_attn_block_id_range: Optional[List[int]] = None, | |
| enable_point_prompt: bool = True, | |
| point_prompt_fourier_dim: int = 128, | |
| enable_box_prompt: bool = True, | |
| box_prompt_fourier_dim: int = 128, | |
| enable_mask_prompt:bool = True, | |
| enable_3dbox_prompt:bool = False, | |
| add_module:bool = True, | |
| ): | |
| super().__init__() | |
| self.out_channels = in_channels | |
| self.num_heads = num_attention_heads | |
| self.inner_dim = width | |
| self.mlp_ratio = 4.0 | |
| self.cross_attention_dim = cross_attention_dim | |
| time_embed_dim, timestep_input_dim = self._set_time_proj( | |
| "positional", | |
| inner_dim=self.inner_dim, | |
| flip_sin_to_cos=False, | |
| freq_shift=0, | |
| time_embedding_dim=None, | |
| ) | |
| self.time_proj = TimestepEmbedding( | |
| timestep_input_dim, time_embed_dim, act_fn="gelu", out_dim=self.inner_dim | |
| ) | |
| self.proj_in = nn.Linear(self.config.in_channels, self.inner_dim, bias=True) | |
| # 调制设置 | |
| self.use_gated_adaln = enable_point_prompt | enable_box_prompt | enable_mask_prompt | enable_3dbox_prompt | |
| if self.use_gated_adaln: | |
| self.prompt_dim = 256 | |
| self.adaln_cond_dim = self.prompt_dim | |
| else: | |
| self.prompt_dim = None | |
| self.null_prompt_embedding = None | |
| self.blocks = nn.ModuleList( | |
| [ | |
| DiTBlock( | |
| dim=self.inner_dim, | |
| num_attention_heads=self.config.num_attention_heads, | |
| adaln_input_dim=self.adaln_cond_dim, | |
| use_self_attention=True, | |
| self_attention_norm_type="fp32_layer_norm", | |
| use_cross_attention=True, | |
| cross_attention_dim=cross_attention_dim, | |
| cross_attention_norm_type=None, | |
| activation_fn="gelu", | |
| norm_type="fp32_layer_norm", # TODO | |
| norm_eps=1e-5, | |
| ff_inner_dim=int(self.inner_dim * self.mlp_ratio), | |
| skip=layer > num_layers // 2, | |
| skip_concat_front=True, | |
| skip_norm_last=True, # this is an error | |
| qk_norm=True, # See http://arxiv.org/abs/2302.05442 for details. | |
| qkv_bias=False, | |
| ) | |
| for layer in range(num_layers) | |
| ] | |
| ) | |
| self.norm_out = LayerNorm(self.inner_dim) | |
| self.proj_out = nn.Linear(self.inner_dim, self.out_channels, bias=True) | |
| self.gradient_checkpointing = False | |
| self.enable_local_cross_attn = enable_local_cross_attn | |
| self.enable_global_cross_attn = enable_global_cross_attn | |
| if global_attn_block_ids is None: | |
| global_attn_block_ids = [] | |
| if global_attn_block_id_range is not None: | |
| global_attn_block_ids = list(range(global_attn_block_id_range[0], global_attn_block_id_range[1] + 1)) | |
| self.global_attn_block_ids = global_attn_block_ids | |
| if len(global_attn_block_ids) > 0: | |
| # Override self-attention processors for global attention blocks | |
| attn_processor_dict = {} | |
| modified_attn_processor = [] | |
| for layer_id in range(num_layers): | |
| for attn_id in [1, 2]: | |
| if layer_id in global_attn_block_ids: | |
| # apply to both self-attention and cross-attention | |
| attn_processor_dict[f'blocks.{layer_id}.attn{attn_id}.processor'] = PartCrafterAttnProcessor() | |
| modified_attn_processor.append(f'blocks.{layer_id}.attn{attn_id}.processor') | |
| else: | |
| attn_processor_dict[f'blocks.{layer_id}.attn{attn_id}.processor'] = TripoSGAttnProcessor2_0() | |
| self.set_attn_processor(attn_processor_dict) | |
| # logger.info(f"Modified {modified_attn_processor} to PartCrafterAttnProcessor") | |
| self.enable_part_embedding = enable_part_embedding | |
| self.max_num_parts = max_num_parts | |
| self.enable_point_prompt = enable_point_prompt | |
| self.point_prompt_fourier_dim = point_prompt_fourier_dim | |
| self.enable_box_prompt = enable_box_prompt | |
| self.box_prompt_fourier_dim = box_prompt_fourier_dim | |
| self.enable_mask_prompt = enable_mask_prompt | |
| self.enable_3dbox_prompt = enable_3dbox_prompt | |
| if add_module: | |
| self.add_modules() | |
| def add_modules(self): | |
| for block in self.blocks: | |
| if hasattr(block, "add_adaln_modules"): | |
| block.add_adaln_modules() | |
| if self.enable_part_embedding: | |
| self.part_embedding = nn.Embedding(self.max_num_parts, self.inner_dim) | |
| self.part_embedding.weight.data.normal_(mean=0.0, std=0.02) | |
| if self.enable_point_prompt: | |
| self.point_prompt_encoder = PointPromptEncoder( | |
| output_dim=self.prompt_dim, | |
| fourier_dim=self.point_prompt_fourier_dim | |
| ) | |
| self.null_point_prompt_emb = nn.Embedding(1, self.prompt_dim) | |
| nn.init.normal_(self.null_point_prompt_emb.weight, std=0.02) | |
| if self.enable_box_prompt: | |
| self.box_prompt_encoder = BoxPromptEncoder( | |
| output_dim=self.prompt_dim, | |
| fourier_dim=self.box_prompt_fourier_dim | |
| ) | |
| self.null_box_prompt_emb = nn.Embedding(1, self.prompt_dim) | |
| nn.init.normal_(self.null_box_prompt_emb.weight, std=0.02) | |
| if self.enable_mask_prompt: | |
| self.mask_prompt_encoder = MaskPromptEncoder(output_dim=self.prompt_dim) | |
| self.null_mask_prompt_emb = nn.Embedding(1, self.prompt_dim) | |
| nn.init.normal_(self.null_mask_prompt_emb.weight, std=0.02) | |
| if self.enable_3dbox_prompt: | |
| self.box3d_prompt_encoder = Box3DPromptEncoder(output_dim=self.prompt_dim) | |
| self.null_box3d_prompt_emb = nn.Embedding(1, self.prompt_dim) | |
| nn.init.normal_(self.null_box3d_prompt_emb.weight, std=0.02) | |
| self.proj_dim = 256 | |
| prompt_nums = int(self.enable_point_prompt + self.enable_box_prompt + self.enable_mask_prompt + self.enable_3dbox_prompt) | |
| self.combiner_mlp = nn.Sequential( | |
| nn.Linear(self.prompt_dim*prompt_nums, self.adaln_cond_dim), | |
| nn.ReLU(), | |
| nn.Linear(self.adaln_cond_dim, self.adaln_cond_dim) | |
| ) | |
| self.visual_bridge = PartVisualBridge(part_dim=self.prompt_dim,image_dim=self.cross_attention_dim) | |
| self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) | |
| self.sinkhorn = SinkhornWithDustbin(init_value=1.0) | |
| def _set_gradient_checkpointing( | |
| self, | |
| enable: bool = False, | |
| gradient_checkpointing_func: Optional[Callable] = None, | |
| ): | |
| # TODO: implement gradient checkpointing | |
| self.gradient_checkpointing = enable | |
| def _set_time_proj( | |
| self, | |
| time_embedding_type: str, | |
| inner_dim: int, | |
| flip_sin_to_cos: bool, | |
| freq_shift: float, | |
| time_embedding_dim: int, | |
| ) -> Tuple[int, int]: | |
| if time_embedding_type == "fourier": | |
| time_embed_dim = time_embedding_dim or inner_dim * 2 | |
| if time_embed_dim % 2 != 0: | |
| raise ValueError( | |
| f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}." | |
| ) | |
| self.time_embed = GaussianFourierProjection( | |
| time_embed_dim // 2, | |
| set_W_to_weight=False, | |
| log=False, | |
| flip_sin_to_cos=flip_sin_to_cos, | |
| ) | |
| timestep_input_dim = time_embed_dim | |
| elif time_embedding_type == "positional": | |
| time_embed_dim = time_embedding_dim or inner_dim * 4 | |
| self.time_embed = Timesteps(inner_dim, flip_sin_to_cos, freq_shift) | |
| timestep_input_dim = inner_dim | |
| else: | |
| raise ValueError( | |
| f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." | |
| ) | |
| return time_embed_dim, timestep_input_dim | |
| # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedTripoSGAttnProcessor2_0 | |
| def fuse_qkv_projections(self): | |
| """ | |
| Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) | |
| are fused. For cross-attention modules, key and value projection matrices are fused. | |
| <Tip warning={true}> | |
| This API is 🧪 experimental. | |
| </Tip> | |
| """ | |
| self.original_attn_processors = None | |
| for _, attn_processor in self.attn_processors.items(): | |
| if "Added" in str(attn_processor.__class__.__name__): | |
| raise ValueError( | |
| "`fuse_qkv_projections()` is not supported for models having added KV projections." | |
| ) | |
| self.original_attn_processors = self.attn_processors | |
| for module in self.modules(): | |
| if isinstance(module, Attention): | |
| module.fuse_projections(fuse=True) | |
| self.set_attn_processor(FusedTripoSGAttnProcessor2_0()) | |
| # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.unfuse_qkv_projections | |
| def unfuse_qkv_projections(self): | |
| """Disables the fused QKV projection if enabled. | |
| <Tip warning={true}> | |
| This API is 🧪 experimental. | |
| </Tip> | |
| """ | |
| if self.original_attn_processors is not None: | |
| self.set_attn_processor(self.original_attn_processors) | |
| # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors | |
| def attn_processors(self) -> Dict[str, AttentionProcessor]: | |
| r""" | |
| Returns: | |
| `dict` of attention processors: A dictionary containing all attention processors used in the model with | |
| indexed by its weight name. | |
| """ | |
| # set recursively | |
| processors = {} | |
| def fn_recursive_add_processors( | |
| name: str, | |
| module: torch.nn.Module, | |
| processors: Dict[str, AttentionProcessor], | |
| ): | |
| if hasattr(module, "get_processor"): | |
| processors[f"{name}.processor"] = module.get_processor() | |
| for sub_name, child in module.named_children(): | |
| fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) | |
| return processors | |
| for name, module in self.named_children(): | |
| fn_recursive_add_processors(name, module, processors) | |
| return processors | |
| # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor | |
| def set_attn_processor( | |
| self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] | |
| ): | |
| r""" | |
| Sets the attention processor to use to compute attention. | |
| Parameters: | |
| processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): | |
| The instantiated processor class or a dictionary of processor classes that will be set as the processor | |
| for **all** `Attention` layers. | |
| If `processor` is a dict, the key needs to define the path to the corresponding cross attention | |
| processor. This is strongly recommended when setting trainable attention processors. | |
| """ | |
| count = len(self.attn_processors.keys()) | |
| if isinstance(processor, dict) and len(processor) != count: | |
| raise ValueError( | |
| f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" | |
| f" number of attention layers: {count}. Please make sure to pass {count} processor classes." | |
| ) | |
| def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): | |
| if hasattr(module, "set_processor"): | |
| if not isinstance(processor, dict): | |
| module.set_processor(processor) | |
| else: | |
| module.set_processor(processor.pop(f"{name}.processor")) | |
| for sub_name, child in module.named_children(): | |
| fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) | |
| for name, module in self.named_children(): | |
| fn_recursive_attn_processor(name, module, processor) | |
| def set_default_attn_processor(self): | |
| """ | |
| Disables custom attention processors and sets the default attention implementation. | |
| """ | |
| self.set_attn_processor(TripoSGAttnProcessor2_0()) | |
| def forward( | |
| self, | |
| hidden_states: Optional[torch.Tensor], # [Total_Parts, T, D] (注意这里 batch 已经是 total parts) | |
| timestep: Union[int, float, torch.LongTensor], | |
| encoder_hidden_states: Optional[torch.Tensor] = None, # 🟢 [Total_Parts, 257, 1024] (DINO, 已扩展) | |
| image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, | |
| attention_kwargs: Optional[Dict[str, Any]] = None, | |
| return_dict: bool = True, | |
| # --- Prompts --- | |
| point_prompt_xy: Optional[torch.Tensor] = None, | |
| box_prompt_xyxy: Optional[torch.Tensor] = None, | |
| mask_prompt: Optional[torch.Tensor] = None, | |
| # --- Dropout Masks --- | |
| keep_point: Optional[torch.Tensor] = None, | |
| keep_box: Optional[torch.Tensor] = None, | |
| keep_mask: Optional[torch.Tensor] = None, | |
| box3d_prompt: Optional[torch.Tensor] = None, | |
| ): | |
| if attention_kwargs is not None: | |
| attention_kwargs = attention_kwargs.copy() | |
| lora_scale = attention_kwargs.pop("scale", 1.0) | |
| num_parts = attention_kwargs["num_parts"] # 必须有，用于 Locator 交互 | |
| else: | |
| raise ValueError("attention_kwargs with 'num_parts' is required.") | |
| total_parts_batch = hidden_states.shape[0] | |
| device = hidden_states.device | |
| if USE_PEFT_BACKEND: | |
| scale_lora_layers(self, lora_scale) | |
| _, T, _ = hidden_states.shape | |
| temb = self.time_embed(timestep).to(hidden_states.dtype) | |
| temb = self.time_proj(temb) | |
| temb = temb.unsqueeze(dim=1) | |
| hidden_states = self.proj_in(hidden_states) | |
| hidden_states = torch.cat([temb, hidden_states], dim=1) | |
| if self.enable_part_embedding: | |
| if isinstance(num_parts, torch.Tensor): | |
| part_embeddings = [] | |
| for num_part in num_parts: | |
| part_embedding = self.part_embedding(torch.arange(num_part, device=device)) | |
| part_embeddings.append(part_embedding) | |
| part_emb_base = torch.cat(part_embeddings, dim=0) | |
| elif isinstance(num_parts, int): | |
| part_emb_base = self.part_embedding(torch.arange(total_parts_batch, device=device)) | |
| # hidden_states = hidden_states + part_emb_base.unsqueeze(dim=1) | |
| hidden_states = torch.cat([part_emb_base.unsqueeze(1), hidden_states], dim=1) | |
| dino_patches = encoder_hidden_states[:, 1:, :] | |
| # --- Point Encoding --- | |
| final_point_emb = None | |
| if self.enable_point_prompt: | |
| final_point_emb = self.null_point_prompt_emb.weight.expand(total_parts_batch, -1) | |
| target_xy = point_prompt_xy | |
| if target_xy is not None: | |
| real_point_emb = self.point_prompt_encoder(target_xy) | |
| if keep_point is not None: | |
| mask_p = keep_point.view(-1, 1) | |
| final_point_emb = torch.where(mask_p, real_point_emb, final_point_emb) | |
| else: | |
| final_point_emb = real_point_emb | |
| # --- Box Encoding --- | |
| final_box_emb = None | |
| if self.enable_box_prompt: | |
| final_box_emb = self.null_box_prompt_emb.weight.expand(total_parts_batch, -1) | |
| if box_prompt_xyxy is not None: | |
| real_box_emb = self.box_prompt_encoder(box_prompt_xyxy) | |
| mask_b = keep_box.view(-1, 1) | |
| final_box_emb = torch.where(mask_b, real_box_emb, final_box_emb) | |
| # --- Mask Encoding --- | |
| final_mask_emb = None | |
| if self.enable_mask_prompt: | |
| final_mask_emb = self.null_mask_prompt_emb.weight.expand(total_parts_batch, -1) | |
| if mask_prompt is not None: | |
| real_mask_emb = self.mask_prompt_encoder(mask_prompt) | |
| mask_m = keep_mask.view(-1, 1) | |
| final_mask_emb = torch.where(mask_m, real_mask_emb, final_mask_emb) | |
| # --- 3D Box Encoding --- | |
| final_box3d_emb = None | |
| if self.enable_3dbox_prompt: | |
| final_box3d_emb = self.null_box3d_prompt_emb.weight.expand(total_parts_batch, -1) | |
| if box3d_prompt is not None: | |
| real_box3d_emb = self.box3d_prompt_encoder(box3d_prompt) | |
| # box3d_m = keep_mask.view(-1, 1) | |
| final_box3d_emb = real_box3d_emb | |
| feature_list = [] | |
| if self.enable_point_prompt: feature_list.append(final_point_emb) | |
| if self.enable_box_prompt: feature_list.append(final_box_emb) | |
| if self.enable_mask_prompt: feature_list.append(final_mask_emb) | |
| if self.enable_3dbox_prompt: feature_list.append(final_box3d_emb) | |
| combined_feature_raw = torch.cat(feature_list, dim=-1) | |
| final_adaln_input = self.combiner_mlp(combined_feature_raw) | |
| aux_loss = torch.tensor(0.0, device=device) | |
| if self.training and hasattr(self, 'sinkhorn'): | |
| if self.enable_point_prompt and keep_point is not None: | |
| student_feature = final_point_emb | |
| # Teachers | |
| mask_teacher_feat = None | |
| box_teacher_feat = None | |
| if self.enable_mask_prompt and mask_prompt is not None: | |
| mask_teacher_feat = real_mask_emb.detach() | |
| if self.enable_box_prompt and box_prompt_xyxy is not None: | |
| box_teacher_feat = real_box_emb.detach() | |
| # Bridge Alignment | |
| if (mask_teacher_feat is not None or box_teacher_feat is not None): | |
| feat_student = self.visual_bridge(student_feature, dino_patches) | |
| feat_student = F.normalize(feat_student, dim=-1) | |
| # Teacher Logits | |
| final_logits = 0.0 | |
| count = 0 | |
| if box_teacher_feat is not None: | |
| feat_box = self.visual_bridge(box_teacher_feat, dino_patches) | |
| feat_box = F.normalize(feat_box, dim=-1) | |
| final_logits += 0.5 * (self.logit_scale.exp() * feat_student @ feat_box.transpose(-2, -1)) | |
| count += 1 | |
| if mask_teacher_feat is not None: | |
| feat_mask = self.visual_bridge(mask_teacher_feat, dino_patches) | |
| feat_mask = F.normalize(feat_mask, dim=-1) | |
| final_logits += 0.5 * (self.logit_scale.exp() * feat_student @ feat_mask.transpose(-2, -1)) | |
| count += 1 | |
| if count == 1: final_logits *= 2.0 | |
| # Sinkhorn Loss | |
| out = self.sinkhorn(final_logits) | |
| log_P = torch.log(out["P_match"] + 1e-8) | |
| labels = torch.arange(log_P.shape[0], device=log_P.device) | |
| aux_loss = F.nll_loss(log_P, labels) | |
| # 🟢 汇总 | |
| all_loss = aux_loss | |
| # prepare negative encoder_hidden_states | |
| negative_encoder_hidden_states = torch.zeros_like(encoder_hidden_states) if encoder_hidden_states is not None else None | |
| skips = [] | |
| for layer, block in enumerate(self.blocks): | |
| skip = None if layer <= self.config.num_layers // 2 else skips.pop() | |
| if ( | |
| (not self.enable_local_cross_attn) | |
| and len(self.global_attn_block_ids) > 0 | |
| and (layer not in self.global_attn_block_ids) | |
| ): | |
| # If in non-global attention block and disable local cross attention, use negative encoder_hidden_states | |
| # Do not inject control signal into non-global attention block | |
| input_encoder_hidden_states = negative_encoder_hidden_states | |
| elif ( | |
| (not self.enable_global_cross_attn) | |
| and len(self.global_attn_block_ids) > 0 | |
| and (layer in self.global_attn_block_ids) | |
| ): | |
| # If in global attention block and disable global cross attention, use negative encoder_hidden_states | |
| # Do not inject control signal into global attention block | |
| input_encoder_hidden_states = negative_encoder_hidden_states | |
| else: | |
| input_encoder_hidden_states = encoder_hidden_states | |
| if len(self.global_attn_block_ids) > 0 and (layer in self.global_attn_block_ids): | |
| # Inject control signal into global attention block | |
| input_attention_kwargs = attention_kwargs | |
| else: | |
| input_attention_kwargs = None | |
| if self.training and self.gradient_checkpointing: | |
| def create_custom_forward(module): | |
| def custom_forward(*inputs): | |
| return module(*inputs) | |
| return custom_forward | |
| ckpt_kwargs: Dict[str, Any] = ( | |
| {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} | |
| ) | |
| hidden_states = torch.utils.checkpoint.checkpoint( | |
| create_custom_forward(block), | |
| hidden_states, | |
| input_encoder_hidden_states, | |
| temb, | |
| final_adaln_input, | |
| image_rotary_emb, | |
| skip, | |
| input_attention_kwargs, | |
| **ckpt_kwargs, | |
| ) | |
| else: | |
| hidden_states = block( | |
| hidden_states, | |
| encoder_hidden_states=input_encoder_hidden_states, | |
| temb=temb, | |
| part_condition = final_adaln_input, | |
| image_rotary_emb=image_rotary_emb, | |
| skip=skip, | |
| attention_kwargs=input_attention_kwargs, | |
| ) # (N, T+1, D) | |
| if layer < self.config.num_layers // 2: | |
| skips.append(hidden_states) | |
| # final layer | |
| hidden_states = self.norm_out(hidden_states) | |
| hidden_states = hidden_states[:, -T:] # (N, T, D) | |
| hidden_states = self.proj_out(hidden_states) | |
| if USE_PEFT_BACKEND: | |
| # remove `lora_scale` from each PEFT layer | |
| unscale_lora_layers(self, lora_scale) | |
| if not return_dict: | |
| return (hidden_states,) | |
| return Transformer1DModelOutput(sample=hidden_states,aux_loss=all_loss) | |
| # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking | |
| def enable_forward_chunking( | |
| self, chunk_size: Optional[int] = None, dim: int = 0 | |
| ) -> None: | |
| """ | |
| Sets the attention processor to use [feed forward | |
| chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers). | |
| Parameters: | |
| chunk_size (`int`, *optional*): | |
| The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually | |
| over each tensor of dim=`dim`. | |
| dim (`int`, *optional*, defaults to `0`): | |
| The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch) | |
| or dim=1 (sequence length). | |
| """ | |
| if dim not in [0, 1]: | |
| raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}") | |
| # By default chunk size is 1 | |
| chunk_size = chunk_size or 1 | |
| def fn_recursive_feed_forward( | |
| module: torch.nn.Module, chunk_size: int, dim: int | |
| ): | |
| if hasattr(module, "set_chunk_feed_forward"): | |
| module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) | |
| for child in module.children(): | |
| fn_recursive_feed_forward(child, chunk_size, dim) | |
| for module in self.children(): | |
| fn_recursive_feed_forward(module, chunk_size, dim) | |
| # Copied from diffusers.models.unets.unet_3d_condition.UNet3DConditionModel.disable_forward_chunking | |
| def disable_forward_chunking(self): | |
| def fn_recursive_feed_forward( | |
| module: torch.nn.Module, chunk_size: int, dim: int | |
| ): | |
| if hasattr(module, "set_chunk_feed_forward"): | |
| module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim) | |
| for child in module.children(): | |
| fn_recursive_feed_forward(child, chunk_size, dim) | |
| for module in self.children(): | |
| fn_recursive_feed_forward(module, None, 0) | |