added dino2 large support and some fixes

comfy/clip_vision.py

@@ -124,8 +124,12 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
                 json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336.json")
         else:
             json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
-    elif "embeddings.patch_embeddings.projection.weight" in sd:
+    
+    # Dinov2
+    elif 'encoder.layer.39.layer_scale2.lambda1' in sd:
         json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
+    elif 'encoder.layer.23.layer_scale2.lambda1' in sd:
+        json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json")
     else:
         return None
 

comfy/image_encoders/dino2.py

@@ -50,12 +50,14 @@ class SwiGLUFFN(torch.nn.Module):
 
 
 class Dino2Block(torch.nn.Module):
-    def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations):
+    def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations, use_swiglu_ffn):
         super().__init__()
         self.attention = Dino2AttentionBlock(dim, num_heads, layer_norm_eps, dtype, device, operations)
         self.layer_scale1 = LayerScale(dim, dtype, device, operations)
         self.layer_scale2 = LayerScale(dim, dtype, device, operations)
-        self.mlp = SwiGLUFFN(dim, dtype, device, operations)
+        if use_swiglu_ffn:
+            self.mlp = SwiGLUFFN(dim, dtype, device, operations)
+        else: self.mlp = torch.nn.Identity()
         self.norm1 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
         self.norm2 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
 
@@ -66,9 +68,10 @@ class Dino2Block(torch.nn.Module):
 
 
 class Dino2Encoder(torch.nn.Module):
-    def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations):
+    def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations, use_swiglu_ffn):
         super().__init__()
-        self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations) for _ in range(num_layers)])
+        self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations, use_swiglu_ffn = use_swiglu_ffn)
+                                          for _ in range(num_layers)])
 
     def forward(self, x, intermediate_output=None):
         optimized_attention = optimized_attention_for_device(x.device, False, small_input=True)
@@ -128,9 +131,10 @@ class Dinov2Model(torch.nn.Module):
         dim = config_dict["hidden_size"]
         heads = config_dict["num_attention_heads"]
         layer_norm_eps = config_dict["layer_norm_eps"]
+        use_swiglu_ffn = config_dict["use_swiglu_ffn"]
 
         self.embeddings = Dino2Embeddings(dim, dtype, device, operations)
-        self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations)
+        self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations, use_swiglu_ffn = use_swiglu_ffn)
         self.layernorm = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
 
     def forward(self, pixel_values, attention_mask=None, intermediate_output=None):

comfy/image_encoders/dino2_large.json

@@ -0,0 +1,22 @@
+{
+  "hidden_size": 1024,
+  "use_mask_token": true,
+  "patch_size": 14,
+  "image_size": 518,
+  "num_channels": 3,
+  "num_attention_heads": 16,
+  "initializer_range": 0.02,
+  "attention_probs_dropout_prob": 0.0,
+  "hidden_dropout_prob": 0.0,
+  "hidden_act": "gelu",
+  "mlp_ratio": 4,
+  "model_type": "dinov2",
+  "num_hidden_layers": 24,
+  "layer_norm_eps": 1e-6,
+  "qkv_bias": true,
+  "use_swiglu_ffn": false,
+  "layerscale_value": 1.0,
+  "drop_path_rate": 0.0,
+  "image_mean": [0.485, 0.456, 0.406],
+  "image_std": [0.229, 0.224, 0.225]
+}

comfy/ldm/hunyuan3dv2_1/hunyuandit.py

@@ -312,6 +312,8 @@ class CrossAttention(nn.Module):
         b, s1, _ = x.shape  
         _, s2, _ = y.shape 
 
+        y = y.to(next(self.to_k.parameters()).dtype)
+
         q = self.to_q(x)
         k = self.to_k(y)
         v = self.to_v(y)
@@ -531,14 +533,17 @@ class HunYuanDiTPlain(nn.Module):
         qk_norm: bool = True,
         qkv_bias: bool = False,
         num_moe_layers: int = 6,
-        guidance_cond_proj_dim = None,
+        guidance_cond_proj_dim = 2048,
         norm_type = 'layer',
         num_experts: int = 8,
         moe_top_k: int = 2,
         use_fp16: bool = False,
