Upload 19 files

update paint-turbo weights

hunyuan3d-paint-v2-0-turbo/unet/diffusion_pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:690a5fc63c4e263ba07dd41dc95f86f702c059c4361b863e2e21af88d8f75714
-size 3722674238
+oid sha256:24e7f1aea8a7c94cee627eb06f5265f19eeff4e19568636c5eaef050cc19ba3d
+size 7325432923

hunyuan3d-paint-v2-0-turbo/unet/modules.py

@@ -22,7 +22,6 @@
 # fine-tuning enabling code and other elements of the foregoing made publicly available
 # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
 
-
 import copy
 import json
 import os
@@ -41,7 +40,9 @@ def _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim:
     # "feed_forward_chunk_size" can be used to save memory
     if hidden_states.shape[chunk_dim] % chunk_size != 0:
         raise ValueError(
-            f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+            f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]}"
+            f"has to be divisible by chunk size: {chunk_size}."
+            f" Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
         )
 
     num_chunks = hidden_states.shape[chunk_dim] // chunk_size
@@ -51,329 +52,16 @@ def _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim:
     )
     return ff_output
 
-class PoseRoPEAttnProcessor2_0:
-    r"""
-    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
-    """
-
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-
-    def get_1d_rotary_pos_embed(
-            self,
-            dim: int,
-            pos: torch.Tensor,
-            theta: float = 10000.0,
-            linear_factor=1.0,
-            ntk_factor=1.0,
-    ):
-        assert dim % 2 == 0
-
-        theta = theta * ntk_factor
-        freqs = (
-            1.0
-            / (theta ** (torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device)[: (dim // 2)] / dim))
-            / linear_factor
-        )  # [D/2]
-        freqs = torch.outer(pos, freqs)  # type: ignore   # [S, D/2]
-        # flux, hunyuan-dit, cogvideox
-        freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()  # [S, D]
-        freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()  # [S, D]
-        return freqs_cos, freqs_sin
-
-
-    def get_3d_rotary_pos_embed(
-            self,
-            position,
-            embed_dim, 
-            voxel_resolution,
-            theta: int = 10000,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        """
-        RoPE for video tokens with 3D structure.
-
-        Args:
-        voxel_resolution (`int`):
-            The grid size of the spatial positional embedding (height, width).
-        theta (`float`):
-            Scaling factor for frequency computation.
-
-        Returns:
-            `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
-        """
-        assert position.shape[-1]==3
-
-        # Compute dimensions for each axis
-        dim_xy = embed_dim // 8 * 3
-        dim_z = embed_dim // 8 * 2
-
-        # Temporal frequencies
-        grid = torch.arange(voxel_resolution, dtype=torch.float32, device=position.device)
-        freqs_xy = self.get_1d_rotary_pos_embed(dim_xy, grid, theta=theta)
-        freqs_z = self.get_1d_rotary_pos_embed(dim_z, grid, theta=theta)
-
-        xy_cos, xy_sin = freqs_xy  # both t_cos and t_sin has shape: voxel_resolution, dim_xy
-        z_cos, z_sin = freqs_z  # both w_cos and w_sin has shape: voxel_resolution, dim_z
-
-        embed_flattn = position.view(-1, position.shape[-1])
-        x_cos = xy_cos[embed_flattn[:,0], :]
-        x_sin = xy_sin[embed_flattn[:,0], :]
-        y_cos = xy_cos[embed_flattn[:,1], :]
-        y_sin = xy_sin[embed_flattn[:,1], :]
-        z_cos = z_cos[embed_flattn[:,2], :]
-        z_sin = z_sin[embed_flattn[:,2], :]
-
-        cos = torch.cat((x_cos, y_cos, z_cos), dim=-1)
-        sin = torch.cat((x_sin, y_sin, z_sin), dim=-1)
-
-        cos = cos.view(*position.shape[:-1], embed_dim)
-        sin = sin.view(*position.shape[:-1], embed_dim)
-        return cos, sin
-
-    def apply_rotary_emb(
-            self, 
-            x: torch.Tensor,
-            freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]]
-        ):
-        cos, sin = freqs_cis  # [S, D]
-        cos, sin = cos.to(x.device), sin.to(x.device)
-        cos = cos.unsqueeze(1)
-        sin = sin.unsqueeze(1)
-
-        x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
-        x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
-
-        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
-
-        return out
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_indices: Dict = None,
-        temb: Optional[torch.Tensor] = None,
-        *args,
-        **kwargs,
-    ) -> torch.Tensor:
-        if len(args) > 0 or kwargs.get("scale", None) is not None:
-            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
-            deprecate("scale", "1.0.0", deprecation_message)
-
-        residual = hidden_states
-        if attn.spatial_norm is not None:
-            hidden_states = attn.spatial_norm(hidden_states, temb)
-
-        input_ndim = hidden_states.ndim
-
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
-        )
-
-        if attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-            # scaled_dot_product_attention expects attention_mask shape to be
-            # (batch, heads, source_length, target_length)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-
-        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = attn.to_q(hidden_states)
-
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.norm_cross:
-            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-        key = attn.to_k(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-        
-        if position_indices is not None:
-            if head_dim in position_indices:
-                image_rotary_emb = position_indices[head_dim]
-            else:
-                image_rotary_emb = self.get_3d_rotary_pos_embed(position_indices['voxel_indices'], head_dim, voxel_resolution=position_indices['voxel_resolution'])
-                position_indices[head_dim] = image_rotary_emb
-            query = self.apply_rotary_emb(query, image_rotary_emb)
-            key = self.apply_rotary_emb(key, image_rotary_emb)
-
-        # the output of sdp = (batch, num_heads, seq_len, head_dim)
-        # TODO: add support for attn.scale when we move to Torch 2.1
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-        )
-
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-        hidden_states = hidden_states.to(query.dtype)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-        if attn.residual_connection:
-            hidden_states = hidden_states + residual
-
-        hidden_states = hidden_states / attn.rescale_output_factor
-
-        return hidden_states
-
-class IPAttnProcessor2_0:
-    r"""
-    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
-    """
-
-    def __init__(self, scale=0.0):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-
-        self.scale = scale
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        ip_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        temb: Optional[torch.Tensor] = None,
-        *args,
-        **kwargs,
-    ) -> torch.Tensor:
-        if len(args) > 0 or kwargs.get("scale", None) is not None:
-            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
-            deprecate("scale", "1.0.0", deprecation_message)
-
-        residual = hidden_states
-        if attn.spatial_norm is not None:
-            hidden_states = attn.spatial_norm(hidden_states, temb)
-
-        input_ndim = hidden_states.ndim
-
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
-        )
-
-        if attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-            # scaled_dot_product_attention expects attention_mask shape to be
-            # (batch, heads, source_length, target_length)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-
-        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = attn.to_q(hidden_states)
-
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.norm_cross:
-            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-        key = attn.to_k(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-        
-
-        # the output of sdp = (batch, num_heads, seq_len, head_dim)
-        # TODO: add support for attn.scale when we move to Torch 2.1
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
-        )
-
-        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-        hidden_states = hidden_states.to(query.dtype)
-
-        # for ip adapter
-        if ip_hidden_states is not None:
-
-            ip_key = attn.to_k_ip(ip_hidden_states)
-            ip_value = attn.to_v_ip(ip_hidden_states)
-
-            ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-            ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-            # the output of sdp = (batch, num_heads, seq_len, head_dim)
-            ip_hidden_states = F.scaled_dot_product_attention(
-                query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
-            )
-
-            ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
-            ip_hidden_states = ip_hidden_states.to(query.dtype)
-
-            hidden_states = hidden_states + self.scale * ip_hidden_states
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-        if attn.residual_connection:
-            hidden_states = hidden_states + residual
-
-        hidden_states = hidden_states / attn.rescale_output_factor
-
-        return hidden_states
-        
 
 class Basic2p5DTransformerBlock(torch.nn.Module):
