# 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 torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Dict, Tuple, Union, Literal, List, Callable
from einops import rearrange
from diffusers.utils import deprecate
from diffusers.models.attention_processor import Attention, AttnProcessor


class AttnUtils:
    """
    Shared utility functions for attention processing.

    This class provides common operations used across different attention processors
    to eliminate code duplication and improve maintainability.
    """

    @staticmethod
    def check_pytorch_compatibility():
        """
        Check PyTorch compatibility for scaled_dot_product_attention.

        Raises:
            ImportError: If PyTorch version doesn't support scaled_dot_product_attention
        """
        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.")

    @staticmethod
    def handle_deprecation_warning(args, kwargs):
        """
        Handle deprecation warning for the 'scale' argument.

        Args:
            args: Positional arguments passed to attention processor
            kwargs: Keyword arguments passed to attention processor
        """
        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)

    @staticmethod
    def prepare_hidden_states(
        hidden_states, attn, temb, spatial_norm_attr="spatial_norm", group_norm_attr="group_norm"
    ):
        """
        Common preprocessing of hidden states for attention computation.

        Args:
            hidden_states: Input hidden states tensor
            attn: Attention module instance
            temb: Optional temporal embedding tensor
            spatial_norm_attr: Attribute name for spatial normalization
            group_norm_attr: Attribute name for group normalization

        Returns:
            Tuple of (processed_hidden_states, residual, input_ndim, shape_info)
        """
        residual = hidden_states

        spatial_norm = getattr(attn, spatial_norm_attr, None)
        if spatial_norm is not None:
            hidden_states = 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)
        else:
            batch_size, channel, height, width = None, None, None, None

        group_norm = getattr(attn, group_norm_attr, None)
        if group_norm is not None:
            hidden_states = group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

        return hidden_states, residual, input_ndim, (batch_size, channel, height, width)

    @staticmethod
    def prepare_attention_mask(attention_mask, attn, sequence_length, batch_size):
        """
        Prepare attention mask for scaled_dot_product_attention.

        Args:
            attention_mask: Input attention mask tensor or None
            attn: Attention module instance
            sequence_length: Length of the sequence
            batch_size: Batch size

        Returns:
            Prepared attention mask tensor reshaped for multi-head attention
        """
        if attention_mask is not None:
            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
        return attention_mask

    @staticmethod
    def reshape_qkv_for_attention(tensor, batch_size, attn_heads, head_dim):
        """
        Reshape Q/K/V tensors for multi-head attention computation.

        Args:
            tensor: Input tensor to reshape
            batch_size: Batch size
            attn_heads: Number of attention heads
            head_dim: Dimension per attention head

        Returns:
            Reshaped tensor with shape [batch_size, attn_heads, seq_len, head_dim]
        """
        return tensor.view(batch_size, -1, attn_heads, head_dim).transpose(1, 2)

    @staticmethod
    def apply_norms(query, key, norm_q, norm_k):
        """
        Apply Q/K normalization layers if available.

        Args:
            query: Query tensor
            key: Key tensor
            norm_q: Query normalization layer (optional)
            norm_k: Key normalization layer (optional)

        Returns:
            Tuple of (normalized_query, normalized_key)
        """
        if norm_q is not None:
            query = norm_q(query)
        if norm_k is not None:
            key = norm_k(key)
        return query, key

    @staticmethod
    def finalize_output(hidden_states, input_ndim, shape_info, attn, residual, to_out):
        """
        Common output processing including projection, dropout, reshaping, and residual connection.

        Args:
            hidden_states: Processed hidden states from attention
            input_ndim: Original input tensor dimensions
            shape_info: Tuple containing original shape information
            attn: Attention module instance
            residual: Residual connection tensor
            to_out: Output projection layers [linear, dropout]

        Returns:
            Final output tensor after all processing steps
        """
        batch_size, channel, height, width = shape_info

        # Apply output projection and dropout
        hidden_states = to_out[0](hidden_states)
        hidden_states = to_out[1](hidden_states)

        # Reshape back if needed
        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)

        # Apply residual connection
        if attn.residual_connection:
            hidden_states = hidden_states + residual

        # Apply rescaling
        hidden_states = hidden_states / attn.rescale_output_factor
        return hidden_states


