hy3dpaint/hunyuanpaintpbr/unet/modules.py

# 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,
            },
        )