-    def __init__(self, transformer: BasicTransformerBlock, layer_name, use_ipa=True, use_ma=True, use_ra=True) -> None:
+    def __init__(self, transformer: BasicTransformerBlock, layer_name, use_ma=True, use_ra=True, is_turbo=False) -> None:
         super().__init__()
         self.transformer = transformer
         self.layer_name = layer_name
-        self.use_ipa = use_ipa
         self.use_ma = use_ma
         self.use_ra = use_ra
+        self.is_turbo = is_turbo
 
-        if use_ipa:
-            self.attn2.set_processor(IPAttnProcessor2_0())
-            cross_attention_dim = 1024
-            self.attn2.to_k_ip = nn.Linear(cross_attention_dim, self.dim, bias=False)
-            self.attn2.to_v_ip = nn.Linear(cross_attention_dim, self.dim, bias=False)
-            
         # multiview attn
         if self.use_ma:
             self.attn_multiview = Attention(
@@ -385,7 +73,6 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
                 cross_attention_dim=None,
                 upcast_attention=self.attn1.upcast_attention,
                 out_bias=True,
-                processor=PoseRoPEAttnProcessor2_0(),
             )
 
         # ref attn
@@ -400,8 +87,8 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
                 upcast_attention=self.attn1.upcast_attention,
                 out_bias=True,
             )