# Base class for attention processors (eliminating initialization duplication)
class BaseAttnProcessor(nn.Module):
    """
    Base class for attention processors with common initialization.

    This base class provides shared parameter initialization and module registration
    functionality to reduce code duplication across different attention processor types.
    """

    def __init__(
        self,
        query_dim: int,
        pbr_setting: List[str] = ["albedo", "mr"],
        cross_attention_dim: Optional[int] = None,
        heads: int = 8,
        kv_heads: Optional[int] = None,
        dim_head: int = 64,
        dropout: float = 0.0,
        bias: bool = False,
        upcast_attention: bool = False,
        upcast_softmax: bool = False,
        cross_attention_norm: Optional[str] = None,
        cross_attention_norm_num_groups: int = 32,
        qk_norm: Optional[str] = None,
        added_kv_proj_dim: Optional[int] = None,
        added_proj_bias: Optional[bool] = True,
        norm_num_groups: Optional[int] = None,
        spatial_norm_dim: Optional[int] = None,
        out_bias: bool = True,
        scale_qk: bool = True,
        only_cross_attention: bool = False,
        eps: float = 1e-5,
        rescale_output_factor: float = 1.0,
        residual_connection: bool = False,
        _from_deprecated_attn_block: bool = False,
        processor: Optional["AttnProcessor"] = None,
        out_dim: int = None,
        out_context_dim: int = None,
        context_pre_only=None,
        pre_only=False,
        elementwise_affine: bool = True,
        is_causal: bool = False,
        **kwargs,
    ):
        """
        Initialize base attention processor with common parameters.

        Args:
            query_dim: Dimension of query features
            pbr_setting: List of PBR material types to process (e.g., ["albedo", "mr"])
            cross_attention_dim: Dimension of cross-attention features (optional)
            heads: Number of attention heads
            kv_heads: Number of key-value heads for grouped query attention (optional)
            dim_head: Dimension per attention head
            dropout: Dropout rate
            bias: Whether to use bias in linear projections
            upcast_attention: Whether to upcast attention computation to float32
            upcast_softmax: Whether to upcast softmax computation to float32
            cross_attention_norm: Type of cross-attention normalization (optional)
            cross_attention_norm_num_groups: Number of groups for cross-attention norm
            qk_norm: Type of query-key normalization (optional)
            added_kv_proj_dim: Dimension for additional key-value projections (optional)
            added_proj_bias: Whether to use bias in additional projections
            norm_num_groups: Number of groups for normalization (optional)
            spatial_norm_dim: Dimension for spatial normalization (optional)
            out_bias: Whether to use bias in output projection
            scale_qk: Whether to scale query-key products
            only_cross_attention: Whether to only perform cross-attention
            eps: Small epsilon value for numerical stability
            rescale_output_factor: Factor to rescale output values
            residual_connection: Whether to use residual connections
            _from_deprecated_attn_block: Flag for deprecated attention blocks
            processor: Optional attention processor instance
            out_dim: Output dimension (optional)
            out_context_dim: Output context dimension (optional)
            context_pre_only: Whether to only process context in pre-processing
            pre_only: Whether to only perform pre-processing
            elementwise_affine: Whether to use element-wise affine transformations
            is_causal: Whether to use causal attention masking
            **kwargs: Additional keyword arguments
        """
        super().__init__()
        AttnUtils.check_pytorch_compatibility()

        # Store common attributes
        self.pbr_setting = pbr_setting
        self.n_pbr_tokens = len(self.pbr_setting)
        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
        self.query_dim = query_dim
        self.use_bias = bias
        self.is_cross_attention = cross_attention_dim is not None
        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
        self.upcast_attention = upcast_attention
        self.upcast_softmax = upcast_softmax
        self.rescale_output_factor = rescale_output_factor
        self.residual_connection = residual_connection
        self.dropout = dropout
        self.fused_projections = False
        self.out_dim = out_dim if out_dim is not None else query_dim
        self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
        self.context_pre_only = context_pre_only
        self.pre_only = pre_only
        self.is_causal = is_causal
        self._from_deprecated_attn_block = _from_deprecated_attn_block
        self.scale_qk = scale_qk
        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
        self.heads = out_dim // dim_head if out_dim is not None else heads
        self.sliceable_head_dim = heads
        self.added_kv_proj_dim = added_kv_proj_dim
        self.only_cross_attention = only_cross_attention
        self.added_proj_bias = added_proj_bias

        # Validation
        if self.added_kv_proj_dim is None and self.only_cross_attention:
            raise ValueError(
                "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None."
                "Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
            )

    def register_pbr_modules(self, module_types: List[str], **kwargs):
        """
        Generic PBR module registration to eliminate code repetition.

        Dynamically registers PyTorch modules for different PBR material types
        based on the specified module types and PBR settings.

        Args:
            module_types: List of module types to register ("qkv", "v_only", "out", "add_kv")
            **kwargs: Additional arguments for module configuration
        """
        for pbr_token in self.pbr_setting:
            if pbr_token == "albedo":
                continue

            for module_type in module_types:
                if module_type == "qkv":
                    self.register_module(
                        f"to_q_{pbr_token}", nn.Linear(self.query_dim, self.inner_dim, bias=self.use_bias)
                    )
                    self.register_module(
                        f"to_k_{pbr_token}", nn.Linear(self.cross_attention_dim, self.inner_dim, bias=self.use_bias)
                    )
                    self.register_module(
                        f"to_v_{pbr_token}", nn.Linear(self.cross_attention_dim, self.inner_dim, bias=self.use_bias)
                    )
                elif module_type == "v_only":
                    self.register_module(
                        f"to_v_{pbr_token}", nn.Linear(self.cross_attention_dim, self.inner_dim, bias=self.use_bias)
                    )
                elif module_type == "out":
                    if not self.pre_only:
                        self.register_module(
                            f"to_out_{pbr_token}",
                            nn.ModuleList(
                                [
                                    nn.Linear(self.inner_dim, self.out_dim, bias=kwargs.get("out_bias", True)),
                                    nn.Dropout(self.dropout),
                                ]
                            ),
                        )
                    else:
                        self.register_module(f"to_out_{pbr_token}", None)
                elif module_type == "add_kv":
                    if self.added_kv_proj_dim is not None:
                        self.register_module(
                            f"add_k_proj_{pbr_token}",
                            nn.Linear(self.added_kv_proj_dim, self.inner_kv_dim, bias=self.added_proj_bias),
                        )
                        self.register_module(
                            f"add_v_proj_{pbr_token}",
                            nn.Linear(self.added_kv_proj_dim, self.inner_kv_dim, bias=self.added_proj_bias),
                        )
                    else:
                        self.register_module(f"add_k_proj_{pbr_token}", None)
                        self.register_module(f"add_v_proj_{pbr_token}", None)


