mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-17 11:22:18 +08:00
styling
This commit is contained in:
Yousef Rafat
parent
6ca9c64270
commit
1d923123b0
@ -36,7 +36,7 @@ def get_indices_weights(in_size, out_size, scale):
|
||||
x0 = x.floor().long()
|
||||
dx = x.unsqueeze(1) - (x0.unsqueeze(1) + torch.arange(-1, 3))
|
||||
|
||||
weights = cubic_kernel(dx)
|
||||
weights = cubic_kernel(dx)
|
||||
weights = weights / weights.sum(dim=1, keepdim=True)
|
||||
|
||||
indices = x0.unsqueeze(1) + torch.arange(-1, 3)
|
||||
@ -52,12 +52,12 @@ def resize_cubic_1d(x, out_size, dim):
|
||||
indices, weights = get_indices_weights(in_size, out_size, scale)
|
||||
|
||||
if dim == 2:
|
||||
x = x.permute(0, 1, 3, 2)
|
||||
x = x.reshape(-1, h)
|
||||
x = x.permute(0, 1, 3, 2)
|
||||
x = x.reshape(-1, h)
|
||||
else:
|
||||
x = x.reshape(-1, w)
|
||||
x = x.reshape(-1, w)
|
||||
|
||||
gathered = x[:, indices]
|
||||
gathered = x[:, indices]
|
||||
out = (gathered * weights.unsqueeze(0)).sum(dim=2)
|
||||
|
||||
if dim == 2:
|
||||
@ -77,7 +77,7 @@ def resize_cubic(img: torch.Tensor, size: tuple) -> torch.Tensor:
|
||||
img = img.unsqueeze(0)
|
||||
|
||||
img = img.permute(0, 3, 1, 2)
|
||||
|
||||
|
||||
out_h, out_w = size
|
||||
img = resize_cubic_1d(img, out_h, dim=2)
|
||||
img = resize_cubic_1d(img, out_w, dim=3)
|
||||
@ -121,7 +121,7 @@ def resize_area(img: torch.Tensor, size: tuple) -> torch.Tensor:
|
||||
# We will build the weighted sums by iterating over contributing input pixels once
|
||||
output = torch.zeros((B, C, out_h, out_w), dtype=torch.float32, device=device)
|
||||
area = torch.zeros((out_h, out_w), dtype=torch.float32, device=device)
|
||||
|
||||
|
||||
max_kernel_h = int(torch.max(y_end_int - y_start_int).item())
|
||||
max_kernel_w = int(torch.max(x_end_int - x_start_int).item())
|
||||
|
||||
@ -129,8 +129,8 @@ def resize_area(img: torch.Tensor, size: tuple) -> torch.Tensor:
|
||||
for dx in range(max_kernel_w):
|
||||
# compute the weights for this offset for all output pixels
|
||||
|
||||
y_idx = y_start_int.unsqueeze(1) + dy
|
||||
x_idx = x_start_int.unsqueeze(0) + dx
|
||||
y_idx = y_start_int.unsqueeze(1) + dy
|
||||
x_idx = x_start_int.unsqueeze(0) + dx
|
||||
|
||||
# clamp indices to image boundaries
|
||||
y_idx_clamped = torch.clamp(y_idx, 0, H - 1)
|
||||
@ -186,10 +186,10 @@ def recenter(image, border_ratio: float = 0.2):
|
||||
|
||||
h = x_max - x_min
|
||||
w = y_max - y_min
|
||||
|
||||
|
||||
if h == 0 or w == 0:
|
||||
raise ValueError('input image is empty')
|
||||
|
||||
|
||||
desired_size = int(size * (1 - border_ratio))
|
||||
scale = desired_size / max(h, w)
|
||||
|
||||
@ -213,31 +213,31 @@ def recenter(image, border_ratio: float = 0.2):
|
||||
mask = mask * 255
|
||||
result = result.clip(0, 255).to(torch.uint8)
|
||||
mask = mask.clip(0, 255).to(torch.uint8)
|
||||
|
||||
|
||||
return result
|
||||
|
||||
def clip_preprocess(image, size=224, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711],
|
||||
crop=True, value_range = (-1, 1), border_ratio: float = None, recenter_size: int = 512):
|
||||
|
||||
if border_ratio is not None:
|
||||
|
||||
|
||||
image = (image * 255).clamp(0, 255).to(torch.uint8)
|
||||
image = [recenter(img, border_ratio = border_ratio) for img in image]
|
||||
|
||||
|
||||
image = torch.stack(image, dim = 0)
|
||||
image = resize_cubic(image, size = (recenter_size, recenter_size))
|
||||
|
||||
|
||||
image = image / 255 * 2 - 1
|
||||
low, high = value_range
|
||||
|
||||
|
||||
image = (image - low) / (high - low)
|
||||
image = image.permute(0, 2, 3, 1)
|
||||
|
||||
|
||||
image = image[:, :, :, :3] if image.shape[3] > 3 else image
|
||||
|
||||
|
||||
mean = torch.tensor(mean, device=image.device, dtype=image.dtype)
|
||||
std = torch.tensor(std, device=image.device, dtype=image.dtype)
|
||||
|
||||
|
||||
image = image.movedim(-1, 1)
|
||||
if not (image.shape[2] == size and image.shape[3] == size):
|
||||
if crop:
|
||||
@ -341,7 +341,7 @@ 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")
|
||||
|
||||
|
||||
# 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")
|
||||
|
||||
@ -34,7 +34,7 @@ class LayerScale(torch.nn.Module):
|
||||
class Dinov2MLP(torch.nn.Module):
|
||||
def __init__(self, hidden_size: int, dtype, device, operations):
|
||||
super().__init__()
|
||||
|
||||
|
||||
mlp_ratio = 4
|
||||
hidden_features = int(hidden_size * mlp_ratio)
|
||||
self.fc1 = operations.Linear(hidden_size, hidden_features, bias = True, device=device, dtype=dtype)
|
||||
@ -45,7 +45,7 @@ class Dinov2MLP(torch.nn.Module):
|
||||
hidden_state = torch.nn.functional.gelu(hidden_state)
|
||||
hidden_state = self.fc2(hidden_state)
|
||||
return hidden_state
|
||||
|
||||
|
||||
class SwiGLUFFN(torch.nn.Module):
|
||||
def __init__(self, dim, dtype, device, operations):
|
||||
super().__init__()
|
||||
@ -71,7 +71,7 @@ class Dino2Block(torch.nn.Module):
|
||||
self.layer_scale2 = LayerScale(dim, dtype, device, operations)
|
||||
if use_swiglu_ffn:
|
||||
self.mlp = SwiGLUFFN(dim, dtype, device, operations)
|
||||
else:
|
||||
else:
|
||||
self.mlp = Dinov2MLP(dim, dtype, device, operations)
|
||||
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)
|
||||
|
||||
@ -46,7 +46,7 @@ def fps(src: torch.Tensor, batch: torch.Tensor, sampling_ratio: float, start_ran
|
||||
# select a random start point
|
||||
if start_random:
|
||||
farthest = torch.randint(0, num_points, (1,), device = src.device)
|
||||
else:
|
||||
else:
|
||||
farthest = torch.tensor([0], device = src.device, dtype = torch.long)
|
||||
|
||||
for i in range(num_samples):
|
||||
@ -134,24 +134,24 @@ class PointCrossAttention(nn.Module):
|
||||
# Split random and sharpedge surface points
|
||||
random_pc, sharpedge_pc = torch.split(point_cloud, [self.pc_size, self.pc_sharpedge_size], dim=1)
|
||||
|
||||
# assert statements
|
||||
# assert statements
|
||||
assert random_pc.shape[1] <= self.pc_size, "Random surface points size must be less than or equal to pc_size"
|
||||
assert sharpedge_pc.shape[1] <= self.pc_sharpedge_size, "Sharpedge surface points size must be less than or equal to pc_sharpedge_size"
|
||||
|
||||
input_random_pc_size = int(num_random_query * self.downsample_ratio)
|
||||
random_query_pc, random_input_pc, random_idx_pc, random_idx_query = \
|
||||
self.subsample(pc = random_pc, num_query = num_random_query, input_pc_size = input_random_pc_size)
|
||||
|
||||
|
||||
input_sharpedge_pc_size = int(num_sharpedge_query * self.downsample_ratio)
|
||||
|
||||
if input_sharpedge_pc_size == 0:
|
||||
sharpedge_input_pc = torch.zeros(B, 0, D, dtype = random_input_pc.dtype).to(point_cloud.device)
|
||||
sharpedge_query_pc = torch.zeros(B, 0, D, dtype= random_query_pc.dtype).to(point_cloud.device)
|
||||
|
||||
else:
|
||||
else:
|
||||
sharpedge_query_pc, sharpedge_input_pc, sharpedge_idx_pc, sharpedge_idx_query = \
|
||||
self.subsample(pc = sharpedge_pc, num_query = num_sharpedge_query, input_pc_size = input_sharpedge_pc_size)
|
||||
|
||||
|
||||
# concat the random and sharpedges
|
||||
query_pc = torch.cat([random_query_pc, sharpedge_query_pc], dim = 1)
|
||||
input_pc = torch.cat([random_input_pc, sharpedge_input_pc], dim = 1)
|
||||
@ -165,11 +165,11 @@ class PointCrossAttention(nn.Module):
|
||||
input_random_surface_features, query_random_features = \
|
||||
self.handle_features(features = random_surface_features, idx_pc = random_idx_pc, batch_size = B,
|
||||
input_pc_size = input_random_pc_size, idx_query = random_idx_query)
|
||||
|
||||
|
||||
if input_sharpedge_pc_size == 0:
|
||||
input_sharpedge_surface_features = torch.zeros(B, 0, self.point_feats,
|
||||
dtype = input_random_surface_features.dtype, device = point_cloud.device)
|
||||
|
||||
|
||||
query_sharpedge_features = torch.zeros(B, 0, self.point_feats,
|
||||
dtype = query_random_features.dtype, device = point_cloud.device)
|
||||
else:
|
||||
@ -197,7 +197,7 @@ class PointCrossAttention(nn.Module):
|
||||
return query.view(B, -1, query.shape[-1]), data.view(B, -1, data.shape[-1])
|
||||
|
||||
def forward(self, point_cloud: torch.Tensor, features: torch.Tensor):
|
||||
|
||||
|
||||
query, data = self.sample_points_and_latents(point_cloud = point_cloud, features = features)
|
||||
|
||||
# apply projections
|
||||
@ -243,16 +243,16 @@ class PointCrossAttention(nn.Module):
|
||||
query_pc = flattent_input_pc[idx_query].view(B, -1, D)
|
||||
|
||||
return query_pc, input_pc, idx_pc, idx_query
|
||||
|
||||
|
||||
def handle_features(self, features, idx_pc, input_pc_size, batch_size: int, idx_query):
|
||||
|
||||
B = batch_size
|
||||
|
||||
input_surface_features = features[:, idx_pc, :]
|
||||
flattent_input_features = input_surface_features.view(B * input_pc_size, -1)
|
||||
query_features = flattent_input_features[idx_query].view(B, -1,
|
||||
query_features = flattent_input_features[idx_query].view(B, -1,
|
||||
flattent_input_features.shape[-1])
|
||||
|
||||
|
||||
return input_surface_features, query_features
|
||||
|
||||
def normalize_mesh(mesh, scale = 0.9999):
|
||||
@ -361,7 +361,7 @@ def load_surface_sharpedge(mesh, num_points=4096, num_sharp_points=4096, sharped
|
||||
surface = assemble_tensor(torch.cat([surf_pts.to(device), fill_pts.to(device)], dim=0),
|
||||
torch.cat([surf_normals.to(device), fill_normals.to(device)], dim=0),
|
||||
label = 0 if sharpedge_flag else None)
|
||||
|
||||
|
||||
sharp_surface = assemble_tensor(torch.from_numpy(sharp_pts), torch.from_numpy(sharp_normals),
|
||||
label = 1 if sharpedge_flag else None)
|
||||
|
||||
@ -401,9 +401,9 @@ class SharpEdgeSurfaceLoader:
|
||||
for obj in mesh.geometry.values():
|
||||
combined = obj if combined is None else combined + obj
|
||||
return combined
|
||||
|
||||
|
||||
return mesh
|
||||
|
||||
|
||||
class DiagonalGaussianDistribution:
|
||||
def __init__(self, params: torch.Tensor, feature_dim: int = -1):
|
||||
|
||||
@ -428,7 +428,7 @@ class VanillaVolumeDecoder():
|
||||
@torch.no_grad()
|
||||
def __call__(self, latents: torch.Tensor, geo_decoder: callable, octree_resolution: int, bounds = 1.01,
|
||||
num_chunks: int = 10_000, enable_pbar: bool = True, **kwargs):
|
||||
|
||||
|
||||
if isinstance(bounds, float):
|
||||
bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
|
||||
|
||||
@ -445,7 +445,7 @@ class VanillaVolumeDecoder():
|
||||
batch_logits = []
|
||||
for start in tqdm(range(0, xyz.shape[0], num_chunks), desc="Volume Decoding",
|
||||
disable=not enable_pbar):
|
||||
|
||||
|
||||
chunk_queries = xyz[start: start + num_chunks, :]
|
||||
chunk_queries = chunk_queries.unsqueeze(0).repeat(latents.shape[0], 1, 1)
|
||||
logits = geo_decoder(queries = chunk_queries, latents = latents)
|
||||
@ -929,7 +929,7 @@ class ShapeVAE(nn.Module):
|
||||
width = width,
|
||||
point_feats = point_feats,
|
||||
fourier_embedder = self.fourier_embedder,
|
||||
pc_sharpedge_size = pc_sharpedge_size)
|
||||
pc_sharpedge_size = pc_sharpedge_size)
|
||||
|
||||
self.post_kl = ops.Linear(embed_dim, width)
|
||||
|
||||
@ -974,10 +974,10 @@ class ShapeVAE(nn.Module):
|
||||
|
||||
pc, feats = surface[:, :, :3], surface[:, :, 3:]
|
||||
latents = self.encoder(pc, feats)
|
||||
|
||||
|
||||
moments = self.pre_kl(latents)
|
||||
posterior = DiagonalGaussianDistribution(moments, feature_dim = -1)
|
||||
|
||||
latents = posterior.sample()
|
||||
|
||||
return latents
|
||||
return latents
|
||||
|
||||
@ -14,7 +14,7 @@ class GELU(nn.Module):
|
||||
|
||||
if gate.device.type == "mps":
|
||||
return F.gelu(gate.to(dtype = torch.float32)).to(dtype = gate.dtype)
|
||||
|
||||
|
||||
return F.gelu(gate)
|
||||
|
||||
def forward(self, hidden_states):
|
||||
@ -28,7 +28,7 @@ class FeedForward(nn.Module):
|
||||
|
||||
def __init__(self, dim: int, dim_out = None, mult: int = 4,
|
||||
dropout: float = 0.0, inner_dim = None):
|
||||
|
||||
|
||||
super().__init__()
|
||||
if inner_dim is None:
|
||||
inner_dim = int(dim * mult)
|
||||
@ -53,11 +53,11 @@ class AddAuxLoss(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, x, loss):
|
||||
# do nothing in forward (no computation)
|
||||
ctx.requires_aux_loss = loss.requires_grad
|
||||
ctx.requires_aux_loss = loss.requires_grad
|
||||
ctx.dtype = loss.dtype
|
||||
|
||||
return x
|
||||
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
# add the aux loss gradients
|
||||
@ -68,7 +68,7 @@ class AddAuxLoss(torch.autograd.Function):
|
||||
grad_loss = torch.ones(1, dtype = ctx.dtype, device = grad_output.device)
|
||||
|
||||
return grad_output, grad_loss
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
|
||||
def __init__(self, embed_dim, num_experts=16, num_experts_per_tok=2, aux_loss_alpha=0.01):
|
||||
@ -133,13 +133,13 @@ class MoEBlock(nn.Module):
|
||||
|
||||
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
|
||||
flat_topk_idx = topk_idx.view(-1)
|
||||
|
||||
|
||||
if self.training:
|
||||
|
||||
hidden_states = hidden_states.repeat_interleave(self.moe_top_k, dim = 0)
|
||||
y = torch.empty_like(hidden_states, dtype = hidden_states.dtype)
|
||||
|
||||
for i, expert in enumerate(self.experts):
|
||||
for i, expert in enumerate(self.experts):
|
||||
tmp = expert(hidden_states[flat_topk_idx == i])
|
||||
y[flat_topk_idx == i] = tmp.to(hidden_states.dtype)
|
||||
|
||||
@ -156,14 +156,14 @@ class MoEBlock(nn.Module):
|
||||
|
||||
@torch.no_grad()
|
||||
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
||||
|
||||
expert_cache = torch.zeros_like(x)
|
||||
|
||||
expert_cache = torch.zeros_like(x)
|
||||
idxs = flat_expert_indices.argsort()
|
||||
|
||||
# no need for .numpy().cpu() here
|
||||
tokens_per_expert = flat_expert_indices.bincount().cumsum(0)
|
||||
token_idxs = idxs // self.moe_top_k
|
||||
|
||||
token_idxs = idxs // self.moe_top_k
|
||||
|
||||
for i, end_idx in enumerate(tokens_per_expert):
|
||||
|
||||
start_idx = 0 if i == 0 else tokens_per_expert[i-1]
|
||||
@ -177,7 +177,7 @@ class MoEBlock(nn.Module):
|
||||
expert_tokens = x[exp_token_idx]
|
||||
expert_out = expert(expert_tokens)
|
||||
|
||||
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
||||
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
||||
|
||||
# use index_add_ with a 1-D index tensor directly avoids building a large [N, D] index map and extra memcopy required by scatter_reduce_
|
||||
# + avoid dtype conversion
|
||||
@ -198,7 +198,7 @@ class Timesteps(nn.Module):
|
||||
half_dim, dtype=torch.float32
|
||||
) / (half_dim - downscale_freq_shift)
|
||||
|
||||
inv_freq = torch.exp(exponent)
|
||||
inv_freq = torch.exp(exponent)
|
||||
|
||||
# pad
|
||||
if num_channels % 2 == 1:
|
||||
@ -212,7 +212,7 @@ class Timesteps(nn.Module):
|
||||
def forward(self, timesteps: torch.Tensor):
|
||||
|
||||
x = timesteps.float().unsqueeze(1) * self.inv_freq.to(timesteps.device).unsqueeze(0)
|
||||
|
||||
|
||||
|
||||
# fused CUDA kernels for sin and cos
|
||||
sin_emb = x.sin()
|
||||
@ -223,7 +223,7 @@ class Timesteps(nn.Module):
|
||||
# scale factor
|
||||
if self.scale != 1.0:
|
||||
emb = emb * self.scale
|
||||
|
||||
|
||||
# If we padded inv_freq for odd, emb is already wide enough; otherwise:
|
||||
if emb.shape[1] > self.num_channels:
|
||||
emb = emb[:, :self.num_channels]
|
||||
@ -232,8 +232,8 @@ class Timesteps(nn.Module):
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
def __init__(self, hidden_size, frequency_embedding_size = 256, cond_proj_dim = None):
|
||||
super().__init__()
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(hidden_size, frequency_embedding_size, bias=True),
|
||||
nn.GELU(),
|
||||
@ -257,8 +257,8 @@ class TimestepEmbedder(nn.Module):
|
||||
time_conditioned = self.mlp(timestep_embed.to(self.mlp[0].weight.device))
|
||||
|
||||
# for broadcasting with image tokens
|
||||
return time_conditioned.unsqueeze(1)
|
||||
|
||||
return time_conditioned.unsqueeze(1)
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, *, width: int):
|
||||
super().__init__()
|
||||
@ -269,7 +269,7 @@ class MLP(nn.Module):
|
||||
|
||||
def forward(self, x):
|
||||
return self.fc2(self.gelu(self.fc1(x)))
|
||||
|
||||
|
||||
class CrossAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@ -285,10 +285,10 @@ class CrossAttention(nn.Module):
|
||||
super().__init__()
|
||||
self.qdim = qdim
|
||||
self.kdim = kdim
|
||||
|
||||
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = self.qdim // num_heads
|
||||
|
||||
|
||||
self.scale = self.head_dim ** -0.5
|
||||
|
||||
self.to_q = nn.Linear(qdim, qdim, bias=qkv_bias)
|
||||
@ -307,11 +307,11 @@ class CrossAttention(nn.Module):
|
||||
self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps) if qk_norm else nn.Identity()
|
||||
self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps) if qk_norm else nn.Identity()
|
||||
self.out_proj = nn.Linear(qdim, qdim, bias=True)
|
||||
|
||||
|
||||
def forward(self, x, y):
|
||||
|
||||
b, s1, _ = x.shape
|
||||
_, s2, _ = y.shape
|
||||
b, s1, _ = x.shape
|
||||
_, s2, _ = y.shape
|
||||
|
||||
y = y.to(next(self.to_k.parameters()).dtype)
|
||||
|
||||
@ -325,9 +325,9 @@ class CrossAttention(nn.Module):
|
||||
kv = kv.view(1, -1, self.num_heads, split_size * 2)
|
||||
k, v = torch.split(kv, split_size, dim=-1)
|
||||
|
||||
q = q.view(b, s1, self.num_heads, self.head_dim)
|
||||
k = k.view(b, s2, self.num_heads, self.head_dim)
|
||||
v = v.view(b, s2, self.num_heads, self.head_dim)
|
||||
q = q.view(b, s1, self.num_heads, self.head_dim)
|
||||
k = k.view(b, s2, self.num_heads, self.head_dim)
|
||||
v = v.view(b, s2, self.num_heads, self.head_dim)
|
||||
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
@ -348,7 +348,7 @@ class CrossAttention(nn.Module):
|
||||
out = self.out_proj(context)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
|
||||
def __init__(
|
||||
@ -370,9 +370,9 @@ class Attention(nn.Module):
|
||||
self.to_k = nn.Linear(dim, dim, bias = qkv_bias)
|
||||
self.to_v = nn.Linear(dim, dim, bias = qkv_bias)
|
||||
|
||||
if use_fp16:
|
||||
if use_fp16:
|
||||
eps = 1.0 / 65504
|
||||
else:
|
||||
else:
|
||||
eps = 1e-6
|
||||
|
||||
self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps) if qk_norm else nn.Identity()
|
||||
@ -392,8 +392,8 @@ class Attention(nn.Module):
|
||||
qkv = qkv_combined.view(1, -1, self.num_heads, split_size * 3)
|
||||
query, key, value = torch.split(qkv, split_size, dim=-1)
|
||||
|
||||
query = query.reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
key = key.reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
query = query.reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
key = key.reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
value = value.reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
|
||||
|
||||
query = self.q_norm(query)
|
||||
@ -412,7 +412,7 @@ class Attention(nn.Module):
|
||||
|
||||
x = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class HunYuanDiTBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
@ -434,11 +434,11 @@ class HunYuanDiTBlock(nn.Module):
|
||||
super().__init__()
|
||||
|
||||
# eps can't be 1e-6 in fp16 mode because of numerical stability issues
|
||||
if use_fp16:
|
||||
if use_fp16:
|
||||
eps = 1.0 / 65504
|
||||
else:
|
||||
else:
|
||||
eps = 1e-6
|
||||
|
||||
|
||||
self.norm1 = norm_layer(hidden_size, elementwise_affine = True, eps = eps)
|
||||
|
||||
self.attn1 = Attention(hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
|
||||
@ -455,7 +455,7 @@ class HunYuanDiTBlock(nn.Module):
|
||||
|
||||
self.attn2 = CrossAttention(hidden_size, text_states_dim, num_heads=num_heads, qkv_bias=qkv_bias,
|
||||
qk_norm=qk_norm, norm_layer=qk_norm_layer, use_fp16 = use_fp16)
|
||||
|
||||
|
||||
self.norm3 = norm_layer(hidden_size, elementwise_affine = True, eps = eps)
|
||||
|
||||
if skip_connection:
|
||||
@ -504,15 +504,15 @@ class HunYuanDiTBlock(nn.Module):
|
||||
hidden_states = hidden_states + self.mlp(mlp_input)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
|
||||
def __init__(self, final_hidden_size, out_channels, use_fp16: bool = False):
|
||||
super().__init__()
|
||||
|
||||
if use_fp16:
|
||||
if use_fp16:
|
||||
eps = 1.0 / 65504
|
||||
else:
|
||||
else:
|
||||
eps = 1e-6
|
||||
|
||||
self.norm_final = nn.LayerNorm(final_hidden_size, elementwise_affine = True, eps = eps)
|
||||
@ -525,7 +525,7 @@ class FinalLayer(nn.Module):
|
||||
return x
|
||||
|
||||
class HunYuanDiTPlain(nn.Module):
|
||||
|
||||
|
||||
# init with the defaults values from https://huggingface.co/tencent/Hunyuan3D-2.1/blob/main/hunyuan3d-dit-v2-1/config.yaml
|
||||
def __init__(
|
||||
self,
|
||||
@ -642,6 +642,6 @@ class HunYuanDiTPlain(nn.Module):
|
||||
|
||||
output = self.final_layer(combined)
|
||||
output = output.movedim(-2, -1) * (-1.0)
|
||||
|
||||
|
||||
cond_emb, uncond_emb = output.chunk(2, dim = 0)
|
||||
return torch.cat([uncond_emb, cond_emb])
|
||||
return torch.cat([uncond_emb, cond_emb])
|
||||
|
||||
@ -1224,7 +1224,7 @@ class Hunyuan3Dv2(BaseModel):
|
||||
if guidance is not None:
|
||||
out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
|
||||
return out
|
||||
|
||||
|
||||
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)
|
||||
|
||||
@ -459,11 +459,11 @@ class VAE:
|
||||
|
||||
# better memory estimations
|
||||
self.memory_used_encode = lambda shape, dtype, num_layers = 8, kv_cache_multiplier = 0:\
|
||||
estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
|
||||
estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
|
||||
|
||||
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.hunyuan3d.vae.ShapeVAE()
|
||||
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
|
||||
@ -1051,7 +1051,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
|
||||
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():
|
||||
|
||||
@ -1101,7 +1101,7 @@ class Hunyuan3Dv2(supported_models_base.BASE):
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
return None
|
||||
|
||||
|
||||
class Hunyuan3Dv2_1(Hunyuan3Dv2):
|
||||
unet_config = {
|
||||
"image_model": "hunyuan3d2_1",
|
||||
|
||||
@ -627,4 +627,4 @@ NODE_CLASS_MAPPINGS = {
|
||||
"VoxelToMeshBasic": VoxelToMeshBasic,
|
||||
"VoxelToMesh": VoxelToMesh,
|
||||
"SaveGLB": SaveGLB,
|
||||
}
|
||||
}
|
||||
|
||||
4
nodes.py
4
nodes.py
@ -1003,7 +1003,7 @@ class CLIPVisionEncode:
|
||||
"border_ratio": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 0.5, "step": 0.01, "visible_if": {"crop": "recenter"},}),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
RETURN_TYPES = ("CLIP_VISION_OUTPUT",)
|
||||
FUNCTION = "encode"
|
||||
|
||||
@ -1014,7 +1014,7 @@ class CLIPVisionEncode:
|
||||
|
||||
if crop == "recenter":
|
||||
crop_image = True
|
||||
else:
|
||||
else:
|
||||
border_ratio = None
|
||||
|
||||
output = clip_vision.encode_image(image, crop=crop_image, border_ratio = border_ratio)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user