-
-        self._initialize_attn_weights()
+        if self.is_turbo:
+            self._initialize_attn_weights()
 
     def _initialize_attn_weights(self):
 
@@ -418,10 +105,6 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
                     for param in layer.parameters():
                         param.zero_()
 
-        if self.use_ipa:
-            self.attn2.to_k_ip.load_state_dict(self.attn2.to_k.state_dict()) 
-            self.attn2.to_v_ip.load_state_dict(self.attn2.to_v.state_dict()) 
-
     def __getattr__(self, name: str):
         try:
             return super().__getattr__(name)
@@ -447,10 +130,16 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
         cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
         num_in_batch = cross_attention_kwargs.pop('num_in_batch', 1)
         mode = cross_attention_kwargs.pop('mode', None)
+        if not self.is_turbo:
+            mva_scale = cross_attention_kwargs.pop('mva_scale', 1.0)
+            ref_scale = cross_attention_kwargs.pop('ref_scale', 1.0)
+        else:
+            position_attn_mask = cross_attention_kwargs.pop("position_attn_mask", None)
+            position_voxel_indices = cross_attention_kwargs.pop("position_voxel_indices", None)
+            mva_scale = 1.0
+            ref_scale = 1.0
+            
         condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None)
-        ip_hidden_states = cross_attention_kwargs.pop("ip_hidden_states", None)
-        position_attn_mask = cross_attention_kwargs.pop("position_attn_mask", None)
-        position_voxel_indices = cross_attention_kwargs.pop("position_voxel_indices", None)
 
         if self.norm_type == "ada_norm":
             norm_hidden_states = self.norm1(hidden_states, timestep)
@@ -470,10 +159,10 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
             norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
         else:
             raise ValueError("Incorrect norm used")
-            
+
         if self.pos_embed is not None:
             norm_hidden_states = self.pos_embed(norm_hidden_states)
-            
+
         # 1. Prepare GLIGEN inputs
         cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
         gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
@@ -484,6 +173,7 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
             attention_mask=attention_mask,
             **cross_attention_kwargs,
         )
+
         if self.norm_type == "ada_norm_zero":
             attn_output = gate_msa.unsqueeze(1) * attn_output
         elif self.norm_type == "ada_norm_single":
@@ -492,13 +182,17 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
         hidden_states = attn_output + hidden_states
         if hidden_states.ndim == 4:
             hidden_states = hidden_states.squeeze(1)
-        
+
         # 1.2 Reference Attention
         if 'w' in mode:
-            condition_embed_dict[self.layer_name] = rearrange(norm_hidden_states, '(b n) l c -> b (n l) c', n=num_in_batch) # B, (N L), C
-
-        if 'r' in mode:
-            condition_embed = condition_embed_dict[self.layer_name].unsqueeze(1).repeat(1,num_in_batch,1,1) # B N L C
+            condition_embed_dict[self.layer_name] = rearrange(
+                norm_hidden_states, '(b n) l c -> b (n l) c',
+                n=num_in_batch
+            )  # B, (N L), C
+
+        if 'r' in mode and self.use_ra:
+            condition_embed = condition_embed_dict[self.layer_name].unsqueeze(1).repeat(1, num_in_batch, 1,
+                                                                                        1)  # B N L C
             condition_embed = rearrange(condition_embed, 'b n l c -> (b n) l c')
 
             attn_output = self.attn_refview(
@@ -507,35 +201,48 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
                 attention_mask=None,
                 **cross_attention_kwargs
             )
