compatibility with previous nodes + block replace

comfy/ldm/hunyuan3d/vae.py

@@ -1005,7 +1005,7 @@ class ShapeVAE(nn.Module):
         self.scale_factor = scale_factor
 
     def decode(self, latents, **kwargs):
-        latents = self.post_kl(latents)
+        latents = self.post_kl(latents.movedim(-2, -1))
         latents = self.transformer(latents)
 
         bounds = kwargs.get("bounds", 1.01)
@@ -1014,7 +1014,7 @@ class ShapeVAE(nn.Module):
         enable_pbar = kwargs.get("enable_pbar", True)
 
         grid_logits = self.volume_decoder(latents, self.geo_decoder, bounds=bounds, num_chunks=num_chunks, octree_resolution=octree_resolution, enable_pbar=enable_pbar)
-        return grid_logits
+        return grid_logits.movedim(-2, -1)
 
     def encode(self, surface):
 

comfy/ldm/hunyuan3dv2_1/hunyuandit.py

@@ -586,9 +586,11 @@ class HunYuanDiTPlain(nn.Module):
 
         self.final_layer = FinalLayer(hidden_size, self.out_channels, use_fp16 = use_fp16)
 
-    def forward(self, x, t, context, **kwargs):
-        
+    def forward(self, x, t, context, transformer_options = {}, **kwargs):
+
+        x = x.movedim(-1, -2)
         main_condition = context
+        t = 1.0 - t
 
         time_embedded = self.t_embedder(t, condition = kwargs.get('guidance_cond'))
 
@@ -597,20 +599,42 @@ class HunYuanDiTPlain(nn.Module):
 
         combined = torch.cat([time_embedded, x_embedded], dim=1)
 
+        def block_wrap(args):
+                    return block(
+                        args["x"],
+                        args["t"],
+                        args["cond"],
+                        skip_tensor=args.get("skip"),)
+
         skip_stack = []
+        patches_replace = transformer_options.get("patches_replace", {})
+        blocks_replace = patches_replace.get("dit", {})
         for idx, block in enumerate(self.blocks):
             if idx <= self.depth // 2:
                 skip_input = None
             else:
                 skip_input = skip_stack.pop()
 
-            combined = block(combined, time_embedded, main_condition, skip_tensor = skip_input)
+            if ("block", idx) in blocks_replace:
+
+                out = blocks_replace[("block", idx)](
+                    {
+                        "x": combined,
+                        "t": time_embedded,
+                        "cond": main_condition,
+                        "skip": skip_input,
+                    },
+                    {"original_block": block_wrap},
+                )
+                combined = out
+            else:
+                combined = block(combined, time_embedded, main_condition, skip_tensor=skip_input)
 
             if idx < self.depth // 2:
                 skip_stack.append(combined)
 
         output = self.final_layer(combined)
-        return output
+        return output.movedim(-2, -1) * (-1.0)
     
 def get_diffusion_checkpoint():
     import requests

comfy/model_base.py

@@ -1203,6 +1203,17 @@ class Hunyuan3Dv2_1(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan3dv2_1.hunyuandit.HunYuanDiTPlain)
 
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        cross_attn = kwargs.get("cross_attn", None)
+        if cross_attn is not None:
+            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+
+        guidance = kwargs.get("guidance", 5.0)
+        if guidance is not None:
+            out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
+        return out
+
 class HiDream(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hidream.model.HiDreamImageTransformer2DModel)

comfy/sd.py

@@ -15,7 +15,6 @@ import comfy.ldm.lightricks.vae.causal_video_autoencoder
 import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
 import comfy.ldm.hunyuan3d.vae
-import comfy.ldm.hunyuan3dv2_1.vae
 import comfy.ldm.ace.vae.music_dcae_pipeline
 import yaml
 import math

comfy_extras/nodes_hunyuan3d.py

@@ -11,16 +11,20 @@ from comfy.cli_args import args
 class EmptyLatentHunyuan3Dv2:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"resolution": ("INT", {"default": 3072, "min": 1, "max": 8192}),
-                             "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."}),
-                             }}
+        return {
+            "required": {
+                "resolution": ("INT", {"default": 3072, "min": 1, "max": 8192}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."}),
+            }
+        }
+
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "generate"
 
     CATEGORY = "latent/3d"
 
     def generate(self, resolution, batch_size):
-        latent = torch.zeros([batch_size, resolution, 64], device=comfy.model_management.intermediate_device())
+        latent = torch.zeros([batch_size, 64, resolution], device=comfy.model_management.intermediate_device())
         return ({"samples": latent, "type": "hunyuan3dv2"}, )
 
 class Hunyuan3Dv2Conditioning:
@@ -95,7 +99,7 @@ class VAEDecodeHunyuan3D:
     def decode(self, vae, samples, num_chunks, octree_resolution):
         voxels = VOXEL(vae.decode(samples["samples"], vae_options={"num_chunks": num_chunks, "octree_resolution": octree_resolution}))
         return (voxels, )
-    
+
 def voxel_to_mesh(voxels, threshold=0.5, device=None):
     if device is None:
         device = torch.device("cpu")
@@ -623,4 +627,4 @@ NODE_CLASS_MAPPINGS = {
     "VoxelToMeshBasic": VoxelToMeshBasic,
     "VoxelToMesh": VoxelToMesh,
     "SaveGLB": SaveGLB,
-}
+}