From ee65d6ea41da8d949d8b43f0687ccedb4e17a0af Mon Sep 17 00:00:00 2001 From: Yousef Rafat <81116377+yousef-rafat@users.noreply.github.com> Date: Sun, 13 Jul 2025 01:36:08 +0300 Subject: [PATCH] added dino2 large support and some fixes --- comfy/clip_vision.py | 6 +++- comfy/image_encoders/dino2.py | 14 +++++--- comfy/image_encoders/dino2_large.json | 22 +++++++++++++ comfy/ldm/hunyuan3dv2_1/hunyuandit.py | 13 ++++++-- comfy/ldm/hunyuan3dv2_1/vae.py | 32 ++++++++++-------- comfy/model_base.py | 36 -------------------- comfy/model_detection.py | 2 +- comfy/sd.py | 47 ++++++++++++++++++++------- comfy_extras/nodes_hunyuan3d.py | 12 ++++--- requirements.txt | 1 + 10 files changed, 108 insertions(+), 77 deletions(-) create mode 100644 comfy/image_encoders/dino2_large.json diff --git a/comfy/clip_vision.py b/comfy/clip_vision.py index 00aab9164..61fbdc741 100644 --- a/comfy/clip_vision.py +++ b/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 diff --git a/comfy/image_encoders/dino2.py b/comfy/image_encoders/dino2.py index 976f98c65..70720fec1 100644 --- a/comfy/image_encoders/dino2.py +++ b/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): diff --git a/comfy/image_encoders/dino2_large.json b/comfy/image_encoders/dino2_large.json new file mode 100644 index 000000000..43fbb58ff --- /dev/null +++ b/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] +} diff --git a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py b/comfy/ldm/hunyuan3dv2_1/hunyuandit.py index 2ea14f47c..f191fa20d 100644 --- a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py +++ b/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) diff --git a/comfy/ldm/hunyuan3dv2_1/vae.py b/comfy/ldm/hunyuan3dv2_1/vae.py index d09b5614a..95493b2ba 100644 --- a/comfy/ldm/hunyuan3dv2_1/vae.py +++ b/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) diff --git a/comfy/model_base.py b/comfy/model_base.py index 6aa034a06..55d1d34a3 100644 --- a/comfy/model_base.py +++ b/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) diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 44404f595..5d4907143 100644 --- a/comfy/model_detection.py +++ b/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 diff --git a/comfy/sd.py b/comfy/sd.py index 33cadfcea..59d4a69a9 100644 --- a/comfy/sd.py +++ b/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) diff --git a/comfy_extras/nodes_hunyuan3d.py b/comfy_extras/nodes_hunyuan3d.py index e91b48a99..cd60a3024 100644 --- a/comfy_extras/nodes_hunyuan3d.py +++ b/comfy_extras/nodes_hunyuan3d.py @@ -132,13 +132,15 @@ class VAEDecodeHunyuan3D: }) return (VOXEL(voxel), None) - mesh = vae.decode(samples["samples"], to_mesh = True, - num_chunks = num_chunks, - octree_resolution = octree_resolution) + mesh = vae.decode(samples["samples"],vae_options={ + "num_chunks": num_chunks, + "octree_resolution": octree_resolution, + "to_mesh": True + }) # ensure batch dim - if mesh.verticies.ndim == 2: - mesh.verticies = mesh.verticies[np.newaxis, ...] + if mesh.vertices.ndim == 2: + mesh.vertices = mesh.vertices[np.newaxis, ...] mesh.faces = mesh.faces[np.newaxis, ...] return (None, mesh) diff --git a/requirements.txt b/requirements.txt index 82e168b52..d1ee0067c 100644 --- a/requirements.txt +++ b/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) \ No newline at end of file