+            if not self.is_turbo:
+                ref_scale_timing = ref_scale
+                if isinstance(ref_scale, torch.Tensor):
+                    ref_scale_timing = ref_scale.unsqueeze(1).repeat(1, num_in_batch).view(-1)
+                    for _ in range(attn_output.ndim - 1):
+                        ref_scale_timing = ref_scale_timing.unsqueeze(-1)
+                        
+            hidden_states = ref_scale_timing * attn_output + hidden_states
 
-            hidden_states = attn_output + hidden_states
             if hidden_states.ndim == 4:
                 hidden_states = hidden_states.squeeze(1)
-            
 
         # 1.3 Multiview Attention
         if num_in_batch > 1 and self.use_ma:
             multivew_hidden_states = rearrange(norm_hidden_states, '(b n) l c -> b (n l) c', n=num_in_batch)
-            position_mask = None
-            if position_attn_mask is not None:
-                if multivew_hidden_states.shape[1] in position_attn_mask:
-                    position_mask = position_attn_mask[multivew_hidden_states.shape[1]]
-            position_indices = None
-            if position_voxel_indices is not None:
-                if multivew_hidden_states.shape[1] in position_voxel_indices:
-                    position_indices = position_voxel_indices[multivew_hidden_states.shape[1]]
-
-            attn_output = self.attn_multiview(
-                multivew_hidden_states,
-                encoder_hidden_states=multivew_hidden_states,
-                attention_mask=position_mask,
-                position_indices=position_indices,
-                **cross_attention_kwargs
-            )
 
-            attn_output = rearrange(attn_output, 'b (n l) c -> (b n) l c', n=num_in_batch)
+            if self.is_turbo:
+                position_mask = None
+                if position_attn_mask is not None:
+                    if multivew_hidden_states.shape[1] in position_attn_mask:
+                        position_mask = position_attn_mask[multivew_hidden_states.shape[1]]
+                position_indices = None
+                if position_voxel_indices is not None:
+                    if multivew_hidden_states.shape[1] in position_voxel_indices:
+                        position_indices = position_voxel_indices[multivew_hidden_states.shape[1]]
+                attn_output = self.attn_multiview(
+                    multivew_hidden_states,
+                    encoder_hidden_states=multivew_hidden_states,
+                    attention_mask=position_mask,
+                    position_indices=position_indices,
+                    **cross_attention_kwargs
+                )
+            else:
+                attn_output = self.attn_multiview(
+                    multivew_hidden_states,
+                    encoder_hidden_states=multivew_hidden_states,
+                    **cross_attention_kwargs
+                )
 
-            hidden_states = attn_output + hidden_states
+            attn_output = rearrange(attn_output, 'b (n l) c -> (b n) l c', n=num_in_batch)
+            
+            hidden_states = mva_scale * attn_output + hidden_states
             if hidden_states.ndim == 4:
                 hidden_states = hidden_states.squeeze(1)
 
@@ -561,25 +268,12 @@ class Basic2p5DTransformerBlock(torch.nn.Module):
             if self.pos_embed is not None and self.norm_type != "ada_norm_single":
                 norm_hidden_states = self.pos_embed(norm_hidden_states)
 
-            if ip_hidden_states is not None:
-                ip_hidden_states = ip_hidden_states.unsqueeze(1).repeat(1,num_in_batch,1,1) # B N L C
-                ip_hidden_states = rearrange(ip_hidden_states, 'b n l c -> (b n) l c')
-
-            if self.use_ipa:
-                attn_output = self.attn2(
-                    norm_hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    ip_hidden_states=ip_hidden_states,
-                    attention_mask=encoder_attention_mask,
-                    **cross_attention_kwargs,
-                )
-            else:
-                attn_output = self.attn2(
-                    norm_hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    attention_mask=encoder_attention_mask,
-                    **cross_attention_kwargs,
-                )
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                **cross_attention_kwargs,
+            )
 
             hidden_states = attn_output + hidden_states
 
@@ -626,8 +320,16 @@ def compute_voxel_grid_mask(position, grid_resolution=8):
     position[valid_mask==False] = 0
 
     
-    position = rearrange(position, 'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', num_h=grid_resolution, num_w=grid_resolution)
-    valid_mask = rearrange(valid_mask, 'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', num_h=grid_resolution, num_w=grid_resolution)
+    position = rearrange(
+        position,
+        'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', 
+        num_h=grid_resolution, num_w=grid_resolution
+    )
+    valid_mask = rearrange(
+        valid_mask, 
+        'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', 
+        num_h=grid_resolution, num_w=grid_resolution
+    )
 
     grid_position = position.sum(dim=(-2, -1))
     count_masked = valid_mask.sum(dim=(-2, -1))