# Rotary Position Embedding utilities (specialized for PoseRoPE)
class RotaryEmbedding:
    """
    Rotary position embedding utilities for 3D spatial attention.

    Provides functions to compute and apply rotary position embeddings (RoPE)
    for 1D, 3D spatial coordinates used in 3D-aware attention mechanisms.
    """

    @staticmethod
    def get_1d_rotary_pos_embed(dim: int, pos: torch.Tensor, theta: float = 10000.0, linear_factor=1.0, ntk_factor=1.0):
        """
        Compute 1D rotary position embeddings.

        Args:
            dim: Embedding dimension (must be even)
            pos: Position tensor
            theta: Base frequency for rotary embeddings
            linear_factor: Linear scaling factor
            ntk_factor: NTK (Neural Tangent Kernel) scaling factor

        Returns:
            Tuple of (cos_embeddings, sin_embeddings)
        """
        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
        )
        freqs = torch.outer(pos, freqs)
        freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()
        freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()
        return freqs_cos, freqs_sin

    @staticmethod
    def get_3d_rotary_pos_embed(position, embed_dim, voxel_resolution, theta: int = 10000):
        """
        Compute 3D rotary position embeddings for spatial coordinates.

        Args:
            position: 3D position tensor with shape [..., 3]
            embed_dim: Embedding dimension
            voxel_resolution: Resolution of the voxel grid
            theta: Base frequency for rotary embeddings

        Returns:
            Tuple of (cos_embeddings, sin_embeddings) for 3D positions
        """
        assert position.shape[-1] == 3
        dim_xy = embed_dim // 8 * 3
        dim_z = embed_dim // 8 * 2

        grid = torch.arange(voxel_resolution, dtype=torch.float32, device=position.device)
        freqs_xy = RotaryEmbedding.get_1d_rotary_pos_embed(dim_xy, grid, theta=theta)
        freqs_z = RotaryEmbedding.get_1d_rotary_pos_embed(dim_z, grid, theta=theta)

        xy_cos, xy_sin = freqs_xy
        z_cos, z_sin = freqs_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

    @staticmethod
    def apply_rotary_emb(x: torch.Tensor, freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]]):
        """
        Apply rotary position embeddings to input tensor.

        Args:
            x: Input tensor to apply rotary embeddings to
            freqs_cis: Tuple of (cos_embeddings, sin_embeddings) or single tensor

        Returns:
            Tensor with rotary position embeddings applied
        """
        cos, sin = freqs_cis
        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)
        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


