Upload 20 files

hunyuan3d-paint-v2-0-turbo/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

hunyuan3d-paint-v2-0-turbo/README.md

@@ -0,0 +1,53 @@
+---
+license: openrail++
+tags:
+- stable-diffusion
+- text-to-image
+---
+
+# SD v2.1-base with Zero Terminal SNR (LAION Aesthetic 6+)
+
+This model is used in [Diffusion Model with Perceptual Loss](https://arxiv.org/abs/2401.00110) paper as the MSE baseline.
+
+This model is trained using zero terminal SNR schedule following [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/abs/2305.08891) paper on LAION aesthetic 6+ data.
+
+This model is finetuned from [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base).
+
+This model is meant for research demonstration, not for production use.
+
+## Usage
+
+```python
+from diffusers import StableDiffusionPipeline
+prompt = "A young girl smiling"
+pipe = StableDiffusionPipeline.from_pretrained("ByteDance/sd2.1-base-zsnr-laionaes6").to("cuda")
+pipe(prompt, guidance_scale=7.5, guidance_rescale=0.7).images[0].save("out.jpg")
+```
+
+## Related Models
+
+* [bytedance/sd2.1-base-zsnr-laionaes5](https://huggingface.co/ByteDance/sd2.1-base-zsnr-laionaes5)
+* [bytedance/sd2.1-base-zsnr-laionaes6](https://huggingface.co/ByteDance/sd2.1-base-zsnr-laionaes6)
+* [bytedance/sd2.1-base-zsnr-laionaes6-perceptual](https://huggingface.co/ByteDance/sd2.1-base-zsnr-laionaes6-perceptual)
+
+
+## Cite as
+```
+@misc{lin2024diffusion,
+      title={Diffusion Model with Perceptual Loss}, 
+      author={Shanchuan Lin and Xiao Yang},
+      year={2024},
+      eprint={2401.00110},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+
+@misc{lin2023common,
+      title={Common Diffusion Noise Schedules and Sample Steps are Flawed}, 
+      author={Shanchuan Lin and Bingchen Liu and Jiashi Li and Xiao Yang},
+      year={2023},
+      eprint={2305.08891},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

hunyuan3d-paint-v2-0-turbo/feature_extractor/preprocessor_config.json

@@ -0,0 +1,20 @@
+{
+  "crop_size": 224,
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_resize": true,
+  "feature_extractor_type": "CLIPFeatureExtractor",
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "size": 224
+}

hunyuan3d-paint-v2-0-turbo/image_encoder/config.json

@@ -0,0 +1,23 @@
+{
+  "_name_or_path": "D:\\.cache\\huggingface\\hub\\models--sudo-ai--zero123plus-v1.1\\snapshots\\36df7de980afd15f80b2e1a4e9a920d7020e2654\\vision_encoder",
+  "architectures": [
+    "CLIPVisionModelWithProjection"
+  ],
+  "attention_dropout": 0.0,
+  "dropout": 0.0,
+  "hidden_act": "gelu",
+  "hidden_size": 1280,
+  "image_size": 224,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "layer_norm_eps": 1e-05,
+  "model_type": "clip_vision_model",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 32,
+  "patch_size": 14,
+  "projection_dim": 1024,
+  "torch_dtype": "float16",
+  "transformers_version": "4.36.0"
+}

hunyuan3d-paint-v2-0-turbo/image_encoder/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae616c24393dd1854372b0639e5541666f7521cbe219669255e865cb7f89466a
+size 1264217240

hunyuan3d-paint-v2-0-turbo/image_encoder/preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 224
+  }
+}

hunyuan3d-paint-v2-0-turbo/model_index.json

@@ -0,0 +1,37 @@
+{
+  "_class_name": "StableDiffusionPipeline",
+  "_diffusers_version": "0.23.1",
+  "feature_extractor": [
+    "transformers",
+    "CLIPImageProcessor"
+  ],
+  "requires_safety_checker": false,
+  "safety_checker": [
+    null,
+    null
+  ],
+  "scheduler": [
+    "diffusers",
+    "DDIMScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "image_encoder": [
+    "transformers",
+    "CLIPVisionModelWithProjection"
+  ],
+  "unet": [
+    "modules",
+    "UNet2p5DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

hunyuan3d-paint-v2-0-turbo/scheduler/scheduler_config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.23.1",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "num_train_timesteps": 1000,
+  "prediction_type": "v_prediction",
+  "set_alpha_to_one": true,
+  "steps_offset": 1,
+  "trained_betas": null,
+  "timestep_spacing": "trailing",
+  "rescale_betas_zero_snr": true
+}

hunyuan3d-paint-v2-0-turbo/text_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "stabilityai/stable-diffusion-2",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_size": 1024,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 23,
+  "pad_token_id": 1,
+  "projection_dim": 512,
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.0.dev0",
+  "vocab_size": 49408
+}

hunyuan3d-paint-v2-0-turbo/text_encoder/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3e254d7b61353497ea0be2c4013df4ea8f739ee88cffa0ba58cd085459ed565
+size 1361671895

hunyuan3d-paint-v2-0-turbo/tokenizer/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "!",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

hunyuan3d-paint-v2-0-turbo/tokenizer/tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+  "add_prefix_space": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "do_lower_case": true,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "errors": "replace",
+  "model_max_length": 77,
+  "name_or_path": "stabilityai/stable-diffusion-2",
+  "pad_token": "<|endoftext|>",
+  "special_tokens_map_file": "./special_tokens_map.json",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

hunyuan3d-paint-v2-0-turbo/unet/config.json

@@ -0,0 +1,45 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.10.0.dev0",
+  "act_fn": "silu",
+  "attention_head_dim": [
+    5,
+    10,
+    20,
+    20
+  ],
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "cross_attention_dim": 1024,
+  "down_block_types": [
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dual_cross_attention": false,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "sample_size": 64,
+  "up_block_types": [
+    "UpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D"
+  ],
+  "use_linear_projection": true
+}

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

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

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

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6acffa4a22f4da61d87f446bfa83e7ac245481c1535fbf25b200fe4462d0b22
+size 3722161032

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

@@ -0,0 +1,926 @@
+# Open Source Model Licensed under the Apache License Version 2.0
+# and Other Licenses of the Third-Party Components therein:
+# The below Model in this distribution may have been modified by THL A29 Limited
+# ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2024 THL A29 Limited.
+
+# Copyright (C) 2024 THL A29 Limited, a Tencent company.  All rights reserved.
+# The below software and/or models in this distribution may have been
+# modified by THL A29 Limited ("Tencent Modifications").
+# All Tencent Modifications are Copyright (C) THL A29 Limited.
+
+# 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 copy
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models import UNet2DConditionModel
+from diffusers.models.attention_processor import Attention
+from diffusers.models.transformers.transformer_2d import BasicTransformerBlock
+from einops import rearrange
+
+
+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
+    if hidden_states.shape[chunk_dim] % chunk_size != 0:
+        raise ValueError(
+            f"`hidden_states` dimension to be chunked: {hidden_states.shape[chunk_dim]} has to be divisible by chunk size: {chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+        )
+
+    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
+
+class PoseRoPEAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+    def get_1d_rotary_pos_embed(
+            self,
+            dim: int,
+            pos: torch.Tensor,
+            theta: float = 10000.0,
+            linear_factor=1.0,
+            ntk_factor=1.0,
+    ):
+        assert dim % 2 == 0
+
+        theta = theta * ntk_factor
+        freqs = (
+            1.0
+            / (theta ** (torch.arange(0, dim, 2, dtype=pos.dtype, device=pos.device)[: (dim // 2)] / dim))
+            / linear_factor
+        )  # [D/2]
+        freqs = torch.outer(pos, freqs)  # type: ignore   # [S, D/2]
+        # flux, hunyuan-dit, cogvideox
+        freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float()  # [S, D]
+        freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float()  # [S, D]
+        return freqs_cos, freqs_sin
+
+
+    def get_3d_rotary_pos_embed(
+            self,
+            position,
+            embed_dim, 
+            voxel_resolution,
+            theta: int = 10000,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        RoPE for video tokens with 3D structure.
+
+        Args:
+        voxel_resolution (`int`):
+            The grid size of the spatial positional embedding (height, width).
+        theta (`float`):
+            Scaling factor for frequency computation.
+
+        Returns:
+            `torch.Tensor`: positional embedding with shape `(temporal_size * grid_size[0] * grid_size[1], embed_dim/2)`.
+        """
+        assert position.shape[-1]==3
+
+        # Compute dimensions for each axis
+        dim_xy = embed_dim // 8 * 3
+        dim_z = embed_dim // 8 * 2
+
+        # Temporal frequencies
+        grid = torch.arange(voxel_resolution, dtype=torch.float32, device=position.device)
+        freqs_xy = self.get_1d_rotary_pos_embed(dim_xy, grid, theta=theta)
+        freqs_z = self.get_1d_rotary_pos_embed(dim_z, grid, theta=theta)
+
+        xy_cos, xy_sin = freqs_xy  # both t_cos and t_sin has shape: voxel_resolution, dim_xy
+        z_cos, z_sin = freqs_z  # both w_cos and w_sin has shape: voxel_resolution, dim_z
+
+        embed_flattn = position.view(-1, position.shape[-1])
+        x_cos = xy_cos[embed_flattn[:,0], :]
+        x_sin = xy_sin[embed_flattn[:,0], :]
+        y_cos = xy_cos[embed_flattn[:,1], :]
+        y_sin = xy_sin[embed_flattn[:,1], :]
+        z_cos = z_cos[embed_flattn[:,2], :]
+        z_sin = z_sin[embed_flattn[:,2], :]
+
+        cos = torch.cat((x_cos, y_cos, z_cos), dim=-1)
+        sin = torch.cat((x_sin, y_sin, z_sin), dim=-1)
+
+        cos = cos.view(*position.shape[:-1], embed_dim)
+        sin = sin.view(*position.shape[:-1], embed_dim)
+        return cos, sin
+
+    def apply_rotary_emb(
+            self, 
+            x: torch.Tensor,
+            freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]]
+        ):
+        cos, sin = freqs_cis  # [S, D]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+
+        x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+        x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_indices: Dict = None,
+        temb: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+        
+        if position_indices is not None:
+            if head_dim in position_indices:
+                image_rotary_emb = position_indices[head_dim]
+            else:
+                image_rotary_emb = self.get_3d_rotary_pos_embed(position_indices['voxel_indices'], head_dim, voxel_resolution=position_indices['voxel_resolution'])
+                position_indices[head_dim] = image_rotary_emb
+            query = self.apply_rotary_emb(query, image_rotary_emb)
+            key = self.apply_rotary_emb(key, image_rotary_emb)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+class IPAttnProcessor2_0:
+    r"""
+    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+    """
+
+    def __init__(self, scale=0.0):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.scale = scale
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        ip_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        if len(args) > 0 or kwargs.get("scale", None) is not None:
+            deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+            deprecate("scale", "1.0.0", deprecation_message)
+
+        residual = hidden_states
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+        
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # for ip adapter
+        if ip_hidden_states is not None:
+
+            ip_key = attn.to_k_ip(ip_hidden_states)
+            ip_value = attn.to_v_ip(ip_hidden_states)
+
+            ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+            ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+            # the output of sdp = (batch, num_heads, seq_len, head_dim)
+            ip_hidden_states = F.scaled_dot_product_attention(
+                query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+            )
+
+            ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+            ip_hidden_states = ip_hidden_states.to(query.dtype)
+
+            hidden_states = hidden_states + self.scale * ip_hidden_states
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+        
+
+class Basic2p5DTransformerBlock(torch.nn.Module):
+    def __init__(self, transformer: BasicTransformerBlock, layer_name, use_ipa=True, use_ma=True, use_ra=True) -> None:
+        super().__init__()
+        self.transformer = transformer
+        self.layer_name = layer_name
+        self.use_ipa = use_ipa
+        self.use_ma = use_ma
+        self.use_ra = use_ra
+
+        if use_ipa:
+            self.attn2.set_processor(IPAttnProcessor2_0())
+            cross_attention_dim = 1024
+            self.attn2.to_k_ip = nn.Linear(cross_attention_dim, self.dim, bias=False)
+            self.attn2.to_v_ip = nn.Linear(cross_attention_dim, self.dim, bias=False)
+            
+        # multiview attn
+        if self.use_ma:
+            self.attn_multiview = Attention(
+                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,
+            )
+
+        self._initialize_attn_weights()
+
+    def _initialize_attn_weights(self):
+
+        if self.use_ma:
+            self.attn_multiview.load_state_dict(self.attn1.state_dict()) 
+            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()) 
+            with torch.no_grad():
+                for layer in self.attn_refview.to_out:
+                    for param in layer.parameters():
+                        param.zero_()
+
+        if self.use_ipa:
+            self.attn2.to_k_ip.load_state_dict(self.attn2.to_k.state_dict()) 
+            self.attn2.to_v_ip.load_state_dict(self.attn2.to_v.state_dict()) 
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        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:
+
+        # 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)
+        condition_embed_dict = cross_attention_kwargs.pop("condition_embed_dict", None)
+        ip_hidden_states = cross_attention_kwargs.pop("ip_hidden_states", None)
+        position_attn_mask = cross_attention_kwargs.pop("position_attn_mask", None)
+        position_voxel_indices = cross_attention_kwargs.pop("position_voxel_indices", None)
+
+        if self.norm_type == "ada_norm":
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        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)
+
+        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:
+            condition_embed = condition_embed_dict[self.layer_name].unsqueeze(1).repeat(1,num_in_batch,1,1) # B N L C
+            condition_embed = rearrange(condition_embed, 'b n l c -> (b n) l c')
+
+            attn_output = self.attn_refview(
+                norm_hidden_states,
+                encoder_hidden_states=condition_embed,
+                attention_mask=None,
+                **cross_attention_kwargs
+            )
+
+            hidden_states = attn_output + hidden_states
+            if hidden_states.ndim == 4:
+                hidden_states = hidden_states.squeeze(1)
+            
+
+        # 1.3 Multiview Attention
+        if num_in_batch > 1 and self.use_ma:
+            multivew_hidden_states = rearrange(norm_hidden_states, '(b n) l c -> b (n l) c', n=num_in_batch)
+            position_mask = None
+            if position_attn_mask is not None:
+                if multivew_hidden_states.shape[1] in position_attn_mask:
+                    position_mask = position_attn_mask[multivew_hidden_states.shape[1]]
+            position_indices = None
+            if position_voxel_indices is not None:
+                if multivew_hidden_states.shape[1] in position_voxel_indices:
+                    position_indices = position_voxel_indices[multivew_hidden_states.shape[1]]
+
+            attn_output = self.attn_multiview(
+                multivew_hidden_states,
+                encoder_hidden_states=multivew_hidden_states,
+                attention_mask=position_mask,
+                position_indices=position_indices,
+                **cross_attention_kwargs
+            )
+
+            attn_output = rearrange(attn_output, 'b (n l) c -> (b n) l c', n=num_in_batch)
+
+            hidden_states = 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)
+
+            if ip_hidden_states is not None:
+                ip_hidden_states = ip_hidden_states.unsqueeze(1).repeat(1,num_in_batch,1,1) # B N L C
+                ip_hidden_states = rearrange(ip_hidden_states, 'b n l c -> (b n) l c')
+
+            if self.use_ipa:
+                attn_output = self.attn2(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    ip_hidden_states=ip_hidden_states,
+                    attention_mask=encoder_attention_mask,
+                    **cross_attention_kwargs,
+                )
+            else:
+                attn_output = self.attn2(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=encoder_attention_mask,
+                    **cross_attention_kwargs,
+                )
+
+            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
+
+@torch.no_grad()
+def compute_voxel_grid_mask(position, grid_resolution=8):
+
+    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*-32
+    #weights = weights.clamp(min=-10000.0)
+    
+    weights = weights< grid_distance
+
+    return weights
+    
+def compute_multi_resolution_mask(position_maps, grid_resolutions=[32, 16, 8]):
+    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):
+
+    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).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 compute_multi_resolution_discrete_voxel_indice(position_maps, grid_resolutions=[64, 32, 16, 8], voxel_resolutions=[512, 256, 128, 64]):
+    voxel_indices = {}
+    with torch.no_grad():
+        for grid_resolution, voxel_resolution in zip(grid_resolutions, voxel_resolutions):
+            voxel_indice = compute_discrete_voxel_indice(position_maps, grid_resolution, voxel_resolution)
+            voxel_indice = rearrange(voxel_indice, 'b n c h w -> b (n h w) c')
+            voxel_indices[voxel_indice.shape[1]] = {'voxel_indices':voxel_indice, 'voxel_resolution':voxel_resolution}
+    return voxel_indices   
+    
+class ImageProjModel(torch.nn.Module):
+    """Projection Model"""
+
+    def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
+        super().__init__()
+
+        self.generator = None
+        self.cross_attention_dim = cross_attention_dim
+        self.clip_extra_context_tokens = clip_extra_context_tokens
+        self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
+        self.norm = torch.nn.LayerNorm(cross_attention_dim)
+
+    def forward(self, image_embeds):
+        embeds = image_embeds
+        clip_extra_context_tokens = self.proj(embeds).reshape(
+            -1, self.clip_extra_context_tokens, self.cross_attention_dim
+        )
+        clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
+        return clip_extra_context_tokens
+        
+class UNet2p5DConditionModel(torch.nn.Module):
+    def __init__(self, unet: UNet2DConditionModel) -> None:
+        super().__init__()
+        self.unet = unet
+        self.unet_dual = copy.deepcopy(unet)
+
+        self.init_camera_embedding()
+        self.init_attention(self.unet, use_ipa=True, use_ma=True, use_ra=True)
+        self.init_attention(self.unet_dual, use_ipa=False, use_ma=False, use_ra=False)
+        self.init_condition()
+
+    @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 = UNet2p5DConditionModel(unet)
+
+        unet.unet.conv_in = torch.nn.Conv2d(
+            12,
+            unet.unet.conv_in.out_channels,
+            kernel_size=unet.unet.conv_in.kernel_size,
+            stride=unet.unet.conv_in.stride,
+            padding=unet.unet.conv_in.padding,
+            dilation=unet.unet.conv_in.dilation,
+            groups=unet.unet.conv_in.groups,
+            bias=unet.unet.conv_in.bias is not None)
+        
+        unet_ckpt = torch.load(unet_ckpt_path, map_location='cpu', weights_only=True)
+        unet.load_state_dict(unet_ckpt, strict=True)
+        unet = unet.to(torch_dtype)
+        return unet
+        
+    def init_condition(self):
+        self.unet.learned_text_clip_gen = nn.Parameter(torch.randn(1,77,1024))
+        self.unet.learned_text_clip_ref = nn.Parameter(torch.randn(1,77,1024))
+
+        self.unet.image_proj_model = ImageProjModel(
+            cross_attention_dim=self.unet.config.cross_attention_dim,
+            clip_embeddings_dim=1024,
+        )
+
+
+    def init_camera_embedding(self):
+        self.max_num_ref_image = 5
+        self.max_num_gen_image = 12*3+4*2
+
+        time_embed_dim = 1280
+        self.unet.class_embedding = nn.Embedding(self.max_num_ref_image+self.max_num_gen_image, time_embed_dim)
+        # 将嵌入层的权重初始化为全零
+        nn.init.zeros_(self.unet.class_embedding.weight)
+    
+    def init_attention(self, unet, use_ipa=True, use_ma=True, use_ra=True):
+
+        for down_block_i, down_block in enumerate(unet.down_blocks):
+            if hasattr(down_block, "has_cross_attention") and down_block.has_cross_attention:
+                for attn_i, attn in enumerate(down_block.attentions):
+                    for transformer_i, transformer in enumerate(attn.transformer_blocks):
+                        if isinstance(transformer, BasicTransformerBlock):
+                            attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'down_{down_block_i}_{attn_i}_{transformer_i}',use_ipa,use_ma,use_ra)
+
+        if hasattr(unet.mid_block, "has_cross_attention") and unet.mid_block.has_cross_attention:
+            for attn_i, attn in enumerate(unet.mid_block.attentions):
+                for transformer_i, transformer in enumerate(attn.transformer_blocks):
+                    if isinstance(transformer, BasicTransformerBlock):
+                        attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'mid_{attn_i}_{transformer_i}',use_ipa,use_ma,use_ra)
+
+        for up_block_i, up_block in enumerate(unet.up_blocks):
+            if hasattr(up_block, "has_cross_attention") and up_block.has_cross_attention:
+                for attn_i, attn in enumerate(up_block.attentions):
+                    for transformer_i, transformer in enumerate(attn.transformer_blocks):
+                        if isinstance(transformer, BasicTransformerBlock):
+                            attn.transformer_blocks[transformer_i] = Basic2p5DTransformerBlock(transformer, f'up_{up_block_i}_{attn_i}_{transformer_i}',use_ipa,use_ma,use_ra)
+
+
+    def __getattr__(self, name: str):
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self.unet, name)
+        
+    def forward(
+        self, sample, timestep, encoder_hidden_states, class_labels=None,
+        *args, cross_attention_kwargs=None, down_intrablock_additional_residuals=None,
+        down_block_res_samples=None, mid_block_res_sample=None,
+        **cached_condition,
+    ):
+        B, N_gen, _, H, W = sample.shape
+        camera_info_gen = cached_condition['camera_info_gen'] + self.max_num_ref_image
+        camera_info_gen = rearrange(camera_info_gen, 'b n -> (b n)')
+        sample = [sample]
+        
+        if 'normal_imgs' in cached_condition:
+            sample.append(cached_condition["normal_imgs"])
+        if 'position_imgs' in cached_condition:
+            sample.append(cached_condition["position_imgs"])
+
+        sample = torch.cat(sample, dim=2)
+        sample = rearrange(sample, 'b n c h w -> (b n) c h w')
+
+        encoder_hidden_states_gen = encoder_hidden_states.unsqueeze(1).repeat(1, N_gen, 1, 1)
+        encoder_hidden_states_gen = rearrange(encoder_hidden_states_gen, 'b n l c -> (b n) l c')
+        
+        
+        use_position_mask = False
+        use_position_rope = True
+
+        position_attn_mask = None
+        if use_position_mask:
+            if 'position_attn_mask' in cached_condition:
+                position_attn_mask = cached_condition['position_attn_mask']
+            else:
+                if 'position_maps' in cached_condition:
+                    position_attn_mask = compute_multi_resolution_mask(cached_condition['position_maps'])
+        
+        position_voxel_indices = None
+        if use_position_rope:
+            if 'position_voxel_indices' in cached_condition:
+                position_voxel_indices = cached_condition['position_voxel_indices']
+            else:
+                if 'position_maps' in cached_condition:
+                    position_voxel_indices = compute_multi_resolution_discrete_voxel_indice(cached_condition['position_maps'])
+
+        if 'ip_hidden_states' in cached_condition:
+            ip_hidden_states = cached_condition['ip_hidden_states']
+        else:
+            if 'clip_embeds' in cached_condition:
+                ip_hidden_states = self.image_proj_model(cached_condition['clip_embeds'])
+            else:
+                ip_hidden_states = None
+            cached_condition['ip_hidden_states'] = ip_hidden_states
+
+        if 'condition_embed_dict' in cached_condition:
+            condition_embed_dict = cached_condition['condition_embed_dict']
+        else:
+            condition_embed_dict = {}
+            ref_latents = cached_condition['ref_latents']
+            N_ref = ref_latents.shape[1]
+            camera_info_ref = cached_condition['camera_info_ref']
+            camera_info_ref = rearrange(camera_info_ref, 'b n -> (b n)')
+            
+            #ref_latents = [ref_latents]
+            #if 'normal_imgs' in cached_condition:
+            #    ref_latents.append(torch.zeros_like(ref_latents[0]))
+            #if 'position_imgs' in cached_condition:
+            #    ref_latents.append(torch.zeros_like(ref_latents[0]))
+            #ref_latents = torch.cat(ref_latents, dim=2)
+            
+            ref_latents = rearrange(ref_latents, 'b n c h w -> (b n) c h w')
+
+            encoder_hidden_states_ref = self.learned_text_clip_ref.unsqueeze(1).repeat(B, N_ref, 1, 1)
+            encoder_hidden_states_ref = rearrange(encoder_hidden_states_ref, 'b n l c -> (b n) l c')
+
+            noisy_ref_latents = ref_latents
+            timestep_ref = 0
+            '''
+            if timestep.dim()>0:
+                timestep_ref = rearrange(timestep, '(b n) -> b n', b=B)[:,:1].repeat(1, N_ref)
+                timestep_ref = rearrange(timestep_ref, 'b n -> (b n)')
+            else:
+                timestep_ref = timestep
+            noise = torch.randn_like(noisy_ref_latents[:,:4,...])
+            if self.training:
+                noisy_ref_latents[:,:4,...] = self.train_sched.add_noise(noisy_ref_latents[:,:4,...], noise, timestep_ref)
+                noisy_ref_latents[:,:4,...] = self.train_sched.scale_model_input(noisy_ref_latents[:,:4,...], timestep_ref)
+            else:
+                noisy_ref_latents[:,:4,...] = self.val_sched.add_noise(noisy_ref_latents[:,:4,...], noise, timestep_ref.reshape(-1))
+                noisy_ref_latents[:,:4,...] = self.val_sched.scale_model_input(noisy_ref_latents[:,:4,...], timestep_ref.reshape(-1))
+            '''
+            self.unet_dual(
+                noisy_ref_latents, timestep_ref,
+                encoder_hidden_states=encoder_hidden_states_ref,
+                #class_labels=camera_info_ref,
+                # **kwargs
+                return_dict=False,
+                cross_attention_kwargs={
+                    'mode':'w', 'num_in_batch':N_ref, 
+                    'condition_embed_dict':condition_embed_dict},
+            )
+            cached_condition['condition_embed_dict'] = condition_embed_dict
+
+        return self.unet(
+            sample, timestep,
+            encoder_hidden_states_gen, *args,
+            class_labels=camera_info_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, 
+                'ip_hidden_states':ip_hidden_states,
+                'condition_embed_dict':condition_embed_dict, 
+                'position_attn_mask':position_attn_mask, 
+                'position_voxel_indices':position_voxel_indices
+            },
+        ) 

hunyuan3d-paint-v2-0-turbo/vae/config.json

@@ -0,0 +1,29 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.10.0.dev0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 768,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

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

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b4889b6b1d4ce7ae320a02dedaeff1780ad77d415ea0d744b476155c6377ddc
+size 334707217