@@ -674,8 +376,16 @@ def compute_discrete_voxel_indice(position, grid_resolution=8, voxel_resolution=
     valid_mask = valid_mask.expand_as(position)
     position[valid_mask==False] = 0
     
-    position = rearrange(position, 'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', num_h=grid_resolution, num_w=grid_resolution)
-    valid_mask = rearrange(valid_mask, 'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', num_h=grid_resolution, num_w=grid_resolution)
+    position = rearrange(
+        position, 
+        'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', 
+        num_h=grid_resolution, num_w=grid_resolution
+    )
+    valid_mask = rearrange(
+        valid_mask, 
+        'b n c (num_h grid_h) (num_w grid_w) -> b n num_h num_w c grid_h grid_w', 
+        num_h=grid_resolution, num_w=grid_resolution
+    )
 
     grid_position = position.sum(dim=(-2, -1))
     count_masked = valid_mask.sum(dim=(-2, -1))
@@ -688,45 +398,36 @@ def compute_discrete_voxel_indice(position, grid_resolution=8, voxel_resolution=
     voxel_indices = torch.round(voxel_indices).long()
     return voxel_indices
     
-def compute_multi_resolution_discrete_voxel_indice(position_maps, grid_resolutions=[64, 32, 16, 8], voxel_resolutions=[512, 256, 128, 64]):
+def compute_multi_resolution_discrete_voxel_indice(
+    position_maps, 
+    grid_resolutions=[64, 32, 16, 8], 
+    voxel_resolutions=[512, 256, 128, 64]
+):
     voxel_indices = {}
     with torch.no_grad():
         for grid_resolution, voxel_resolution in zip(grid_resolutions, voxel_resolutions):
             voxel_indice = compute_discrete_voxel_indice(position_maps, grid_resolution, voxel_resolution)
             voxel_indice = rearrange(voxel_indice, 'b n c h w -> b (n h w) c')
             voxel_indices[voxel_indice.shape[1]] = {'voxel_indices':voxel_indice, 'voxel_resolution':voxel_resolution}
-    return voxel_indices   
-    
-class ImageProjModel(torch.nn.Module):
-    """Projection Model"""
-
-    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
-        super().__init__()
-
-        self.generator = None
-        self.cross_attention_dim = cross_attention_dim
-        self.clip_extra_context_tokens = clip_extra_context_tokens
-        self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
-        self.norm = torch.nn.LayerNorm(cross_attention_dim)
+    return voxel_indices
 
-    def forward(self, image_embeds):
-        embeds = image_embeds
-        clip_extra_context_tokens = self.proj(embeds).reshape(
-            -1, self.clip_extra_context_tokens, self.cross_attention_dim
-        )
-        clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
-        return clip_extra_context_tokens
-        
 class UNet2p5DConditionModel(torch.nn.Module):
     def __init__(self, unet: UNet2DConditionModel) -> None:
         super().__init__()
         self.unet = unet
-        self.unet_dual = copy.deepcopy(unet)
 
-        self.init_camera_embedding()
-        self.init_attention(self.unet, use_ipa=True, use_ma=True, use_ra=True)
-        self.init_attention(self.unet_dual, use_ipa=False, use_ma=False, use_ra=False)
+        self.use_ma = True
+        self.use_ra = True
+        self.use_camera_embedding = True
+        self.use_dual_stream = True
+        self.is_turbo = False
+
+        if self.use_dual_stream:
+            self.unet_dual = copy.deepcopy(unet)
+            self.init_attention(self.unet_dual)
+        self.init_attention(self.unet, use_ma=self.use_ma, use_ra=self.use_ra, is_turbo=self.is_turbo)
         self.init_condition()
+        self.init_camera_embedding()
 
     @staticmethod
     def from_pretrained(pretrained_model_name_or_path, **kwargs):
@@ -737,170 +438,158 @@ class UNet2p5DConditionModel(torch.nn.Module):
             config = json.load(file)
         unet = UNet2DConditionModel(**config)
         unet = UNet2p5DConditionModel(unet)
-
-        unet.unet.conv_in = torch.nn.Conv2d(
-            12,
-            unet.unet.conv_in.out_channels,
-            kernel_size=unet.unet.conv_in.kernel_size,
-            stride=unet.unet.conv_in.stride,
-            padding=unet.unet.conv_in.padding,
-            dilation=unet.unet.conv_in.dilation,
-            groups=unet.unet.conv_in.groups,
-            bias=unet.unet.conv_in.bias is not None)
-        
         unet_ckpt = torch.load(unet_ckpt_path, map_location='cpu', weights_only=True)
         unet.load_state_dict(unet_ckpt, strict=True)
         unet = unet.to(torch_dtype)
         return unet
-        
-    def init_condition(self):
-        self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1,77,1024))
-        self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1,77,1024))
 
