mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-17 11:50:05 +08:00
added dino2 large support and some fixes
This commit is contained in:
Yousef Rafat
parent
dff570e364
commit
ee65d6ea41
@ -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
|
||||
|
||||
|
||||
@ -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):
|
||||
|
||||
22
comfy/image_encoders/dino2_large.json
Normal file
22
comfy/image_encoders/dino2_large.json
Normal file
@ -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]
|
||||
}
|
||||
@ -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)
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
47
comfy/sd.py
47
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)
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -28,3 +28,4 @@ soundfile
|
||||
av>=14.2.0
|
||||
pydantic~=2.0
|
||||
pydantic-settings~=2.0
|
||||
scikit-image # for 3D mesh generation (marching cubes)
|
||||
Loading…
x
Reference in New Issue
Block a user