# Core attention processing logic (eliminating major duplication)
class AttnCore:
    """
    Core attention processing logic shared across processors.

    This class provides the fundamental attention computation pipeline
    that can be reused across different attention processor implementations.
    """

    @staticmethod
    def process_attention_base(
        attn: Attention,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        temb: Optional[torch.Tensor] = None,
        get_qkv_fn: Callable = None,
        apply_rope_fn: Optional[Callable] = None,
        **kwargs,
    ):
        """
        Generic attention processing core shared across different processors.

        This function implements the common attention computation pipeline including:
        1. Hidden state preprocessing
        2. Attention mask preparation
        3. Q/K/V computation via provided function
        4. Tensor reshaping for multi-head attention
        5. Optional normalization and RoPE application
        6. Scaled dot-product attention computation

        Args:
            attn: Attention module instance
            hidden_states: Input hidden states tensor
            encoder_hidden_states: Optional encoder hidden states for cross-attention
            attention_mask: Optional attention mask tensor
            temb: Optional temporal embedding tensor
            get_qkv_fn: Function to compute Q, K, V tensors
            apply_rope_fn: Optional function to apply rotary position embeddings
            **kwargs: Additional keyword arguments passed to subfunctions

        Returns:
            Tuple containing (attention_output, residual, input_ndim, shape_info,
            batch_size, num_heads, head_dim)
        """
        # Prepare hidden states
        hidden_states, residual, input_ndim, shape_info = AttnUtils.prepare_hidden_states(hidden_states, attn, temb)

        batch_size, sequence_length, _ = (
            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
        )

        # Prepare attention mask
        attention_mask = AttnUtils.prepare_attention_mask(attention_mask, attn, sequence_length, batch_size)

        # Get Q, K, V
        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)

        query, key, value = get_qkv_fn(attn, hidden_states, encoder_hidden_states, **kwargs)

        # Reshape for attention
        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads

        query = AttnUtils.reshape_qkv_for_attention(query, batch_size, attn.heads, head_dim)
        key = AttnUtils.reshape_qkv_for_attention(key, batch_size, attn.heads, head_dim)
        value = AttnUtils.reshape_qkv_for_attention(value, batch_size, attn.heads, value.shape[-1] // attn.heads)

        # Apply normalization
        query, key = AttnUtils.apply_norms(query, key, getattr(attn, "norm_q", None), getattr(attn, "norm_k", None))

        # Apply RoPE if provided
        if apply_rope_fn is not None:
            query, key = apply_rope_fn(query, key, head_dim, **kwargs)

        # Compute attention
        hidden_states = F.scaled_dot_product_attention(
            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
        )

        return hidden_states, residual, input_ndim, shape_info, batch_size, attn.heads, head_dim


# Specific processor implementations (minimal unique code)
class PoseRoPEAttnProcessor2_0:
    """
    Attention processor with Rotary Position Encoding (RoPE) for 3D spatial awareness.

    This processor extends standard attention with 3D rotary position embeddings
    to provide spatial awareness for 3D scene understanding tasks.
    """

    def __init__(self):
        """Initialize the RoPE attention processor."""
        AttnUtils.check_pytorch_compatibility()

    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,
        n_pbrs=1,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        """
        Apply RoPE-enhanced attention computation.

        Args:
            attn: Attention module instance
            hidden_states: Input hidden states tensor
            encoder_hidden_states: Optional encoder hidden states for cross-attention
            attention_mask: Optional attention mask tensor
            position_indices: Dictionary containing 3D position information for RoPE
            temb: Optional temporal embedding tensor
            n_pbrs: Number of PBR material types
            *args: Additional positional arguments
            **kwargs: Additional keyword arguments

        Returns:
            Attention output tensor with applied rotary position encodings
        """
        AttnUtils.handle_deprecation_warning(args, kwargs)

        def get_qkv(attn, hidden_states, encoder_hidden_states, **kwargs):
            return attn.to_q(hidden_states), attn.to_k(encoder_hidden_states), attn.to_v(encoder_hidden_states)

        def apply_rope(query, key, head_dim, **kwargs):
            if position_indices is not None:
                if head_dim in position_indices:
                    image_rotary_emb = position_indices[head_dim]
                else:
                    image_rotary_emb = RotaryEmbedding.get_3d_rotary_pos_embed(
                        rearrange(
                            position_indices["voxel_indices"].unsqueeze(1).repeat(1, n_pbrs, 1, 1),
                            "b n_pbrs l c -> (b n_pbrs) l c",
                        ),
                        head_dim,
                        voxel_resolution=position_indices["voxel_resolution"],
                    )
                    position_indices[head_dim] = image_rotary_emb

                query = RotaryEmbedding.apply_rotary_emb(query, image_rotary_emb)
                key = RotaryEmbedding.apply_rotary_emb(key, image_rotary_emb)
            return query, key

        # Core attention processing
        hidden_states, residual, input_ndim, shape_info, batch_size, heads, head_dim = AttnCore.process_attention_base(
            attn,
            hidden_states,
            encoder_hidden_states,
            attention_mask,
            temb,
            get_qkv_fn=get_qkv,
            apply_rope_fn=apply_rope,
            position_indices=position_indices,
            n_pbrs=n_pbrs,
        )

        # Finalize output
        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, heads * head_dim)
        hidden_states = hidden_states.to(hidden_states.dtype)

        return AttnUtils.finalize_output(hidden_states, input_ndim, shape_info, attn, residual, attn.to_out)