-        self.unet.image_proj_model = ImageProjModel(
-            cross_attention_dim=self.unet.config.cross_attention_dim,
-            clip_embeddings_dim=1024,
-        )
+    def init_condition(self):
+        self.unet.conv_in = torch.nn.Conv2d(
+            12,
+            self.unet.conv_in.out_channels,
+            kernel_size=self.unet.conv_in.kernel_size,
+            stride=self.unet.conv_in.stride,
+            padding=self.unet.conv_in.padding,
+            dilation=self.unet.conv_in.dilation,
+            groups=self.unet.conv_in.groups,
+            bias=self.unet.conv_in.bias is not None)
 
+        self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1, 77, 1024))
+        self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1, 77, 1024))
 
     def init_camera_embedding(self):
-        self.max_num_ref_image = 5
-        self.max_num_gen_image = 12*3+4*2
 
-        time_embed_dim = 1280
-        self.unet.class_embedding = nn.Embedding(self.max_num_ref_image+self.max_num_gen_image, time_embed_dim)
-        # 将嵌入层的权重初始化为全零
-        nn.init.zeros_(self.unet.class_embedding.weight)
-    
-    def init_attention(self, unet, use_ipa=True, use_ma=True, use_ra=True):
+        if self.use_camera_embedding:
+            time_embed_dim = 1280
+            self.max_num_ref_image = 5
+            self.max_num_gen_image = 12 * 3 + 4 * 2
+            self.unet.class_embedding = nn.Embedding(self.max_num_ref_image + self.max_num_gen_image, time_embed_dim)
+
+    def init_attention(self, unet, use_ma=False, use_ra=False, is_turbo=False):
 
         for down_block_i, down_block in enumerate(unet.down_blocks):
             if hasattr(down_block, "has_cross_attention") and down_block.has_cross_attention:
                 for attn_i, attn in enumerate(down_block.attentions):
                     for transformer_i, transformer in enumerate(attn.transformer_blocks):
                         if isinstance(transformer, BasicTransformerBlock):
-                            attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'down_{down_block_i}_{attn_i}_{transformer_i}',use_ipa,use_ma,use_ra)
+                            attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(
+                                transformer,
+                                f'down_{down_block_i}_{attn_i}_{transformer_i}',
+                                use_ma, use_ra, is_turbo
+                            )
 
         if hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention:
             for attn_i, attn in enumerate(unet.mid_block.attentions):
                 for transformer_i, transformer in enumerate(attn.transformer_blocks):
                     if isinstance(transformer, BasicTransformerBlock):
-                        attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'mid_{attn_i}_{transformer_i}',use_ipa,use_ma,use_ra)
+                        attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(
+                            transformer,
+                            f'mid_{attn_i}_{transformer_i}',
+                            use_ma, use_ra, is_turbo
+                        )
 
         for up_block_i, up_block in enumerate(unet.up_blocks):
             if hasattr(up_block, "has_cross_attention") and up_block.has_cross_attention:
                 for attn_i, attn in enumerate(up_block.attentions):
                     for transformer_i, transformer in enumerate(attn.transformer_blocks):
                         if isinstance(transformer, BasicTransformerBlock):
-                            attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'up_{up_block_i}_{attn_i}_{transformer_i}',use_ipa,use_ma,use_ra)
-
+                            attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(
+                                transformer,
+                                f'up_{up_block_i}_{attn_i}_{transformer_i}',
+                                use_ma, use_ra, is_turbo
+                            )
 
     def __getattr__(self, name: str):
         try:
             return super().__getattr__(name)
         except AttributeError:
             return getattr(self.unet, name)