+        dtype = None,
         **kwargs
         ):
 
+        self.dtype = dtype
+
         super().__init__()
 
         self.depth = depth
@@ -581,11 +586,13 @@ class HunYuanDiTPlain(nn.Module):
 
         self.final_layer = FinalLayer(hidden_size, self.out_channels, use_fp16 = use_fp16)
 
-    def forward(self, x, t, contexts, **kwargs):
+    def forward(self, x, t, context, **kwargs):
         
-        main_condition = contexts['main']
+        main_condition = context
 
         time_embedded = self.t_embedder(t, condition = kwargs.get('guidance_cond'))
+
+        x = x.to(dtype = next(self.x_embedder.parameters()).dtype)
         x_embedded = self.x_embedder(x)
 
         combined = torch.cat([time_embedded, x_embedded], dim=1)

comfy/ldm/hunyuan3dv2_1/vae.py

@@ -5,10 +5,13 @@ import numpy as np
 from skimage import measure
 from dataclasses import dataclass
 import torch.nn as nn
-from hunyuan3d.vae import (
+
+import sys, os; 
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../..")))
+
+from comfy.ldm.hunyuan3d.vae import (
     CrossAttentionDecoder, Transformer, ResidualCrossAttentionBlock, FourierEmbedder, VanillaVolumeDecoder
 )
-
 def fps(src: Tensor, batch: Tensor, sampling_ratio: float, start_random: bool = True):
 
     # manually create the pointer vector
@@ -91,12 +94,11 @@ class PointCrossAttention(nn.Module):
         self.self_attn = None
         if layers > 0:
             self.self_attn = Transformer(
-                n_ctx = num_latents,
                 width = width,
                 heads = heads,
                 qkv_bias = qkv_bias,
                 qk_norm = qk_norm,
-                depth = layers
+                layers = layers
             )
 
         if use_ln_post:
@@ -281,7 +283,7 @@ class Latent2MeshOutput():
     vertices: None
     faces: None
 
-class SufraceExtractor():
+class SurfaceExtractor():
     def compute_box_stat(self, bounds, octree_resolution: int):
 
         # if float, turn it into a cube
@@ -293,14 +295,14 @@ class SufraceExtractor():
         grid_size = [int(octree_resolution) + 1, int(octree_resolution) + 1, int(octree_resolution) + 1]
         return grid_size, bbox_min, bbox_size
     
-    def run(self, grid_logit, *, bounds, octree_res, level: float = 0.0, **kwargs):
+    def run(self, grid_logit, *, bounds, octree_resolution, level: float = 0.0, **kwargs):
         # grid_logit from volume decoder
         # use marching cube algo to turn an sdf to a mesh
         vertices, faces, _, _ = measure.marching_cubes(grid_logit.cpu().numpy(),
                                            level,
                                            method = "lewiner")
         
-        grid_size, bbox_min, bbox_size = self.compute_box_stat(bounds = bounds, octree_resolution = octree_res)
+        grid_size, bbox_min, bbox_size = self.compute_box_stat(bounds = bounds, octree_resolution = octree_resolution)
         vertices = vertices / grid_size * bbox_size + bbox_min
 
         return vertices, faces
@@ -308,6 +310,8 @@ class SufraceExtractor():
     def __call__(self, grid_logits, **kwds):
         
         outputs = []
+        veritces_list = []
+        faces_list = []
         # loop over the batches
         for i in range(grid_logits.shape[0]):
             try:
@@ -315,7 +319,9 @@ class SufraceExtractor():
                 vertices, faces = self.run(grid_logits[i], **kwds)
                 vertices = vertices.astype(np.float32)
                 faces = np.ascontiguousarray(faces)
-                outputs.append(Latent2MeshOutput(vertices = vertices, faces = faces))
+                #outputs.append(Latent2MeshOutput(vertices = vertices, faces = faces))
+                veritces_list.append(vertices)
+                faces_list.append(faces)
 
             except Exception:
                 import traceback
@@ -552,9 +558,8 @@ class VAE(nn.Module):
                                            pc_sharpedge_size = pc_sharpedge_size) 
         
         self.transformer = Transformer(
-            n_ctx=num_latents,
             width=width,
-            depth=num_decoder_layers,
+            layers=num_decoder_layers,
             heads=heads,
             qkv_bias=qkv_bias,
             qk_norm=qk_norm,
@@ -564,7 +569,6 @@ class VAE(nn.Module):
         self.geo_decoder = CrossAttentionDecoder(
             fourier_embedder = self.fourier_embedder,
             out_channels = 1,
-            num_latents = num_latents,
             mlp_expand_ratio = geo_decoder_mlp_expand_ratio,
             downsample_ratio = geo_decoder_downsample_ratio,
             enable_ln_post = geo_decoder_ln_post,
@@ -578,7 +582,7 @@ class VAE(nn.Module):
         self.post_kl = nn.Linear(embed_dim, width)
 
         self.volume_decoder = VanillaVolumeDecoder()
-        self.surface_extractor = SufraceExtractor()
+        self.surface_extractor = SurfaceExtractor()
         
 
     def forward(self):
@@ -596,8 +600,8 @@ class VAE(nn.Module):
 
         return latents
 
-    def decode(self, latents, to_mesh: bool = True, **kwargs):
+    def decode(self, latents, **kwargs):
-
+        to_mesh = kwargs.pop("to_mesh", True)
         latents = self.post_kl(latents)
         latents = self.transformer(latents)
 

comfy/model_base.py

@@ -1203,42 +1203,6 @@ 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 get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=torch.float32):
-
-        assert len(w.shape) == 1
-        w = w * 1000.0
-
-        half_dim = embedding_dim // 2
-        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
-
-        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
-        emb = w.to(dtype)[:, None] * emb[None, :]
-
-        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-        if embedding_dim % 2 == 1:  # zero pad
-            emb = torch.nn.functional.pad(emb, (0, 1))
-
-        assert emb.shape == (w.shape[0], embedding_dim)
-
-        return emb
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-
-        guidance = kwargs.get("guidance", 5.0)
-        if guidance is not None:
-            
-            guidance_scale = torch.tensor([guidance], dtype = torch.float32, device = self.device)
-
-            guidance_embed = self.get_guidance_scale_embedding(guidance_scale,
-                                                               self.model.hidden_size,
-                                                               dtype = next(self.model.parameters()).dtype)
-            
-            out['guidance_cond'] = comfy.conds.CONDRegular(guidance_embed)
-
-        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/model_detection.py

@@ -399,7 +399,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["num_heads"] = 16
         dit_config["depth"] = count_blocks(state_dict_keys, f"{key_prefix}blocks.{{}}")
         dit_config["qkv_bias"] = False
-        dit_config["guidance_cond_proj_dim"] = f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys
+        dit_config["guidance_cond_proj_dim"] = None#f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys
         return dit_config
 
     if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys:  # HiDream

comfy/sd.py

@@ -431,17 +431,6 @@ class VAE:
                 self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
                 self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
                 self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
-            elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd:
-                self.latent_dim = 1
-                ln_post = "geo_decoder.ln_post.weight" in sd
-                inner_size = sd["geo_decoder.output_proj.weight"].shape[1]
-                downsample_ratio = sd["post_kl.weight"].shape[0] // inner_size
-                mlp_expand = sd["geo_decoder.cross_attn_decoder.mlp.c_fc.weight"].shape[0] // inner_size
-                self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)  # TODO
-                self.memory_used_decode = lambda shape, dtype: (1024 * 1024 * 1024 * 2.0) * model_management.dtype_size(dtype)  # TODO
-                ddconfig = {"embed_dim": 64, "num_freqs": 8, "include_pi": False, "heads": 16, "width": 1024, "num_decoder_layers": 16, "qkv_bias": False, "qk_norm": True, "geo_decoder_mlp_expand_ratio": mlp_expand, "geo_decoder_downsample_ratio": downsample_ratio, "geo_decoder_ln_post": ln_post}
-                self.first_stage_model = comfy.ldm.hunyuan3d.vae.ShapeVAE(**ddconfig)
-                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
 
             # Hunyuan 3d v2 2.1
             elif 'geo_decoder.cross_attn_decoder.mlp.c_proj.weight' in sd:
@@ -462,9 +451,22 @@ class VAE:
                 self.memory_used_decode = lambda shape, dtype, num_layers = 16, kv_cache_multiplier = 2: \
                     estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
 
-                self.first_stage_model = comfy.ldm.hunyuan3dv2_1.vae.ShapeVAE()
+                self.first_stage_model = comfy.ldm.hunyuan3dv2_1.vae.VAE()
                 self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
 
+            elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd:
+                self.latent_dim = 1
+                ln_post = "geo_decoder.ln_post.weight" in sd
+                inner_size = sd["geo_decoder.output_proj.weight"].shape[1]
+                downsample_ratio = sd["post_kl.weight"].shape[0] // inner_size
+                mlp_expand = sd["geo_decoder.cross_attn_decoder.mlp.c_fc.weight"].shape[0] // inner_size
+                self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)  # TODO
+                self.memory_used_decode = lambda shape, dtype: (1024 * 1024 * 1024 * 2.0) * model_management.dtype_size(dtype)  # TODO
+                ddconfig = {"embed_dim": 64, "num_freqs": 8, "include_pi": False, "heads": 16, "width": 1024, "num_decoder_layers": 16, "qkv_bias": False, "qk_norm": True, "geo_decoder_mlp_expand_ratio": mlp_expand, "geo_decoder_downsample_ratio": downsample_ratio, "geo_decoder_ln_post": ln_post}
+                self.first_stage_model = comfy.ldm.hunyuan3d.vae.ShapeVAE(**ddconfig)
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+
+
             elif "vocoder.backbone.channel_layers.0.0.bias" in sd: #Ace Step Audio
                 self.first_stage_model = comfy.ldm.ace.vae.music_dcae_pipeline.MusicDCAE(source_sample_rate=44100)
                 self.memory_used_encode = lambda shape, dtype: (shape[2] * 330) * model_management.dtype_size(dtype)