class SelfAttnProcessor2_0(BaseAttnProcessor):
    """
    Self-attention processor with PBR (Physically Based Rendering) material support.

    This processor handles multiple PBR material types (e.g., albedo, metallic-roughness)
    with separate attention computation paths for each material type.
    """

    def __init__(self, **kwargs):
        """
        Initialize self-attention processor with PBR support.

        Args:
            **kwargs: Arguments passed to BaseAttnProcessor initialization
        """
        super().__init__(**kwargs)
        self.register_pbr_modules(["qkv", "out", "add_kv"], **kwargs)

    def process_single(
        self,
        attn: Attention,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        temb: Optional[torch.Tensor] = None,
        token: Literal["albedo", "mr"] = "albedo",
        multiple_devices=False,
        *args,
        **kwargs,
    ):
        """
        Process attention for a single PBR material type.

        Args:
            attn: Attention module instance
            hidden_states: Input hidden states tensor
            encoder_hidden_states: Optional encoder hidden states for cross-attention
            attention_mask: Optional attention mask tensor
            temb: Optional temporal embedding tensor
            token: PBR material type to process ("albedo", "mr", etc.)
            multiple_devices: Whether to use multiple GPU devices
            *args: Additional positional arguments
            **kwargs: Additional keyword arguments

        Returns:
            Processed attention output for the specified PBR material type
        """
        target = attn if token == "albedo" else attn.processor
        token_suffix = "" if token == "albedo" else "_" + token

        # Device management (if needed)
        if multiple_devices:
            device = torch.device("cuda:0") if token == "albedo" else torch.device("cuda:1")
            for attr in [f"to_q{token_suffix}", f"to_k{token_suffix}", f"to_v{token_suffix}", f"to_out{token_suffix}"]:
                getattr(target, attr).to(device)

        def get_qkv(attn, hidden_states, encoder_hidden_states, **kwargs):
            return (
                getattr(target, f"to_q{token_suffix}")(hidden_states),
                getattr(target, f"to_k{token_suffix}")(encoder_hidden_states),
                getattr(target, f"to_v{token_suffix}")(encoder_hidden_states),
            )

        # Core processing using shared logic
        hidden_states, residual, input_ndim, shape_info, batch_size, heads, head_dim = AttnCore.process_attention_base(
            attn, hidden_states, encoder_hidden_states, attention_mask, temb, get_qkv_fn=get_qkv
        )

        # Finalize
        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, heads * head_dim)
        hidden_states = hidden_states.to(hidden_states.dtype)

        return AttnUtils.finalize_output(
            hidden_states, input_ndim, shape_info, attn, residual, getattr(target, f"to_out{token_suffix}")
        )

    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        temb: Optional[torch.Tensor] = None,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        """
        Apply self-attention with PBR material processing.

        Processes multiple PBR material types sequentially, applying attention
        computation for each material type separately and combining results.

        Args:
            attn: Attention module instance
            hidden_states: Input hidden states tensor with PBR dimension
            encoder_hidden_states: Optional encoder hidden states for cross-attention
            attention_mask: Optional attention mask tensor
            temb: Optional temporal embedding tensor
            *args: Additional positional arguments
            **kwargs: Additional keyword arguments

        Returns:
            Combined attention output for all PBR material types
        """
        AttnUtils.handle_deprecation_warning(args, kwargs)

        B = hidden_states.size(0)
        pbr_hidden_states = torch.split(hidden_states, 1, dim=1)

        # Process each PBR setting
        results = []
        for token, pbr_hs in zip(self.pbr_setting, pbr_hidden_states):
            processed_hs = rearrange(pbr_hs, "b n_pbrs n l c -> (b n_pbrs n) l c").to("cuda:0")
            result = self.process_single(attn, processed_hs, None, attention_mask, temb, token, False)
            results.append(result)

        outputs = [rearrange(result, "(b n_pbrs n) l c -> b n_pbrs n l c", b=B, n_pbrs=1) for result in results]
        return torch.cat(outputs, dim=1)