-        
+
     def forward(
-        self, sample, timestep, encoder_hidden_states, class_labels=None,
-        *args, cross_attention_kwargs=None, down_intrablock_additional_residuals=None,
+        self, sample, timestep, encoder_hidden_states,
+        *args, down_intrablock_additional_residuals=None,
         down_block_res_samples=None, mid_block_res_sample=None,
         **cached_condition,
     ):
         B, N_gen, _, H, W = sample.shape
-        camera_info_gen = cached_condition['camera_info_gen'] + self.max_num_ref_image
-        camera_info_gen = rearrange(camera_info_gen, 'b n -> (b n)')
+        assert H == W
+
+        if self.use_camera_embedding:
+            camera_info_gen = cached_condition['camera_info_gen'] + self.max_num_ref_image
+            camera_info_gen = rearrange(camera_info_gen, 'b n -> (b n)')
+        else:
+            camera_info_gen = None
+
         sample = [sample]
-        
         if 'normal_imgs' in cached_condition:
             sample.append(cached_condition["normal_imgs"])
         if 'position_imgs' in cached_condition:
             sample.append(cached_condition["position_imgs"])
-
         sample = torch.cat(sample, dim=2)
+
         sample = rearrange(sample, 'b n c h w -> (b n) c h w')
 
         encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(1).repeat(1, N_gen, 1, 1)
         encoder_hidden_states_gen = rearrange(encoder_hidden_states_gen, 'b n l c -> (b n) l c')
-        
-        
-        use_position_mask = False
-        use_position_rope = True
-
-        position_attn_mask = None
-        if use_position_mask:
-            if 'position_attn_mask' in cached_condition:
-                position_attn_mask = cached_condition['position_attn_mask']
-            else:
-                if 'position_maps' in cached_condition:
-                    position_attn_mask = compute_multi_resolution_mask(cached_condition['position_maps'])
-        
-        position_voxel_indices = None
-        if use_position_rope:
-            if 'position_voxel_indices' in cached_condition:
-                position_voxel_indices = cached_condition['position_voxel_indices']
-            else:
-                if 'position_maps' in cached_condition:
-                    position_voxel_indices = compute_multi_resolution_discrete_voxel_indice(cached_condition['position_maps'])
 
-        if 'ip_hidden_states' in cached_condition:
-            ip_hidden_states = cached_condition['ip_hidden_states']
-        else:
-            if 'clip_embeds' in cached_condition:
-                ip_hidden_states = self.image_proj_model(cached_condition['clip_embeds'])
+        if self.use_ra:
+            if 'condition_embed_dict' in cached_condition:
+                condition_embed_dict = cached_condition['condition_embed_dict']
             else:
-                ip_hidden_states = None
-            cached_condition['ip_hidden_states'] = ip_hidden_states
-
-        if 'condition_embed_dict' in cached_condition:
-            condition_embed_dict = cached_condition['condition_embed_dict']
+                condition_embed_dict = {}
+                ref_latents = cached_condition['ref_latents']
+                N_ref = ref_latents.shape[1]
+                if self.use_camera_embedding:
+                    camera_info_ref = cached_condition['camera_info_ref']
+                    camera_info_ref = rearrange(camera_info_ref, 'b n -> (b n)')
+                else:
+                    camera_info_ref = None
+
+                ref_latents = rearrange(ref_latents, 'b n c h w -> (b n) c h w')
+
+                encoder_hidden_states_ref = self.unet.learned_text_clip_ref.unsqueeze(1).repeat(B, N_ref, 1, 1)
+                encoder_hidden_states_ref = rearrange(encoder_hidden_states_ref, 'b n l c -> (b n) l c')
+
+                noisy_ref_latents = ref_latents
+                timestep_ref = 0
+
+                if self.use_dual_stream:
+                    unet_ref = self.unet_dual
+                else:
+                    unet_ref = self.unet
+                unet_ref(
+                    noisy_ref_latents, timestep_ref,
+                    encoder_hidden_states=encoder_hidden_states_ref,
+                    class_labels=camera_info_ref,
+                    # **kwargs
+                    return_dict=False,
+                    cross_attention_kwargs={
+                        'mode': 'w', 'num_in_batch': N_ref,
+                        'condition_embed_dict': condition_embed_dict},
+                )
+                cached_condition['condition_embed_dict'] = condition_embed_dict
         else:
