Upload hunyuan3d-paintpbr-v2-1/unet/modules.py with huggingface_hub

hunyuan3d-paintpbr-v2-1/unet/modules.py

@@ -0,0 +1,1102 @@
+# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
+# except for the third-party components listed below.
+# Hunyuan 3D does not impose any additional limitations beyond what is outlined
+# in the repsective licenses of these third-party components.
+# Users must comply with all terms and conditions of original licenses of these third-party
+# components and must ensure that the usage of the third party components adheres to
+# all relevant laws and regulations.
+
+# For avoidance of doubts, Hunyuan 3D means the large language models and
+# their software and algorithms, including trained model weights, parameters (including
+# optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
+# fine-tuning enabling code and other elements of the foregoing made publicly available
+# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+import os
+import json
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from typing import Any, Callable, Dict, List, Optional, Union, Tuple, Literal
+import diffusers
+from diffusers.utils import deprecate
+from diffusers import (
+    DDPMScheduler,
+    EulerAncestralDiscreteScheduler,
+    UNet2DConditionModel,
+)
+from diffusers.models import UNet2DConditionModel
+from diffusers.models.attention_processor import Attention, AttnProcessor
+from diffusers.models.transformers.transformer_2d import BasicTransformerBlock
+from .attn_processor import SelfAttnProcessor2_0, RefAttnProcessor2_0, PoseRoPEAttnProcessor2_0
+
+from transformers import AutoImageProcessor, AutoModel
+
+
+class Dino_v2(nn.Module):
+
+    """Wrapper for DINOv2 vision transformer (frozen weights).
+    
+    Provides feature extraction for reference images.
+    
+    Args:
+        dino_v2_path: Custom path to DINOv2 model weights (uses default if None)
+    """
+
+
+    def __init__(self, dino_v2_path):
+        super(Dino_v2, self).__init__()
+        self.dino_processor = AutoImageProcessor.from_pretrained(dino_v2_path)
+        self.dino_v2 = AutoModel.from_pretrained(dino_v2_path)
+
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.dino_v2.eval()
+
+    def forward(self, images):
+
+        """Processes input images through DINOv2 ViT.
+        
+        Handles both tensor input (B, N, C, H, W) and PIL image lists.
+        Extracts patch embeddings and flattens spatial dimensions.
+        
+        Returns:
+            torch.Tensor: Feature vectors [B, N*(num_patches), feature_dim]
+        """
+
+        if isinstance(images, torch.Tensor):
+            batch_size = images.shape[0]
+            dino_proceesed_images = self.dino_processor(
+                images=rearrange(images, "b n c h w -> (b n) c h w"), return_tensors="pt", do_rescale=False
+            ).pixel_values
+        else:
+            batch_size = 1
+            dino_proceesed_images = self.dino_processor(images=images, return_tensors="pt").pixel_values
+            dino_proceesed_images = torch.stack(
+                [torch.from_numpy(np.array(image)) for image in dino_proceesed_images], dim=0
+            )
+        dino_param = next(self.dino_v2.parameters())
+        dino_proceesed_images = dino_proceesed_images.to(dino_param)
+        dino_hidden_states = self.dino_v2(dino_proceesed_images)[0]
+        dino_hidden_states = rearrange(dino_hidden_states.to(dino_param), "(b n) l c -> b (n l) c", b=batch_size)
+
+        return dino_hidden_states
+
+
+def _chunked_feed_forward(ff: nn.Module, hidden_states: torch.Tensor, chunk_dim: int, chunk_size: int):
+    # "feed_forward_chunk_size" can be used to save memory
+
+    """Memory-efficient feedforward execution via chunking.
+    
+    Divides input along specified dimension for sequential processing.
+    
+    Args:
+        ff: Feedforward module to apply
+        hidden_states: Input tensor
+        chunk_dim: Dimension to split
+        chunk_size: Size of each chunk
+        
+    Returns:
+        torch.Tensor: Reassembled output tensor
+    """
+
+    if hidden_states.shape[chunk_dim] % chunk_size != 0:
+        raise ValueError(
+            f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]}"
+            f"has to be divisible by chunk size: {chunk_size}."
+            "Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+        )
+
+    num_chunks = hidden_states.shape[chunk_dim] // chunk_size
+    ff_output = torch.cat(
+        [ff(hid_slice) for hid_slice in hidden_states.chunk(num_chunks, dim=chunk_dim)],
+        dim=chunk_dim,
+    )
+    return ff_output
+
+
+@torch.no_grad()
+def compute_voxel_grid_mask(position, grid_resolution=8):
+
+    """Generates view-to-view attention mask based on 3D position similarity.
+    
+    Uses voxel grid downsampling to determine spatially adjacent regions.
+    Mask indicates where features should interact across different views.
+    
+    Args:
+        position: Position maps [B, N, 3, H, W] (normalized 0-1)
+        grid_resolution: Spatial reduction factor
+        
+    Returns:
+        torch.Tensor: Attention mask [B, N*grid_res**2, N*grid_res**2]
+    """
+
+    position = position.half()
+    B, N, _, H, W = position.shape
+    assert H % grid_resolution == 0 and W % grid_resolution == 0
+
+    valid_mask = (position != 1).all(dim=2, keepdim=True)
+    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,
+    )
+
+    grid_position = position.sum(dim=(-2, -1))
+    count_masked = valid_mask.sum(dim=(-2, -1))
+
+    grid_position = grid_position / count_masked.clamp(min=1)
+    grid_position[count_masked < 5] = 0
+
+    grid_position = grid_position.permute(0, 1, 4, 2, 3)
+    grid_position = rearrange(grid_position, "b n c h w -> b n (h w) c")
+
+    grid_position_expanded_1 = grid_position.unsqueeze(2).unsqueeze(4)  # 形状变为 B, N, 1, L, 1, 3
+    grid_position_expanded_2 = grid_position.unsqueeze(1).unsqueeze(3)  # 形状变为 B, 1, N, 1, L, 3
+
+    # 计算欧氏距离
+    distances = torch.norm(grid_position_expanded_1 - grid_position_expanded_2, dim=-1)  # 形状为 B, N, N, L, L
+
+    weights = distances
+    grid_distance = 1.73 / grid_resolution
+    weights = weights < grid_distance
+
+    return weights
+
+
+def compute_multi_resolution_mask(position_maps, grid_resolutions=[32, 16, 8]):
+
+    """Generates attention masks at multiple spatial resolutions.
+    
+    Creates pyramid of position-based masks for hierarchical attention.
+    
+    Args:
+        position_maps: Position maps [B, N, 3, H, W]
+        grid_resolutions: List of downsampling factors
+        
+    Returns:
+        dict: Resolution-specific masks keyed by flattened dimension size
+    """
+
+    position_attn_mask = {}
+    with torch.no_grad():
+        for grid_resolution in grid_resolutions:
+            position_mask = compute_voxel_grid_mask(position_maps, grid_resolution)
+            position_mask = rearrange(position_mask, "b ni nj li lj -> b (ni li) (nj lj)")
+            position_attn_mask[position_mask.shape[1]] = position_mask
+    return position_attn_mask
+
+
+@torch.no_grad()
+def compute_discrete_voxel_indice(position, grid_resolution=8, voxel_resolution=128):
+
+    """Quantizes position maps to discrete voxel indices.
+    
+    Creates sparse 3D coordinate representations for efficient hashing.
+    
+    Args:
+        position: Position maps [B, N, 3, H, W]
+        grid_resolution: Spatial downsampling factor
+        voxel_resolution: Quantization resolution
+        
+    Returns:
+        torch.Tensor: Voxel indices [B, N, grid_res, grid_res, 3]
+    """
+
+    position = position.half()
+    B, N, _, H, W = position.shape
+    assert H % grid_resolution == 0 and W % grid_resolution == 0
+
+    valid_mask = (position != 1).all(dim=2, keepdim=True)
+    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,
+    )
+
+    grid_position = position.sum(dim=(-2, -1))
+    count_masked = valid_mask.sum(dim=(-2, -1))
+
+    grid_position = grid_position / count_masked.clamp(min=1)
+    voxel_mask_thres = (H // grid_resolution) * (W // grid_resolution) // (4 * 4)
+    grid_position[count_masked < voxel_mask_thres] = 0
+
+    grid_position = grid_position.permute(0, 1, 4, 2, 3).clamp(0, 1)  # B N C H W
+    voxel_indices = grid_position * (voxel_resolution - 1)
+    voxel_indices = torch.round(voxel_indices).long()
+    return voxel_indices
+
+
+def calc_multires_voxel_idxs(position_maps, grid_resolutions=[64, 32, 16, 8], voxel_resolutions=[512, 256, 128, 64]):
+
+    """Generates multi-resolution voxel indices for position encoding.
+    
+    Creates pyramid of quantized position representations.
+    
+    Args:
+        position_maps: Input position maps
+        grid_resolutions: Spatial resolution levels
+        voxel_resolutions: Quantization levels
+        
+    Returns:
+        dict: Voxel indices keyed by flattened dimension size, with resolution metadata
+    """
+
+    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 Basic2p5DTransformerBlock(torch.nn.Module):
+
+
+    """Enhanced transformer block for multiview 2.5D image generation.
+    
+    Extends standard transformer blocks with:
+    - Material-specific attention (MDA)
+    - Multiview attention (MA)
+    - Reference attention (RA)
+    - DINO feature integration
+    
+    Args:
+        transformer: Base transformer block
+        layer_name: Identifier for layer
+        use_ma: Enable multiview attention
+        use_ra: Enable reference attention
+        use_mda: Enable material-aware attention
+        use_dino: Enable DINO feature integration
+        pbr_setting: List of PBR materials
+    """
+
+    def __init__(
+        self,
+        transformer: BasicTransformerBlock,
+        layer_name,
+        use_ma=True,
+        use_ra=True,
+        use_mda=True,
+        use_dino=True,
+        pbr_setting=None,
+    ) -> None:
+        
+        """
+        Initialization:
+        1. Material-Dimension Attention (MDA):
+           - Processes each PBR material with separate projection weights
+           - Uses custom SelfAttnProcessor2_0 with material awareness
+           
+        2. Multiview Attention (MA):
+           - Adds cross-view attention with PoseRoPE
+           - Initialized as zero-initialized residual pathway
+           
+        3. Reference Attention (RA):
+           - Conditions on reference view features
+           - Uses RefAttnProcessor2_0 for material-specific conditioning
+           
+        4. DINO Attention:
+           - Incorporates DINO-ViT features
+           - Initialized as zero-initialized residual pathway
+        """
+
+        super().__init__()
+        self.transformer = transformer
+        self.layer_name = layer_name
+        self.use_ma = use_ma
+        self.use_ra = use_ra
+        self.use_mda = use_mda
+        self.use_dino = use_dino
+        self.pbr_setting = pbr_setting
+
+        if self.use_mda:
+            self.attn1.set_processor(
+                SelfAttnProcessor2_0(
+                    query_dim=self.dim,
+                    heads=self.num_attention_heads,
+                    dim_head=self.attention_head_dim,
+                    dropout=self.dropout,
+                    bias=self.attention_bias,
+                    cross_attention_dim=None,
+                    upcast_attention=self.attn1.upcast_attention,
+                    out_bias=True,
+                    pbr_setting=self.pbr_setting,
+                )
+            )
+
+        # multiview attn
+        if self.use_ma:
+            self.attn_multiview = Attention(
+                query_dim=self.dim,
+                heads=self.num_attention_heads,
+                dim_head=self.attention_head_dim,
+                dropout=self.dropout,
+                bias=self.attention_bias,
+                cross_attention_dim=None,
+                upcast_attention=self.attn1.upcast_attention,
+                out_bias=True,
+                processor=PoseRoPEAttnProcessor2_0(),
+            )
+
+        # ref attn
+        if self.use_ra:
+            self.attn_refview = Attention(
+                query_dim=self.dim,
+                heads=self.num_attention_heads,
+                dim_head=self.attention_head_dim,
+                dropout=self.dropout,
+                bias=self.attention_bias,
+                cross_attention_dim=None,
+                upcast_attention=self.attn1.upcast_attention,
+                out_bias=True,
+                processor=RefAttnProcessor2_0(
+                    query_dim=self.dim,
+                    heads=self.num_attention_heads,
+                    dim_head=self.attention_head_dim,
+                    dropout=self.dropout,
+                    bias=self.attention_bias,
+                    cross_attention_dim=None,
+                    upcast_attention=self.attn1.upcast_attention,
+                    out_bias=True,
+                    pbr_setting=self.pbr_setting,
+                ),
+            )
+
+        # dino attn
+        if self.use_dino:
+            self.attn_dino = Attention(
+                query_dim=self.dim,
+                heads=self.num_attention_heads,
+                dim_head=self.attention_head_dim,
+                dropout=self.dropout,
+                bias=self.attention_bias,
+                cross_attention_dim=self.cross_attention_dim,
+                upcast_attention=self.attn2.upcast_attention,
+                out_bias=True,
+            )
+
+        self._initialize_attn_weights()
+
+    def _initialize_attn_weights(self):
+
+        """Initializes specialized attention heads with base weights.
+        
+        Uses weight sharing strategy:
+        - Copies base transformer weights to specialized heads
+        - Initializes newly-added parameters to zero
+        """
+
+        if self.use_mda:
+            for token in self.pbr_setting:
+                if token == "albedo":
+                    continue
+                getattr(self.attn1.processor, f"to_q_{token}").load_state_dict(self.attn1.to_q.state_dict())
+                getattr(self.attn1.processor, f"to_k_{token}").load_state_dict(self.attn1.to_k.state_dict())
+                getattr(self.attn1.processor, f"to_v_{token}").load_state_dict(self.attn1.to_v.state_dict())
+                getattr(self.attn1.processor, f"to_out_{token}").load_state_dict(self.attn1.to_out.state_dict())
+
+        if self.use_ma:
+            self.attn_multiview.load_state_dict(self.attn1.state_dict(), strict=False)
+            with torch.no_grad():
+                for layer in self.attn_multiview.to_out:
+                    for param in layer.parameters():
+                        param.zero_()
+
+        if self.use_ra:
+            self.attn_refview.load_state_dict(self.attn1.state_dict(), strict=False)
+            for token in self.pbr_setting:
+                if token == "albedo":
+                    continue
+                getattr(self.attn_refview.processor, f"to_v_{token}").load_state_dict(
+                    self.attn_refview.to_q.state_dict()
+                )
+                getattr(self.attn_refview.processor, f"to_out_{token}").load_state_dict(
+                    self.attn_refview.to_out.state_dict()
+                )
+            with torch.no_grad():
+                for layer in self.attn_refview.to_out:
+                    for param in layer.parameters():
+                        param.zero_()
+                for token in self.pbr_setting:
+                    if token == "albedo":
+                        continue
+                    for layer in getattr(self.attn_refview.processor, f"to_out_{token}"):
+                        for param in layer.parameters():
+                            param.zero_()
+
+        if self.use_dino:
+            self.attn_dino.load_state_dict(self.attn2.state_dict(), strict=False)
+            with torch.no_grad():
+                for layer in self.attn_dino.to_out:
+                    for param in layer.parameters():
+                        param.zero_()
+
+        if self.use_dino:
+            self.attn_dino.load_state_dict(self.attn2.state_dict(), strict=False)
+            with torch.no_grad():
+                for layer in self.attn_dino.to_out:
+                    for param in layer.parameters():
+                        param.zero_()
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.transformer, name)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+
+        """Forward pass with multi-mechanism attention.
+        
+        Processing stages:
+        1. Material-aware self-attention (MDA)
+        2. Reference attention (RA)
+        3. Multiview attention (MA) with position-aware attention
+        4. Text conditioning (base attention)
+        5. DINO feature conditioning (optional)
+        6. Position-aware conditioning
+        7. Feed-forward network
+        
+        Args:
+            hidden_states: Input features [B * N_materials * N_views, Seq_len, Feat_dim]
+            See base transformer for other parameters
+            
+        Returns:
+            torch.Tensor: Output features
+        """
+        # [Full multi-mechanism processing pipeline...]
+        # Key processing stages:
+        # 1. Material-aware self-attention (handles albedo/mr separation)
+        # 2. Reference attention (conditioned on reference features)
+        # 3. View-to-view attention with geometric constraints
+        # 4. Text-to-image cross-attention
+        # 5. DINO feature fusion (when enabled)
+        # 6. Positional conditioning (RoPE-style)
+        # 7. Feed-forward network with conditional normalization
+
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 0. Self-Attention
+        batch_size = hidden_states.shape[0]
+
+        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)
+        mva_scale = cross_attention_kwargs.pop("mva_scale", 1.0)
+        ref_scale = cross_attention_kwargs.pop("ref_scale", 1.0)
+        condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None)
+        dino_hidden_states = cross_attention_kwargs.pop("dino_hidden_states", None)
+        position_voxel_indices = cross_attention_kwargs.pop("position_voxel_indices", None)
+        N_pbr = len(self.pbr_setting) if self.pbr_setting is not None else 1
+
+        if self.norm_type == "ada_norm":
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.norm_type == "ada_norm_zero":
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]:
+            norm_hidden_states = self.norm1(hidden_states)
+        elif self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif self.norm_type == "ada_norm_single":
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+            ).chunk(6, dim=1)
+            norm_hidden_states = self.norm1(hidden_states)
+            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)
+
+        if self.use_mda:
+            mda_norm_hidden_states = rearrange(
+                norm_hidden_states, "(b n_pbr n) l c -> b n_pbr n l c", n=num_in_batch, n_pbr=N_pbr
+            )
+            attn_output = self.attn1(
+                mda_norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+                attention_mask=attention_mask,
+                **cross_attention_kwargs,
+            )
+            attn_output = rearrange(attn_output, "b n_pbr n l c -> (b n_pbr n) l c")
+        else:
+            attn_output = self.attn1(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+                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":
+            attn_output = gate_msa * attn_output
+
+        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 and self.use_ra:
+            condition_embed = condition_embed_dict[self.layer_name]
+
+            #! Only using albedo features for reference attention
+            ref_norm_hidden_states = rearrange(
+                norm_hidden_states, "(b n_pbr n) l c -> b n_pbr (n l) c", n=num_in_batch, n_pbr=N_pbr
+            )[:, 0, ...]
+
+            attn_output = self.attn_refview(
+                ref_norm_hidden_states,
+                encoder_hidden_states=condition_embed,
+                attention_mask=None,
+                **cross_attention_kwargs,
+            )  # b (n l) c
+            attn_output = rearrange(attn_output, "b n_pbr (n l) c -> (b n_pbr n) l c", n=num_in_batch, n_pbr=N_pbr)
+
+            ref_scale_timing = ref_scale
+            if isinstance(ref_scale, torch.Tensor):
+                ref_scale_timing = ref_scale.unsqueeze(1).repeat(1, num_in_batch * N_pbr).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
+            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_pbr n) l c -> (b n_pbr) (n l) c", n_pbr=N_pbr, n=num_in_batch
+            )
+            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,
+                position_indices=position_indices,
+                n_pbrs=N_pbr,
+                **cross_attention_kwargs,
+            )
+
+            attn_output = rearrange(attn_output, "(b n_pbr) (n l) c -> (b n_pbr n) l c", n_pbr=N_pbr, n=num_in_batch)
+
+            hidden_states = mva_scale * attn_output + hidden_states
+            if hidden_states.ndim == 4:
+                hidden_states = hidden_states.squeeze(1)
+
+        # 1.2 GLIGEN Control
+        if gligen_kwargs is not None:
+            hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
+
+        # 3. Cross-Attention
+        if self.attn2 is not None:
+            if self.norm_type == "ada_norm":
+                norm_hidden_states = self.norm2(hidden_states, timestep)
+            elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]:
+                norm_hidden_states = self.norm2(hidden_states)
+            elif self.norm_type == "ada_norm_single":
+                # For PixArt norm2 isn't applied here:
+                # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
+                norm_hidden_states = hidden_states
+            elif self.norm_type == "ada_norm_continuous":
+                norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"])
+            else:
+                raise ValueError("Incorrect norm")
+
+            if self.pos_embed is not None and self.norm_type != "ada_norm_single":
+                norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+            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
+
+            # dino attn
+            if self.use_dino:
+                dino_hidden_states = dino_hidden_states.unsqueeze(1).repeat(1, N_pbr * num_in_batch, 1, 1)
+                dino_hidden_states = rearrange(dino_hidden_states, "b n l c -> (b n) l c")
+                attn_output = self.attn_dino(
+                    norm_hidden_states,
+                    encoder_hidden_states=dino_hidden_states,
+                    attention_mask=None,
+                    **cross_attention_kwargs,
+                )
+
+                hidden_states = attn_output + hidden_states
+
+        # 4. Feed-forward
+        # i2vgen doesn't have this norm 🤷‍♂️
+        if self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif not self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm3(hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm2(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+        elif self.norm_type == "ada_norm_single":
+            ff_output = gate_mlp * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
+
+
+class ImageProjModel(torch.nn.Module):
+
+    """Projects image embeddings into cross-attention space.
+    
+    Transforms CLIP embeddings into additional context tokens for conditioning.
+    
+    Args:
+        cross_attention_dim: Dimension of attention space
+        clip_embeddings_dim: Dimension of input CLIP embeddings
+        clip_extra_context_tokens: Number of context tokens to generate
+    """
+
+    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)
+
+    def forward(self, image_embeds):
+
+        """Projects image embeddings to cross-attention context tokens.
+        
+        Args:
+            image_embeds: Input embeddings [B, N, C] or [B, C]
+            
+        Returns:
+            torch.Tensor: Context tokens [B, N*clip_extra_context_tokens, cross_attention_dim]
+        """
+
+        embeds = image_embeds
+        num_token = 1
+        if embeds.dim() == 3:
+            num_token = embeds.shape[1]
+            embeds = rearrange(embeds, "b n c -> (b n) c")
+
+        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)
+
+        clip_extra_context_tokens = rearrange(clip_extra_context_tokens, "(b nt) n c -> b (nt n) c", nt=num_token)
+
+        return clip_extra_context_tokens
+
+
+class UNet2p5DConditionModel(torch.nn.Module):
+
+    """2.5D UNet extension for multiview PBR generation.
+    
+    Enhances standard 2D UNet with:
+    - Multiview attention mechanisms
+    - Material-aware processing
+    - Position-aware conditioning
+    - Dual-stream reference processing
+    
+    Args:
+        unet: Base 2D UNet model
+        train_sched: Training scheduler (DDPM)
+        val_sched: Validation scheduler (EulerAncestral)
+    """
+
+    def __init__(
+        self,
+        unet: UNet2DConditionModel,
+        train_sched: DDPMScheduler = None,
+        val_sched: EulerAncestralDiscreteScheduler = None,
+    ) -> None:
+        super().__init__()
+        self.unet = unet
+        self.train_sched = train_sched
+        self.val_sched = val_sched
+
+        self.use_ma = True
+        self.use_ra = True
+        self.use_mda = True
+        self.use_dino = True
+        self.use_position_rope = True
+        self.use_learned_text_clip = True
+        self.use_dual_stream = True
+        self.pbr_setting = ["albedo", "mr"]
+        self.pbr_token_channels = 77
+
+        if self.use_dual_stream and self.use_ra:
+            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,
+            use_dino=self.use_dino,
+            use_mda=self.use_mda,
+            pbr_setting=self.pbr_setting,
+        )
+        self.init_condition(use_dino=self.use_dino)
+
+    @staticmethod
+    def from_pretrained(pretrained_model_name_or_path, **kwargs):
+        torch_dtype = kwargs.pop("torch_dtype", torch.float32)
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        unet_ckpt_path = os.path.join(pretrained_model_name_or_path, "diffusion_pytorch_model.bin")
+        with open(config_path, "r", encoding="utf-8") as file:
+            config = json.load(file)
+        unet = UNet2DConditionModel(**config)
+        unet_2p5d = UNet2p5DConditionModel(unet)
+        unet_2p5d.unet.conv_in = torch.nn.Conv2d(
+            12,
+            unet.conv_in.out_channels,
+            kernel_size=unet.conv_in.kernel_size,
+            stride=unet.conv_in.stride,
+            padding=unet.conv_in.padding,
+            dilation=unet.conv_in.dilation,
+            groups=unet.conv_in.groups,
+            bias=unet.conv_in.bias is not None,
+        )
+        unet_ckpt = torch.load(unet_ckpt_path, map_location="cpu", weights_only=True)
+        unet_2p5d.load_state_dict(unet_ckpt, strict=True)
+        unet_2p5d = unet_2p5d.to(torch_dtype)
+        return unet_2p5d
+
+    def init_condition(self, use_dino):
+
+        """Initializes conditioning mechanisms for multiview PBR generation.
+        
+        Sets up:
+        1. Learned text embeddings: Material-specific tokens (albedo, mr) initialized to zeros
+        2. DINO projector: Model to process DINO-ViT features for cross-attention
+        
+        Args:
+            use_dino: Flag to enable DINO feature integration
+        """
+
+        if self.use_learned_text_clip:
+            for token in self.pbr_setting:
+                self.unet.register_parameter(
+                    f"learned_text_clip_{token}", nn.Parameter(torch.zeros(self.pbr_token_channels, 1024))
+                )
+            self.unet.learned_text_clip_ref = nn.Parameter(torch.zeros(self.pbr_token_channels, 1024))
+
+        if use_dino:
+            self.unet.image_proj_model_dino = ImageProjModel(
+                cross_attention_dim=self.unet.config.cross_attention_dim,
+                clip_embeddings_dim=1536,
+                clip_extra_context_tokens=4,
+            )
+
+    def init_attention(self, unet, use_ma=False, use_ra=False, use_mda=False, use_dino=False, pbr_setting=None):
+
+        """Recursively replaces standard transformers with enhanced 2.5D blocks.
+        
+        Processes UNet architecture:
+        1. Downsampling blocks: Replaces transformers in attention layers
+        2. Middle block: Upgrades central transformers
+        3. Upsampling blocks: Modifies decoder transformers
+        
+        Args:
+            unet: UNet model to enhance
+            use_ma: Enable multiview attention
+            use_ra: Enable reference attention
+            use_mda: Enable material-specific attention
+            use_dino: Enable DINO feature integration
+            pbr_setting: List of PBR materials
+        """
+
+        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_ma,
+                                use_ra,
+                                use_mda,
+                                use_dino,
+                                pbr_setting,
+                            )
+
+        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_ma, use_ra, use_mda, use_dino, pbr_setting
+                        )
+
+        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_ma,
+                                use_ra,
+                                use_mda,
+                                use_dino,
+                                pbr_setting,
+                            )
+
+    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,
+        *args,
+        added_cond_kwargs=None,
+        cross_attention_kwargs=None,
+        down_intrablock_additional_residuals=None,
+        down_block_res_samples=None,
+        mid_block_res_sample=None,
+        **cached_condition,
+    ):
+
+        """Forward pass with multiview/material conditioning.
+        
+        Key stages:
+        1. Input preparation (concat normal/position maps)
+        2. Reference feature extraction (dual-stream)
+        3. Position encoding (voxel indices)
+        4. DINO feature projection
+        5. Main UNet processing with attention conditioning
+        
+        Args:
+            sample: Input latents [B, N_pbr, N_gen, C, H, W]
+            cached_condition: Dictionary containing:
+                - embeds_normal: Normal map embeddings
+                - embeds_position: Position map embeddings
+                - ref_latents: Reference image latents
+                - dino_hidden_states: DINO features
+                - position_maps: 3D position maps
+                - mva_scale: Multiview attention scale
+                - ref_scale: Reference attention scale
+                
+        Returns:
+            torch.Tensor: Output features
+        """
+
+        B, N_pbr, N_gen, _, H, W = sample.shape
+        assert H == W
+
+        if "cache" not in cached_condition:
+            cached_condition["cache"] = {}
+
+        sample = [sample]
+        if "embeds_normal" in cached_condition:
+            sample.append(cached_condition["embeds_normal"].unsqueeze(1).repeat(1, N_pbr, 1, 1, 1, 1))
+        if "embeds_position" in cached_condition:
+            sample.append(cached_condition["embeds_position"].unsqueeze(1).repeat(1, N_pbr, 1, 1, 1, 1))
+        sample = torch.cat(sample, dim=-3)
+
+        sample = rearrange(sample, "b n_pbr n c h w -> (b n_pbr n) c h w")
+
+        encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(-3).repeat(1, 1, N_gen, 1, 1)
+        encoder_hidden_states_gen = rearrange(encoder_hidden_states_gen, "b n_pbr n l c -> (b n_pbr n) l c")
+
+        if added_cond_kwargs is not None:
+            text_embeds_gen = added_cond_kwargs["text_embeds"].unsqueeze(1).repeat(1, N_gen, 1)
+            text_embeds_gen = rearrange(text_embeds_gen, "b n c -> (b n) c")
+            time_ids_gen = added_cond_kwargs["time_ids"].unsqueeze(1).repeat(1, N_gen, 1)
+            time_ids_gen = rearrange(time_ids_gen, "b n c -> (b n) c")
+            added_cond_kwargs_gen = {"text_embeds": text_embeds_gen, "time_ids": time_ids_gen}
+        else:
+            added_cond_kwargs_gen = None
+
+        if self.use_position_rope:
+            if "position_voxel_indices" in cached_condition["cache"]:
+                position_voxel_indices = cached_condition["cache"]["position_voxel_indices"]
+            else:
+                if "position_maps" in cached_condition:
+                    position_voxel_indices = calc_multires_voxel_idxs(
+                        cached_condition["position_maps"],
+                        grid_resolutions=[H, H // 2, H // 4, H // 8],
+                        voxel_resolutions=[H * 8, H * 4, H * 2, H],
+                    )
+                    cached_condition["cache"]["position_voxel_indices"] = position_voxel_indices
+        else:
+            position_voxel_indices = None
+
+        if self.use_dino:
+            if "dino_hidden_states_proj" in cached_condition["cache"]:
+                dino_hidden_states = cached_condition["cache"]["dino_hidden_states_proj"]
+            else:
+                assert "dino_hidden_states" in cached_condition
+                dino_hidden_states = cached_condition["dino_hidden_states"]
+                dino_hidden_states = self.image_proj_model_dino(dino_hidden_states)
+                cached_condition["cache"]["dino_hidden_states_proj"] = dino_hidden_states
+        else:
+            dino_hidden_states = None
+
+        if self.use_ra:
+            if "condition_embed_dict" in cached_condition["cache"]:
+                condition_embed_dict = cached_condition["cache"]["condition_embed_dict"]
+            else:
+                condition_embed_dict = {}
+                ref_latents = cached_condition["ref_latents"]
+                N_ref = ref_latents.shape[1]
+
+                if not self.use_dual_stream:
+                    ref_latents = [ref_latents]
+                    if "embeds_normal" in cached_condition:
+                        ref_latents.append(torch.zeros_like(ref_latents[0]))
+                    if "embeds_position" 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")
+
+                encoder_hidden_states_ref = self.unet.learned_text_clip_ref.repeat(B, N_ref, 1, 1)
+
+                encoder_hidden_states_ref = rearrange(encoder_hidden_states_ref, "b n l c -> (b n) l c")
+
+                if added_cond_kwargs is not None:
+                    text_embeds_ref = added_cond_kwargs["text_embeds"].unsqueeze(1).repeat(1, N_ref, 1)
+                    text_embeds_ref = rearrange(text_embeds_ref, "b n c -> (b n) c")
+                    time_ids_ref = added_cond_kwargs["time_ids"].unsqueeze(1).repeat(1, N_ref, 1)
+                    time_ids_ref = rearrange(time_ids_ref, "b n c -> (b n) c")
+                    added_cond_kwargs_ref = {
+                        "text_embeds": text_embeds_ref,
+                        "time_ids": time_ids_ref,
+                    }
+                else:
+                    added_cond_kwargs_ref = None
+
+                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=None,
+                    added_cond_kwargs=added_cond_kwargs_ref,
+                    # **kwargs
+                    return_dict=False,
+                    cross_attention_kwargs={
+                        "mode": "w",
+                        "num_in_batch": N_ref,
+                        "condition_embed_dict": condition_embed_dict,
+                    },
+                )
+                cached_condition["cache"]["condition_embed_dict"] = condition_embed_dict
+        else:
+            condition_embed_dict = None
+
+        mva_scale = cached_condition.get("mva_scale", 1.0)
+        ref_scale = cached_condition.get("ref_scale", 1.0)
+
+        return self.unet(
+            sample,
+            timestep,
+            encoder_hidden_states_gen,
+            *args,
+            class_labels=None,
+            added_cond_kwargs=added_cond_kwargs_gen,
+            down_intrablock_additional_residuals=(
+                [sample.to(dtype=self.unet.dtype) for sample in down_intrablock_additional_residuals]
+                if down_intrablock_additional_residuals is not None
+                else None
+            ),
+            down_block_additional_residuals=(
+                [sample.to(dtype=self.unet.dtype) for sample in down_block_res_samples]
+                if down_block_res_samples is not None
+                else None
+            ),
+            mid_block_additional_residual=(
+                mid_block_res_sample.to(dtype=self.unet.dtype) if mid_block_res_sample is not None else None
+            ),
+            return_dict=False,
+            cross_attention_kwargs={
+                "mode": "r",
+                "num_in_batch": N_gen,
+                "dino_hidden_states": dino_hidden_states,
+                "condition_embed_dict": condition_embed_dict,
+                "mva_scale": mva_scale,
+                "ref_scale": ref_scale,
+                "position_voxel_indices": position_voxel_indices,
+            },
+        )