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)
@maybe_allow_in_graph
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
@register_to_config
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.
This API is 🧪 experimental.
"""
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.
This API is 🧪 experimental.
"""
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
@property
# 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)