-            condition_embed_dict = {}
-            ref_latents = cached_condition['ref_latents']
-            N_ref = ref_latents.shape[1]
-            camera_info_ref = cached_condition['camera_info_ref']
-            camera_info_ref = rearrange(camera_info_ref, 'b n -> (b n)')
-            
-            #ref_latents = [ref_latents]
-            #if 'normal_imgs' in cached_condition:
-            #    ref_latents.append(torch.zeros_like(ref_latents[0]))
-            #if 'position_imgs' in cached_condition:
-            #    ref_latents.append(torch.zeros_like(ref_latents[0]))
-            #ref_latents = torch.cat(ref_latents, dim=2)
-            
-            ref_latents = rearrange(ref_latents, 'b n c h w -> (b n) c h w')
+            condition_embed_dict = None
 
-            encoder_hidden_states_ref = self.learned_text_clip_ref.unsqueeze(1).repeat(B, N_ref, 1, 1)
-            encoder_hidden_states_ref = rearrange(encoder_hidden_states_ref, 'b n l c -> (b n) l c')
-
-            noisy_ref_latents = ref_latents
-            timestep_ref = 0
-            '''
-            if timestep.dim()>0:
-                timestep_ref = rearrange(timestep, '(b n) -> b n', b=B)[:,:1].repeat(1, N_ref)
-                timestep_ref = rearrange(timestep_ref, 'b n -> (b n)')
-            else:
-                timestep_ref = timestep
-            noise = torch.randn_like(noisy_ref_latents[:,:4,...])
-            if self.training:
-                noisy_ref_latents[:,:4,...] = self.train_sched.add_noise(noisy_ref_latents[:,:4,...], noise, timestep_ref)
-                noisy_ref_latents[:,:4,...] = self.train_sched.scale_model_input(noisy_ref_latents[:,:4,...], timestep_ref)
-            else:
-                noisy_ref_latents[:,:4,...] = self.val_sched.add_noise(noisy_ref_latents[:,:4,...], noise, timestep_ref.reshape(-1))
-                noisy_ref_latents[:,:4,...] = self.val_sched.scale_model_input(noisy_ref_latents[:,:4,...], timestep_ref.reshape(-1))
-            '''
-            self.unet_dual(
-                noisy_ref_latents, timestep_ref,
-                encoder_hidden_states=encoder_hidden_states_ref,
-                #class_labels=camera_info_ref,
-                # **kwargs
-                return_dict=False,
-                cross_attention_kwargs={
-                    'mode':'w', 'num_in_batch':N_ref, 
-                    'condition_embed_dict':condition_embed_dict},
-            )
-            cached_condition['condition_embed_dict'] = condition_embed_dict
+        mva_scale = cached_condition.get('mva_scale', 1.0)
+        ref_scale = cached_condition.get('ref_scale', 1.0)
 
+        if self.is_turbo:
+            cross_attention_kwargs_ = {
+                'mode': 'r', 'num_in_batch': N_gen,
+                'condition_embed_dict': condition_embed_dict,
+                'position_attn_mask':position_attn_mask, 
+                'position_voxel_indices':position_voxel_indices,
+                'mva_scale': mva_scale,
+                'ref_scale': ref_scale,
+            }
+        else:
+            cross_attention_kwargs_ = {
+                'mode': 'r', 'num_in_batch': N_gen,
+                'condition_embed_dict': condition_embed_dict,
+                'mva_scale': mva_scale,
+                'ref_scale': ref_scale,
+            }
         return self.unet(
             sample, timestep,
             encoder_hidden_states_gen, *args,
@@ -916,11 +605,6 @@ class UNet2p5DConditionModel(torch.nn.Module):
                 if mid_block_res_sample is not None else None
             ),
             return_dict=False,
-            cross_attention_kwargs={
-                'mode':'r', 'num_in_batch':N_gen, 
-                'ip_hidden_states':ip_hidden_states,
-                'condition_embed_dict':condition_embed_dict, 
-                'position_attn_mask':position_attn_mask, 
-                'position_voxel_indices':position_voxel_indices
-            },
-        ) 
+            cross_attention_kwargs=cross_attention_kwargs_,
+        )
+