class RefAttnProcessor2_0(BaseAttnProcessor):
    """
    Reference attention processor with shared value computation across PBR materials.

    This processor computes query and key once, but uses separate value projections
    for different PBR material types, enabling efficient multi-material processing.
    """

    def __init__(self, **kwargs):
        """
        Initialize reference attention processor.

        Args:
            **kwargs: Arguments passed to BaseAttnProcessor initialization
        """
        super().__init__(**kwargs)
        self.pbr_settings = self.pbr_setting  # Alias for compatibility
        self.register_pbr_modules(["v_only", "out"], **kwargs)

    def __call__(
        self,
        attn: Attention,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        temb: Optional[torch.Tensor] = None,
        *args,
        **kwargs,
    ) -> torch.Tensor:
        """
        Apply reference attention with shared Q/K and separate V projections.

        This method computes query and key tensors once and reuses them across
        all PBR material types, while using separate value projections for each
        material type to maintain material-specific information.

        Args:
            attn: Attention module instance
            hidden_states: Input hidden states tensor
            encoder_hidden_states: Optional encoder hidden states for cross-attention
            attention_mask: Optional attention mask tensor
            temb: Optional temporal embedding tensor
            *args: Additional positional arguments
            **kwargs: Additional keyword arguments

        Returns:
            Stacked attention output for all PBR material types
        """
        AttnUtils.handle_deprecation_warning(args, kwargs)

        def get_qkv(attn, hidden_states, encoder_hidden_states, **kwargs):
            query = attn.to_q(hidden_states)
            key = attn.to_k(encoder_hidden_states)

            # Concatenate values from all PBR settings
            value_list = [attn.to_v(encoder_hidden_states)]
            for token in ["_" + token for token in self.pbr_settings if token != "albedo"]:
                value_list.append(getattr(attn.processor, f"to_v{token}")(encoder_hidden_states))
            value = torch.cat(value_list, dim=-1)

            return query, key, value

        # Core processing
        hidden_states, residual, input_ndim, shape_info, batch_size, heads, head_dim = AttnCore.process_attention_base(
            attn, hidden_states, encoder_hidden_states, attention_mask, temb, get_qkv_fn=get_qkv
        )

        # Split and process each PBR setting output
        hidden_states_list = torch.split(hidden_states, head_dim, dim=-1)
        output_hidden_states_list = []

        for i, hs in enumerate(hidden_states_list):
            hs = hs.transpose(1, 2).reshape(batch_size, -1, heads * head_dim).to(hs.dtype)
            token_suffix = "_" + self.pbr_settings[i] if self.pbr_settings[i] != "albedo" else ""
            target = attn if self.pbr_settings[i] == "albedo" else attn.processor

            hs = AttnUtils.finalize_output(
                hs, input_ndim, shape_info, attn, residual, getattr(target, f"to_out{token_suffix}")
            )
            output_hidden_states_list.append(hs)

        return torch.stack(output_hidden_states_list, dim=1)