diff --git a/comfy/ldm/hunyuan3d/vae.py b/comfy/ldm/hunyuan3d/vae.py index f6738a2b7..5be9152a2 100644 --- a/comfy/ldm/hunyuan3d/vae.py +++ b/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): diff --git a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py b/comfy/ldm/hunyuan3dv2_1/hunyuandit.py index f191fa20d..06502262c 100644 --- a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py +++ b/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 diff --git a/comfy/model_base.py b/comfy/model_base.py index 55d1d34a3..b1148b979 100644 --- a/comfy/model_base.py +++ b/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) diff --git a/comfy/sd.py b/comfy/sd.py index 4904d10f8..19cee0afc 100644 --- a/comfy/sd.py +++ b/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 diff --git a/comfy_extras/nodes_hunyuan3d.py b/comfy_extras/nodes_hunyuan3d.py index 947be34e7..c6a9abb3c 100644 --- a/comfy_extras/nodes_hunyuan3d.py +++ b/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, -} +} \ No newline at end of file