@@ -1039,6 +1041,27 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
     model = None
     model_patcher = None
 
+    if isinstance(sd, dict) and all(k in sd for k in ["model", "vae", "conditioner"]):
+        from collections import OrderedDict
+        import gc
+        
+        merged_sd = OrderedDict()
+
+        for k, v in sd["model"].items():
+            merged_sd[f"model.{k}"] = v
+
+        for k, v in sd["vae"].items():
+            merged_sd[f"vae.{k}"] = v
+
+        for key, value in sd["conditioner"].items():
+            merged_sd[f"conditioner.{key}"] = value
+
+        sd = merged_sd
+
+        del merged_sd
+        gc.collect()
+        torch.cuda.empty_cache()
+
     diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
     parameters = comfy.utils.calculate_parameters(sd, diffusion_model_prefix)
     weight_dtype = comfy.utils.weight_dtype(sd, diffusion_model_prefix)

comfy_extras/nodes_hunyuan3d.py

@@ -132,13 +132,15 @@ class VAEDecodeHunyuan3D:
             })
             return (VOXEL(voxel), None)
 
-        mesh = vae.decode(samples["samples"], to_mesh = True,
+        mesh = vae.decode(samples["samples"],vae_options={
-                         num_chunks = num_chunks,
+                "num_chunks": num_chunks,
-                         octree_resolution = octree_resolution)
+                "octree_resolution": octree_resolution,
+                "to_mesh": True
+        })
         
         # ensure batch dim
-        if mesh.verticies.ndim == 2:
+        if mesh.vertices.ndim == 2:
-            mesh.verticies = mesh.verticies[np.newaxis, ...]
+            mesh.vertices = mesh.vertices[np.newaxis, ...]
             mesh.faces = mesh.faces[np.newaxis, ...]
 
         return (None, mesh)

requirements.txt

@@ -28,3 +28,4 @@ soundfile
 av>=14.2.0
 pydantic~=2.0
 pydantic-settings~=2.0
+scikit-image  # for 3D mesh generation (marching cubes)