diff --git a/.github/workflows/stable-release.yml b/.github/workflows/stable-release.yml
index e5444d36..0bdd5a3b 100644
--- a/.github/workflows/stable-release.yml
+++ b/.github/workflows/stable-release.yml
@@ -17,12 +17,12 @@ on:
description: 'Python minor version'
required: true
type: string
- default: "11"
+ default: "12"
python_patch:
description: 'Python patch version'
required: true
type: string
- default: "9"
+ default: "7"
jobs:
diff --git a/.github/workflows/windows_release_dependencies.yml b/.github/workflows/windows_release_dependencies.yml
index 8d3a8665..85e6a52f 100644
--- a/.github/workflows/windows_release_dependencies.yml
+++ b/.github/workflows/windows_release_dependencies.yml
@@ -12,7 +12,7 @@ on:
description: 'extra dependencies'
required: false
type: string
- default: "\"numpy<2\""
+ default: ""
cu:
description: 'cuda version'
required: true
@@ -23,13 +23,13 @@ on:
description: 'python minor version'
required: true
type: string
- default: "11"
+ default: "12"
python_patch:
description: 'python patch version'
required: true
type: string
- default: "9"
+ default: "7"
# push:
# branches:
# - master
diff --git a/.github/workflows/windows_release_package.yml b/.github/workflows/windows_release_package.yml
index 84d99b8a..11e724ba 100644
--- a/.github/workflows/windows_release_package.yml
+++ b/.github/workflows/windows_release_package.yml
@@ -13,13 +13,13 @@ on:
description: 'python minor version'
required: true
type: string
- default: "11"
+ default: "12"
python_patch:
description: 'python patch version'
required: true
type: string
- default: "9"
+ default: "7"
# push:
# branches:
# - master
diff --git a/README.md b/README.md
index 745cf1a3..8e5cef45 100644
--- a/README.md
+++ b/README.md
@@ -127,6 +127,8 @@ To run it on services like paperspace, kaggle or colab you can use my [Jupyter N
## Manual Install (Windows, Linux)
+Note that some dependencies do not yet support python 3.13 so using 3.12 is recommended.
+
Git clone this repo.
Put your SD checkpoints (the huge ckpt/safetensors files) in: models/checkpoints
diff --git a/app/frontend_management.py b/app/frontend_management.py
index 9c832e46..6f20e439 100644
--- a/app/frontend_management.py
+++ b/app/frontend_management.py
@@ -151,6 +151,15 @@ class FrontendManager:
return cls.DEFAULT_FRONTEND_PATH
repo_owner, repo_name, version = cls.parse_version_string(version_string)
+
+ if version.startswith("v"):
+ expected_path = str(Path(cls.CUSTOM_FRONTENDS_ROOT) / f"{repo_owner}_{repo_name}" / version.lstrip("v"))
+ if os.path.exists(expected_path):
+ logging.info(f"Using existing copy of specific frontend version tag: {repo_owner}/{repo_name}@{version}")
+ return expected_path
+
+ logging.info(f"Initializing frontend: {repo_owner}/{repo_name}@{version}, requesting version details from GitHub...")
+
provider = provider or FrontEndProvider(repo_owner, repo_name)
release = provider.get_release(version)
diff --git a/comfy/controlnet.py b/comfy/controlnet.py
index 1d24afa6..d2744e42 100644
--- a/comfy/controlnet.py
+++ b/comfy/controlnet.py
@@ -60,7 +60,7 @@ class StrengthType(Enum):
LINEAR_UP = 2
class ControlBase:
- def __init__(self, device=None):
+ def __init__(self):
self.cond_hint_original = None
self.cond_hint = None
self.strength = 1.0
@@ -72,10 +72,6 @@ class ControlBase:
self.compression_ratio = 8
self.upscale_algorithm = 'nearest-exact'
self.extra_args = {}
-
- if device is None:
- device = comfy.model_management.get_torch_device()
- self.device = device
self.previous_controlnet = None
self.extra_conds = []
self.strength_type = StrengthType.CONSTANT
@@ -185,8 +181,8 @@ class ControlBase:
class ControlNet(ControlBase):
- def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False):
- super().__init__(device)
+ def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False):
+ super().__init__()
self.control_model = control_model
self.load_device = load_device
if control_model is not None:
@@ -242,7 +238,7 @@ class ControlNet(ControlBase):
to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))
self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)
- self.cond_hint = self.cond_hint.to(device=self.device, dtype=dtype)
+ self.cond_hint = self.cond_hint.to(device=x_noisy.device, dtype=dtype)
if x_noisy.shape[0] != self.cond_hint.shape[0]:
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
@@ -341,8 +337,8 @@ class ControlLoraOps:
class ControlLora(ControlNet):
- def __init__(self, control_weights, global_average_pooling=False, device=None, model_options={}): #TODO? model_options
- ControlBase.__init__(self, device)
+ def __init__(self, control_weights, global_average_pooling=False, model_options={}): #TODO? model_options
+ ControlBase.__init__(self)
self.control_weights = control_weights
self.global_average_pooling = global_average_pooling
self.extra_conds += ["y"]
@@ -662,12 +658,15 @@ def load_controlnet(ckpt_path, model=None, model_options={}):
class T2IAdapter(ControlBase):
def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None):
- super().__init__(device)
+ super().__init__()
self.t2i_model = t2i_model
self.channels_in = channels_in
self.control_input = None
self.compression_ratio = compression_ratio
self.upscale_algorithm = upscale_algorithm
+ if device is None:
+ device = comfy.model_management.get_torch_device()
+ self.device = device
def scale_image_to(self, width, height):
unshuffle_amount = self.t2i_model.unshuffle_amount
diff --git a/comfy/float.py b/comfy/float.py
index 0c8f1d6c..521316fd 100644
--- a/comfy/float.py
+++ b/comfy/float.py
@@ -41,6 +41,8 @@ def manual_stochastic_round_to_float8(x, dtype, generator=None):
(2.0 ** (-EXPONENT_BIAS + 1)) * abs_x
)
+ inf = torch.finfo(dtype)
+ torch.clamp(sign, min=inf.min, max=inf.max, out=sign)
return sign
diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py
index d5fc05cb..605b1092 100644
--- a/comfy/k_diffusion/sampling.py
+++ b/comfy/k_diffusion/sampling.py
@@ -166,6 +166,8 @@ def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None,
@torch.no_grad()
def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
+ if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
+ return sample_euler_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
"""Ancestral sampling with Euler method steps."""
extra_args = {} if extra_args is None else extra_args
noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
@@ -183,6 +185,29 @@ def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, dis
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
return x
+@torch.no_grad()
+def sample_euler_ancestral_RF(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1., noise_sampler=None):
+ """Ancestral sampling with Euler method steps."""
+ extra_args = {} if extra_args is None else extra_args
+ noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+ s_in = x.new_ones([x.shape[0]])
+ for i in trange(len(sigmas) - 1, disable=disable):
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
+ # sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
+ downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
+ sigma_down = sigmas[i+1] * downstep_ratio
+ alpha_ip1 = 1 - sigmas[i+1]
+ alpha_down = 1 - sigma_down
+ renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5
+ if callback is not None:
+ callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+ # Euler method
+ sigma_down_i_ratio = sigma_down / sigmas[i]
+ x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * denoised
+ if sigmas[i + 1] > 0 and eta > 0:
+ x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
+ return x
@torch.no_grad()
def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
@@ -1192,7 +1217,6 @@ def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disabl
d = to_d(x, sigma_hat, temp[0])
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
- dt = sigmas[i + 1] - sigma_hat
# Euler method
x = denoised + d * sigmas[i + 1]
return x
@@ -1219,7 +1243,6 @@ def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=No
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
d = to_d(x, sigmas[i], temp[0])
# Euler method
- dt = sigma_down - sigmas[i]
x = denoised + d * sigma_down
if sigmas[i + 1] > 0:
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
@@ -1250,7 +1273,6 @@ def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback
if sigma_down == 0:
# Euler method
d = to_d(x, sigmas[i], temp[0])
- dt = sigma_down - sigmas[i]
x = denoised + d * sigma_down
else:
# DPM-Solver++(2S)
@@ -1298,4 +1320,4 @@ def sample_dpmpp_2m_cfg_pp(model, x, sigmas, extra_args=None, callback=None, dis
denoised_mix = -torch.exp(-h) * uncond_denoised - torch.expm1(-h) * (1 / (2 * r)) * (denoised - old_uncond_denoised)
x = denoised + denoised_mix + torch.exp(-h) * x
old_uncond_denoised = uncond_denoised
- return x
\ No newline at end of file
+ return x
diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py
index 191c7091..a48f60c7 100644
--- a/comfy/latent_formats.py
+++ b/comfy/latent_formats.py
@@ -175,3 +175,30 @@ class Flux(SD3):
def process_out(self, latent):
return (latent / self.scale_factor) + self.shift_factor
+
+class Mochi(LatentFormat):
+ latent_channels = 12
+
+ def __init__(self):
+ self.scale_factor = 1.0
+ self.latents_mean = torch.tensor([-0.06730895953510081, -0.038011381506090416, -0.07477820912866141,
+ -0.05565264470995561, 0.012767231469026969, -0.04703542746246419,
+ 0.043896967884726704, -0.09346305707025976, -0.09918314763016893,
+ -0.008729793427399178, -0.011931556316503654, -0.0321993391887285]).view(1, self.latent_channels, 1, 1, 1)
+ self.latents_std = torch.tensor([0.9263795028493863, 0.9248894543193766, 0.9393059390890617,
+ 0.959253732819592, 0.8244560132752793, 0.917259975397747,
+ 0.9294154431013696, 1.3720942357788521, 0.881393668867029,
+ 0.9168315692124348, 0.9185249279345552, 0.9274757570805041]).view(1, self.latent_channels, 1, 1, 1)
+
+ self.latent_rgb_factors = None #TODO
+ self.taesd_decoder_name = None #TODO
+
+ def process_in(self, latent):
+ latents_mean = self.latents_mean.to(latent.device, latent.dtype)
+ latents_std = self.latents_std.to(latent.device, latent.dtype)
+ return (latent - latents_mean) * self.scale_factor / latents_std
+
+ def process_out(self, latent):
+ latents_mean = self.latents_mean.to(latent.device, latent.dtype)
+ latents_std = self.latents_std.to(latent.device, latent.dtype)
+ return latent * latents_std / self.scale_factor + latents_mean
diff --git a/comfy/ldm/common_dit.py b/comfy/ldm/common_dit.py
index 5aebaf9e..cb6b7414 100644
--- a/comfy/ldm/common_dit.py
+++ b/comfy/ldm/common_dit.py
@@ -13,9 +13,15 @@ try:
except:
rms_norm_torch = None
-def rms_norm(x, weight, eps=1e-6):
+def rms_norm(x, weight=None, eps=1e-6):
if rms_norm_torch is not None and not (torch.jit.is_tracing() or torch.jit.is_scripting()):
- return rms_norm_torch(x, weight.shape, weight=comfy.ops.cast_to(weight, dtype=x.dtype, device=x.device), eps=eps)
+ if weight is None:
+ return rms_norm_torch(x, (x.shape[-1],), eps=eps)
+ else:
+ return rms_norm_torch(x, weight.shape, weight=comfy.ops.cast_to(weight, dtype=x.dtype, device=x.device), eps=eps)
else:
- rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + eps)
- return (x * rrms) * comfy.ops.cast_to(weight, dtype=x.dtype, device=x.device)
+ r = x * torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + eps)
+ if weight is None:
+ return r
+ else:
+ return r * comfy.ops.cast_to(weight, dtype=x.dtype, device=x.device)
diff --git a/comfy/ldm/genmo/joint_model/asymm_models_joint.py b/comfy/ldm/genmo/joint_model/asymm_models_joint.py
new file mode 100644
index 00000000..c36a0006
--- /dev/null
+++ b/comfy/ldm/genmo/joint_model/asymm_models_joint.py
@@ -0,0 +1,541 @@
+#original code from https://github.com/genmoai/models under apache 2.0 license
+#adapted to ComfyUI
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+# from flash_attn import flash_attn_varlen_qkvpacked_func
+from comfy.ldm.modules.attention import optimized_attention
+
+from .layers import (
+ FeedForward,
+ PatchEmbed,
+ RMSNorm,
+ TimestepEmbedder,
+)
+
+from .rope_mixed import (
+ compute_mixed_rotation,
+ create_position_matrix,
+)
+from .temporal_rope import apply_rotary_emb_qk_real
+from .utils import (
+ AttentionPool,
+ modulate,
+)
+
+import comfy.ldm.common_dit
+import comfy.ops
+
+
+def modulated_rmsnorm(x, scale, eps=1e-6):
+ # Normalize and modulate
+ x_normed = comfy.ldm.common_dit.rms_norm(x, eps=eps)
+ x_modulated = x_normed * (1 + scale.unsqueeze(1))
+
+ return x_modulated
+
+
+def residual_tanh_gated_rmsnorm(x, x_res, gate, eps=1e-6):
+ # Apply tanh to gate
+ tanh_gate = torch.tanh(gate).unsqueeze(1)
+
+ # Normalize and apply gated scaling
+ x_normed = comfy.ldm.common_dit.rms_norm(x_res, eps=eps) * tanh_gate
+
+ # Apply residual connection
+ output = x + x_normed
+
+ return output
+
+class AsymmetricAttention(nn.Module):
+ def __init__(
+ self,
+ dim_x: int,
+ dim_y: int,
+ num_heads: int = 8,
+ qkv_bias: bool = True,
+ qk_norm: bool = False,
+ attn_drop: float = 0.0,
+ update_y: bool = True,
+ out_bias: bool = True,
+ attend_to_padding: bool = False,
+ softmax_scale: Optional[float] = None,
+ device: Optional[torch.device] = None,
+ dtype=None,
+ operations=None,
+ ):
+ super().__init__()
+ self.dim_x = dim_x
+ self.dim_y = dim_y
+ self.num_heads = num_heads
+ self.head_dim = dim_x // num_heads
+ self.attn_drop = attn_drop
+ self.update_y = update_y
+ self.attend_to_padding = attend_to_padding
+ self.softmax_scale = softmax_scale
+ if dim_x % num_heads != 0:
+ raise ValueError(
+ f"dim_x={dim_x} should be divisible by num_heads={num_heads}"
+ )
+
+ # Input layers.
+ self.qkv_bias = qkv_bias
+ self.qkv_x = operations.Linear(dim_x, 3 * dim_x, bias=qkv_bias, device=device, dtype=dtype)
+ # Project text features to match visual features (dim_y -> dim_x)
+ self.qkv_y = operations.Linear(dim_y, 3 * dim_x, bias=qkv_bias, device=device, dtype=dtype)
+
+ # Query and key normalization for stability.
+ assert qk_norm
+ self.q_norm_x = RMSNorm(self.head_dim, device=device, dtype=dtype)
+ self.k_norm_x = RMSNorm(self.head_dim, device=device, dtype=dtype)
+ self.q_norm_y = RMSNorm(self.head_dim, device=device, dtype=dtype)
+ self.k_norm_y = RMSNorm(self.head_dim, device=device, dtype=dtype)
+
+ # Output layers. y features go back down from dim_x -> dim_y.
+ self.proj_x = operations.Linear(dim_x, dim_x, bias=out_bias, device=device, dtype=dtype)
+ self.proj_y = (
+ operations.Linear(dim_x, dim_y, bias=out_bias, device=device, dtype=dtype)
+ if update_y
+ else nn.Identity()
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor, # (B, N, dim_x)
+ y: torch.Tensor, # (B, L, dim_y)
+ scale_x: torch.Tensor, # (B, dim_x), modulation for pre-RMSNorm.
+ scale_y: torch.Tensor, # (B, dim_y), modulation for pre-RMSNorm.
+ crop_y,
+ **rope_rotation,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ rope_cos = rope_rotation.get("rope_cos")
+ rope_sin = rope_rotation.get("rope_sin")
+ # Pre-norm for visual features
+ x = modulated_rmsnorm(x, scale_x) # (B, M, dim_x) where M = N / cp_group_size
+
+ # Process visual features
+ # qkv_x = self.qkv_x(x) # (B, M, 3 * dim_x)
+ # assert qkv_x.dtype == torch.bfloat16
+ # qkv_x = all_to_all_collect_tokens(
+ # qkv_x, self.num_heads
+ # ) # (3, B, N, local_h, head_dim)
+
+ # Process text features
+ y = modulated_rmsnorm(y, scale_y) # (B, L, dim_y)
+ q_y, k_y, v_y = self.qkv_y(y).view(y.shape[0], y.shape[1], 3, self.num_heads, -1).unbind(2) # (B, N, local_h, head_dim)
+
+ q_y = self.q_norm_y(q_y)
+ k_y = self.k_norm_y(k_y)
+
+ # Split qkv_x into q, k, v
+ q_x, k_x, v_x = self.qkv_x(x).view(x.shape[0], x.shape[1], 3, self.num_heads, -1).unbind(2) # (B, N, local_h, head_dim)
+ q_x = self.q_norm_x(q_x)
+ q_x = apply_rotary_emb_qk_real(q_x, rope_cos, rope_sin)
+ k_x = self.k_norm_x(k_x)
+ k_x = apply_rotary_emb_qk_real(k_x, rope_cos, rope_sin)
+
+ q = torch.cat([q_x, q_y[:, :crop_y]], dim=1).transpose(1, 2)
+ k = torch.cat([k_x, k_y[:, :crop_y]], dim=1).transpose(1, 2)
+ v = torch.cat([v_x, v_y[:, :crop_y]], dim=1).transpose(1, 2)
+
+ xy = optimized_attention(q,
+ k,
+ v, self.num_heads, skip_reshape=True)
+
+ x, y = torch.tensor_split(xy, (q_x.shape[1],), dim=1)
+ x = self.proj_x(x)
+ o = torch.zeros(y.shape[0], q_y.shape[1], y.shape[-1], device=y.device, dtype=y.dtype)
+ o[:, :y.shape[1]] = y
+
+ y = self.proj_y(o)
+ # print("ox", x)
+ # print("oy", y)
+ return x, y
+
+
+class AsymmetricJointBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size_x: int,
+ hidden_size_y: int,
+ num_heads: int,
+ *,
+ mlp_ratio_x: float = 8.0, # Ratio of hidden size to d_model for MLP for visual tokens.
+ mlp_ratio_y: float = 4.0, # Ratio of hidden size to d_model for MLP for text tokens.
+ update_y: bool = True, # Whether to update text tokens in this block.
+ device: Optional[torch.device] = None,
+ dtype=None,
+ operations=None,
+ **block_kwargs,
+ ):
+ super().__init__()
+ self.update_y = update_y
+ self.hidden_size_x = hidden_size_x
+ self.hidden_size_y = hidden_size_y
+ self.mod_x = operations.Linear(hidden_size_x, 4 * hidden_size_x, device=device, dtype=dtype)
+ if self.update_y:
+ self.mod_y = operations.Linear(hidden_size_x, 4 * hidden_size_y, device=device, dtype=dtype)
+ else:
+ self.mod_y = operations.Linear(hidden_size_x, hidden_size_y, device=device, dtype=dtype)
+
+ # Self-attention:
+ self.attn = AsymmetricAttention(
+ hidden_size_x,
+ hidden_size_y,
+ num_heads=num_heads,
+ update_y=update_y,
+ device=device,
+ dtype=dtype,
+ operations=operations,
+ **block_kwargs,
+ )
+
+ # MLP.
+ mlp_hidden_dim_x = int(hidden_size_x * mlp_ratio_x)
+ assert mlp_hidden_dim_x == int(1536 * 8)
+ self.mlp_x = FeedForward(
+ in_features=hidden_size_x,
+ hidden_size=mlp_hidden_dim_x,
+ multiple_of=256,
+ ffn_dim_multiplier=None,
+ device=device,
+ dtype=dtype,
+ operations=operations,
+ )
+
+ # MLP for text not needed in last block.
+ if self.update_y:
+ mlp_hidden_dim_y = int(hidden_size_y * mlp_ratio_y)
+ self.mlp_y = FeedForward(
+ in_features=hidden_size_y,
+ hidden_size=mlp_hidden_dim_y,
+ multiple_of=256,
+ ffn_dim_multiplier=None,
+ device=device,
+ dtype=dtype,
+ operations=operations,
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ c: torch.Tensor,
+ y: torch.Tensor,
+ **attn_kwargs,
+ ):
+ """Forward pass of a block.
+
+ Args:
+ x: (B, N, dim) tensor of visual tokens
+ c: (B, dim) tensor of conditioned features
+ y: (B, L, dim) tensor of text tokens
+ num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens
+
+ Returns:
+ x: (B, N, dim) tensor of visual tokens after block
+ y: (B, L, dim) tensor of text tokens after block
+ """
+ N = x.size(1)
+
+ c = F.silu(c)
+ mod_x = self.mod_x(c)
+ scale_msa_x, gate_msa_x, scale_mlp_x, gate_mlp_x = mod_x.chunk(4, dim=1)
+
+ mod_y = self.mod_y(c)
+ if self.update_y:
+ scale_msa_y, gate_msa_y, scale_mlp_y, gate_mlp_y = mod_y.chunk(4, dim=1)
+ else:
+ scale_msa_y = mod_y
+
+ # Self-attention block.
+ x_attn, y_attn = self.attn(
+ x,
+ y,
+ scale_x=scale_msa_x,
+ scale_y=scale_msa_y,
+ **attn_kwargs,
+ )
+
+ assert x_attn.size(1) == N
+ x = residual_tanh_gated_rmsnorm(x, x_attn, gate_msa_x)
+ if self.update_y:
+ y = residual_tanh_gated_rmsnorm(y, y_attn, gate_msa_y)
+
+ # MLP block.
+ x = self.ff_block_x(x, scale_mlp_x, gate_mlp_x)
+ if self.update_y:
+ y = self.ff_block_y(y, scale_mlp_y, gate_mlp_y)
+
+ return x, y
+
+ def ff_block_x(self, x, scale_x, gate_x):
+ x_mod = modulated_rmsnorm(x, scale_x)
+ x_res = self.mlp_x(x_mod)
+ x = residual_tanh_gated_rmsnorm(x, x_res, gate_x) # Sandwich norm
+ return x
+
+ def ff_block_y(self, y, scale_y, gate_y):
+ y_mod = modulated_rmsnorm(y, scale_y)
+ y_res = self.mlp_y(y_mod)
+ y = residual_tanh_gated_rmsnorm(y, y_res, gate_y) # Sandwich norm
+ return y
+
+
+class FinalLayer(nn.Module):
+ """
+ The final layer of DiT.
+ """
+
+ def __init__(
+ self,
+ hidden_size,
+ patch_size,
+ out_channels,
+ device: Optional[torch.device] = None,
+ dtype=None,
+ operations=None,
+ ):
+ super().__init__()
+ self.norm_final = operations.LayerNorm(
+ hidden_size, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype
+ )
+ self.mod = operations.Linear(hidden_size, 2 * hidden_size, device=device, dtype=dtype)
+ self.linear = operations.Linear(
+ hidden_size, patch_size * patch_size * out_channels, device=device, dtype=dtype
+ )
+
+ def forward(self, x, c):
+ c = F.silu(c)
+ shift, scale = self.mod(c).chunk(2, dim=1)
+ x = modulate(self.norm_final(x), shift, scale)
+ x = self.linear(x)
+ return x
+
+
+class AsymmDiTJoint(nn.Module):
+ """
+ Diffusion model with a Transformer backbone.
+
+ Ingests text embeddings instead of a label.
+ """
+
+ def __init__(
+ self,
+ *,
+ patch_size=2,
+ in_channels=4,
+ hidden_size_x=1152,
+ hidden_size_y=1152,
+ depth=48,
+ num_heads=16,
+ mlp_ratio_x=8.0,
+ mlp_ratio_y=4.0,
+ use_t5: bool = False,
+ t5_feat_dim: int = 4096,
+ t5_token_length: int = 256,
+ learn_sigma=True,
+ patch_embed_bias: bool = True,
+ timestep_mlp_bias: bool = True,
+ attend_to_padding: bool = False,
+ timestep_scale: Optional[float] = None,
+ use_extended_posenc: bool = False,
+ posenc_preserve_area: bool = False,
+ rope_theta: float = 10000.0,
+ image_model=None,
+ device: Optional[torch.device] = None,
+ dtype=None,
+ operations=None,
+ **block_kwargs,
+ ):
+ super().__init__()
+
+ self.dtype = dtype
+ self.learn_sigma = learn_sigma
+ self.in_channels = in_channels
+ self.out_channels = in_channels * 2 if learn_sigma else in_channels
+ self.patch_size = patch_size
+ self.num_heads = num_heads
+ self.hidden_size_x = hidden_size_x
+ self.hidden_size_y = hidden_size_y
+ self.head_dim = (
+ hidden_size_x // num_heads
+ ) # Head dimension and count is determined by visual.
+ self.attend_to_padding = attend_to_padding
+ self.use_extended_posenc = use_extended_posenc
+ self.posenc_preserve_area = posenc_preserve_area
+ self.use_t5 = use_t5
+ self.t5_token_length = t5_token_length
+ self.t5_feat_dim = t5_feat_dim
+ self.rope_theta = (
+ rope_theta # Scaling factor for frequency computation for temporal RoPE.
+ )
+
+ self.x_embedder = PatchEmbed(
+ patch_size=patch_size,
+ in_chans=in_channels,
+ embed_dim=hidden_size_x,
+ bias=patch_embed_bias,
+ dtype=dtype,
+ device=device,
+ operations=operations
+ )
+ # Conditionings
+ # Timestep
+ self.t_embedder = TimestepEmbedder(
+ hidden_size_x, bias=timestep_mlp_bias, timestep_scale=timestep_scale, dtype=dtype, device=device, operations=operations
+ )
+
+ if self.use_t5:
+ # Caption Pooling (T5)
+ self.t5_y_embedder = AttentionPool(
+ t5_feat_dim, num_heads=8, output_dim=hidden_size_x, dtype=dtype, device=device, operations=operations
+ )
+
+ # Dense Embedding Projection (T5)
+ self.t5_yproj = operations.Linear(
+ t5_feat_dim, hidden_size_y, bias=True, dtype=dtype, device=device
+ )
+
+ # Initialize pos_frequencies as an empty parameter.
+ self.pos_frequencies = nn.Parameter(
+ torch.empty(3, self.num_heads, self.head_dim // 2, dtype=dtype, device=device)
+ )
+
+ assert not self.attend_to_padding
+
+ # for depth 48:
+ # b = 0: AsymmetricJointBlock, update_y=True
+ # b = 1: AsymmetricJointBlock, update_y=True
+ # ...
+ # b = 46: AsymmetricJointBlock, update_y=True
+ # b = 47: AsymmetricJointBlock, update_y=False. No need to update text features.
+ blocks = []
+ for b in range(depth):
+ # Joint multi-modal block
+ update_y = b < depth - 1
+ block = AsymmetricJointBlock(
+ hidden_size_x,
+ hidden_size_y,
+ num_heads,
+ mlp_ratio_x=mlp_ratio_x,
+ mlp_ratio_y=mlp_ratio_y,
+ update_y=update_y,
+ attend_to_padding=attend_to_padding,
+ device=device,
+ dtype=dtype,
+ operations=operations,
+ **block_kwargs,
+ )
+
+ blocks.append(block)
+ self.blocks = nn.ModuleList(blocks)
+
+ self.final_layer = FinalLayer(
+ hidden_size_x, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations
+ )
+
+ def embed_x(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ x: (B, C=12, T, H, W) tensor of visual tokens
+
+ Returns:
+ x: (B, C=3072, N) tensor of visual tokens with positional embedding.
+ """
+ return self.x_embedder(x) # Convert BcTHW to BCN
+
+ def prepare(
+ self,
+ x: torch.Tensor,
+ sigma: torch.Tensor,
+ t5_feat: torch.Tensor,
+ t5_mask: torch.Tensor,
+ ):
+ """Prepare input and conditioning embeddings."""
+ # Visual patch embeddings with positional encoding.
+ T, H, W = x.shape[-3:]
+ pH, pW = H // self.patch_size, W // self.patch_size
+ x = self.embed_x(x) # (B, N, D), where N = T * H * W / patch_size ** 2
+ assert x.ndim == 3
+ B = x.size(0)
+
+
+ pH, pW = H // self.patch_size, W // self.patch_size
+ N = T * pH * pW
+ assert x.size(1) == N
+ pos = create_position_matrix(
+ T, pH=pH, pW=pW, device=x.device, dtype=torch.float32
+ ) # (N, 3)
+ rope_cos, rope_sin = compute_mixed_rotation(
+ freqs=comfy.ops.cast_to(self.pos_frequencies, dtype=x.dtype, device=x.device), pos=pos
+ ) # Each are (N, num_heads, dim // 2)
+
+ c_t = self.t_embedder(1 - sigma, out_dtype=x.dtype) # (B, D)
+
+ t5_y_pool = self.t5_y_embedder(t5_feat, t5_mask) # (B, D)
+
+ c = c_t + t5_y_pool
+
+ y_feat = self.t5_yproj(t5_feat) # (B, L, t5_feat_dim) --> (B, L, D)
+
+ return x, c, y_feat, rope_cos, rope_sin
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ timestep: torch.Tensor,
+ context: List[torch.Tensor],
+ attention_mask: List[torch.Tensor],
+ num_tokens=256,
+ packed_indices: Dict[str, torch.Tensor] = None,
+ rope_cos: torch.Tensor = None,
+ rope_sin: torch.Tensor = None,
+ control=None, **kwargs
+ ):
+ y_feat = context
+ y_mask = attention_mask
+ sigma = timestep
+ """Forward pass of DiT.
+
+ Args:
+ x: (B, C, T, H, W) tensor of spatial inputs (images or latent representations of images)
+ sigma: (B,) tensor of noise standard deviations
+ y_feat: List((B, L, y_feat_dim) tensor of caption token features. For SDXL text encoders: L=77, y_feat_dim=2048)
+ y_mask: List((B, L) boolean tensor indicating which tokens are not padding)
+ packed_indices: Dict with keys for Flash Attention. Result of compute_packed_indices.
+ """
+ B, _, T, H, W = x.shape
+
+ x, c, y_feat, rope_cos, rope_sin = self.prepare(
+ x, sigma, y_feat, y_mask
+ )
+ del y_mask
+
+ for i, block in enumerate(self.blocks):
+ x, y_feat = block(
+ x,
+ c,
+ y_feat,
+ rope_cos=rope_cos,
+ rope_sin=rope_sin,
+ crop_y=num_tokens,
+ ) # (B, M, D), (B, L, D)
+ del y_feat # Final layers don't use dense text features.
+
+ x = self.final_layer(x, c) # (B, M, patch_size ** 2 * out_channels)
+ x = rearrange(
+ x,
+ "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)",
+ T=T,
+ hp=H // self.patch_size,
+ wp=W // self.patch_size,
+ p1=self.patch_size,
+ p2=self.patch_size,
+ c=self.out_channels,
+ )
+
+ return -x
diff --git a/comfy/ldm/genmo/joint_model/layers.py b/comfy/ldm/genmo/joint_model/layers.py
new file mode 100644
index 00000000..51d97955
--- /dev/null
+++ b/comfy/ldm/genmo/joint_model/layers.py
@@ -0,0 +1,164 @@
+#original code from https://github.com/genmoai/models under apache 2.0 license
+#adapted to ComfyUI
+
+import collections.abc
+import math
+from itertools import repeat
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import comfy.ldm.common_dit
+
+
+# From PyTorch internals
+def _ntuple(n):
+ def parse(x):
+ if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+ return tuple(x)
+ return tuple(repeat(x, n))
+
+ return parse
+
+
+to_2tuple = _ntuple(2)
+
+
+class TimestepEmbedder(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ frequency_embedding_size: int = 256,
+ *,
+ bias: bool = True,
+ timestep_scale: Optional[float] = None,
+ dtype=None,
+ device=None,
+ operations=None,
+ ):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ operations.Linear(frequency_embedding_size, hidden_size, bias=bias, dtype=dtype, device=device),
+ nn.SiLU(),
+ operations.Linear(hidden_size, hidden_size, bias=bias, dtype=dtype, device=device),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+ self.timestep_scale = timestep_scale
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ half = dim // 2
+ freqs = torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
+ freqs.mul_(-math.log(max_period) / half).exp_()
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat(
+ [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+ )
+ return embedding
+
+ def forward(self, t, out_dtype):
+ if self.timestep_scale is not None:
+ t = t * self.timestep_scale
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype=out_dtype)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class FeedForward(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_size: int,
+ multiple_of: int,
+ ffn_dim_multiplier: Optional[float],
+ device: Optional[torch.device] = None,
+ dtype=None,
+ operations=None,
+ ):
+ super().__init__()
+ # keep parameter count and computation constant compared to standard FFN
+ hidden_size = int(2 * hidden_size / 3)
+ # custom dim factor multiplier
+ if ffn_dim_multiplier is not None:
+ hidden_size = int(ffn_dim_multiplier * hidden_size)
+ hidden_size = multiple_of * ((hidden_size + multiple_of - 1) // multiple_of)
+
+ self.hidden_dim = hidden_size
+ self.w1 = operations.Linear(in_features, 2 * hidden_size, bias=False, device=device, dtype=dtype)
+ self.w2 = operations.Linear(hidden_size, in_features, bias=False, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x, gate = self.w1(x).chunk(2, dim=-1)
+ x = self.w2(F.silu(x) * gate)
+ return x
+
+
+class PatchEmbed(nn.Module):
+ def __init__(
+ self,
+ patch_size: int = 16,
+ in_chans: int = 3,
+ embed_dim: int = 768,
+ norm_layer: Optional[Callable] = None,
+ flatten: bool = True,
+ bias: bool = True,
+ dynamic_img_pad: bool = False,
+ dtype=None,
+ device=None,
+ operations=None,
+ ):
+ super().__init__()
+ self.patch_size = to_2tuple(patch_size)
+ self.flatten = flatten
+ self.dynamic_img_pad = dynamic_img_pad
+
+ self.proj = operations.Conv2d(
+ in_chans,
+ embed_dim,
+ kernel_size=patch_size,
+ stride=patch_size,
+ bias=bias,
+ device=device,
+ dtype=dtype,
+ )
+ assert norm_layer is None
+ self.norm = (
+ norm_layer(embed_dim, device=device) if norm_layer else nn.Identity()
+ )
+
+ def forward(self, x):
+ B, _C, T, H, W = x.shape
+ if not self.dynamic_img_pad:
+ assert H % self.patch_size[0] == 0, f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
+ assert W % self.patch_size[1] == 0, f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
+ else:
+ pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
+ pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
+ x = F.pad(x, (0, pad_w, 0, pad_h))
+
+ x = rearrange(x, "B C T H W -> (B T) C H W", B=B, T=T)
+ x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size, padding_mode='circular')
+ x = self.proj(x)
+
+ # Flatten temporal and spatial dimensions.
+ if not self.flatten:
+ raise NotImplementedError("Must flatten output.")
+ x = rearrange(x, "(B T) C H W -> B (T H W) C", B=B, T=T)
+
+ x = self.norm(x)
+ return x
+
+
+class RMSNorm(torch.nn.Module):
+ def __init__(self, hidden_size, eps=1e-5, device=None, dtype=None):
+ super().__init__()
+ self.eps = eps
+ self.weight = torch.nn.Parameter(torch.empty(hidden_size, device=device, dtype=dtype))
+ self.register_parameter("bias", None)
+
+ def forward(self, x):
+ return comfy.ldm.common_dit.rms_norm(x, self.weight, self.eps)
diff --git a/comfy/ldm/genmo/joint_model/rope_mixed.py b/comfy/ldm/genmo/joint_model/rope_mixed.py
new file mode 100644
index 00000000..dee3fa21
--- /dev/null
+++ b/comfy/ldm/genmo/joint_model/rope_mixed.py
@@ -0,0 +1,88 @@
+#original code from https://github.com/genmoai/models under apache 2.0 license
+
+# import functools
+import math
+
+import torch
+
+
+def centers(start: float, stop, num, dtype=None, device=None):
+ """linspace through bin centers.
+
+ Args:
+ start (float): Start of the range.
+ stop (float): End of the range.
+ num (int): Number of points.
+ dtype (torch.dtype): Data type of the points.
+ device (torch.device): Device of the points.
+
+ Returns:
+ centers (Tensor): Centers of the bins. Shape: (num,).
+ """
+ edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device)
+ return (edges[:-1] + edges[1:]) / 2
+
+
+# @functools.lru_cache(maxsize=1)
+def create_position_matrix(
+ T: int,
+ pH: int,
+ pW: int,
+ device: torch.device,
+ dtype: torch.dtype,
+ *,
+ target_area: float = 36864,
+):
+ """
+ Args:
+ T: int - Temporal dimension
+ pH: int - Height dimension after patchify
+ pW: int - Width dimension after patchify
+
+ Returns:
+ pos: [T * pH * pW, 3] - position matrix
+ """
+ # Create 1D tensors for each dimension
+ t = torch.arange(T, dtype=dtype)
+
+ # Positionally interpolate to area 36864.
+ # (3072x3072 frame with 16x16 patches = 192x192 latents).
+ # This automatically scales rope positions when the resolution changes.
+ # We use a large target area so the model is more sensitive
+ # to changes in the learned pos_frequencies matrix.
+ scale = math.sqrt(target_area / (pW * pH))
+ w = centers(-pW * scale / 2, pW * scale / 2, pW)
+ h = centers(-pH * scale / 2, pH * scale / 2, pH)
+
+ # Use meshgrid to create 3D grids
+ grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
+
+ # Stack and reshape the grids.
+ pos = torch.stack([grid_t, grid_h, grid_w], dim=-1) # [T, pH, pW, 3]
+ pos = pos.view(-1, 3) # [T * pH * pW, 3]
+ pos = pos.to(dtype=dtype, device=device)
+
+ return pos
+
+
+def compute_mixed_rotation(
+ freqs: torch.Tensor,
+ pos: torch.Tensor,
+):
+ """
+ Project each 3-dim position into per-head, per-head-dim 1D frequencies.
+
+ Args:
+ freqs: [3, num_heads, num_freqs] - learned rotation frequency (for t, row, col) for each head position
+ pos: [N, 3] - position of each token
+ num_heads: int
+
+ Returns:
+ freqs_cos: [N, num_heads, num_freqs] - cosine components
+ freqs_sin: [N, num_heads, num_freqs] - sine components
+ """
+ assert freqs.ndim == 3
+ freqs_sum = torch.einsum("Nd,dhf->Nhf", pos.to(freqs), freqs)
+ freqs_cos = torch.cos(freqs_sum)
+ freqs_sin = torch.sin(freqs_sum)
+ return freqs_cos, freqs_sin
diff --git a/comfy/ldm/genmo/joint_model/temporal_rope.py b/comfy/ldm/genmo/joint_model/temporal_rope.py
new file mode 100644
index 00000000..88f5d6d2
--- /dev/null
+++ b/comfy/ldm/genmo/joint_model/temporal_rope.py
@@ -0,0 +1,34 @@
+#original code from https://github.com/genmoai/models under apache 2.0 license
+
+# Based on Llama3 Implementation.
+import torch
+
+
+def apply_rotary_emb_qk_real(
+ xqk: torch.Tensor,
+ freqs_cos: torch.Tensor,
+ freqs_sin: torch.Tensor,
+) -> torch.Tensor:
+ """
+ Apply rotary embeddings to input tensors using the given frequency tensor without complex numbers.
+
+ Args:
+ xqk (torch.Tensor): Query and/or Key tensors to apply rotary embeddings. Shape: (B, S, *, num_heads, D)
+ Can be either just query or just key, or both stacked along some batch or * dim.
+ freqs_cos (torch.Tensor): Precomputed cosine frequency tensor.
+ freqs_sin (torch.Tensor): Precomputed sine frequency tensor.
+
+ Returns:
+ torch.Tensor: The input tensor with rotary embeddings applied.
+ """
+ # Split the last dimension into even and odd parts
+ xqk_even = xqk[..., 0::2]
+ xqk_odd = xqk[..., 1::2]
+
+ # Apply rotation
+ cos_part = (xqk_even * freqs_cos - xqk_odd * freqs_sin).type_as(xqk)
+ sin_part = (xqk_even * freqs_sin + xqk_odd * freqs_cos).type_as(xqk)
+
+ # Interleave the results back into the original shape
+ out = torch.stack([cos_part, sin_part], dim=-1).flatten(-2)
+ return out
diff --git a/comfy/ldm/genmo/joint_model/utils.py b/comfy/ldm/genmo/joint_model/utils.py
new file mode 100644
index 00000000..41190242
--- /dev/null
+++ b/comfy/ldm/genmo/joint_model/utils.py
@@ -0,0 +1,102 @@
+#original code from https://github.com/genmoai/models under apache 2.0 license
+#adapted to ComfyUI
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.Tensor:
+ """
+ Pool tokens in x using mask.
+
+ NOTE: We assume x does not require gradients.
+
+ Args:
+ x: (B, L, D) tensor of tokens.
+ mask: (B, L) boolean tensor indicating which tokens are not padding.
+
+ Returns:
+ pooled: (B, D) tensor of pooled tokens.
+ """
+ assert x.size(1) == mask.size(1) # Expected mask to have same length as tokens.
+ assert x.size(0) == mask.size(0) # Expected mask to have same batch size as tokens.
+ mask = mask[:, :, None].to(dtype=x.dtype)
+ mask = mask / mask.sum(dim=1, keepdim=True).clamp(min=1)
+ pooled = (x * mask).sum(dim=1, keepdim=keepdim)
+ return pooled
+
+
+class AttentionPool(nn.Module):
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ output_dim: int = None,
+ device: Optional[torch.device] = None,
+ dtype=None,
+ operations=None,
+ ):
+ """
+ Args:
+ spatial_dim (int): Number of tokens in sequence length.
+ embed_dim (int): Dimensionality of input tokens.
+ num_heads (int): Number of attention heads.
+ output_dim (int): Dimensionality of output tokens. Defaults to embed_dim.
+ """
+ super().__init__()
+ self.num_heads = num_heads
+ self.to_kv = operations.Linear(embed_dim, 2 * embed_dim, device=device, dtype=dtype)
+ self.to_q = operations.Linear(embed_dim, embed_dim, device=device, dtype=dtype)
+ self.to_out = operations.Linear(embed_dim, output_dim or embed_dim, device=device, dtype=dtype)
+
+ def forward(self, x, mask):
+ """
+ Args:
+ x (torch.Tensor): (B, L, D) tensor of input tokens.
+ mask (torch.Tensor): (B, L) boolean tensor indicating which tokens are not padding.
+
+ NOTE: We assume x does not require gradients.
+
+ Returns:
+ x (torch.Tensor): (B, D) tensor of pooled tokens.
+ """
+ D = x.size(2)
+
+ # Construct attention mask, shape: (B, 1, num_queries=1, num_keys=1+L).
+ attn_mask = mask[:, None, None, :].bool() # (B, 1, 1, L).
+ attn_mask = F.pad(attn_mask, (1, 0), value=True) # (B, 1, 1, 1+L).
+
+ # Average non-padding token features. These will be used as the query.
+ x_pool = pool_tokens(x, mask, keepdim=True) # (B, 1, D)
+
+ # Concat pooled features to input sequence.
+ x = torch.cat([x_pool, x], dim=1) # (B, L+1, D)
+
+ # Compute queries, keys, values. Only the mean token is used to create a query.
+ kv = self.to_kv(x) # (B, L+1, 2 * D)
+ q = self.to_q(x[:, 0]) # (B, D)
+
+ # Extract heads.
+ head_dim = D // self.num_heads
+ kv = kv.unflatten(2, (2, self.num_heads, head_dim)) # (B, 1+L, 2, H, head_dim)
+ kv = kv.transpose(1, 3) # (B, H, 2, 1+L, head_dim)
+ k, v = kv.unbind(2) # (B, H, 1+L, head_dim)
+ q = q.unflatten(1, (self.num_heads, head_dim)) # (B, H, head_dim)
+ q = q.unsqueeze(2) # (B, H, 1, head_dim)
+
+ # Compute attention.
+ x = F.scaled_dot_product_attention(
+ q, k, v, attn_mask=attn_mask, dropout_p=0.0
+ ) # (B, H, 1, head_dim)
+
+ # Concatenate heads and run output.
+ x = x.squeeze(2).flatten(1, 2) # (B, D = H * head_dim)
+ x = self.to_out(x)
+ return x
diff --git a/comfy/ldm/genmo/vae/model.py b/comfy/ldm/genmo/vae/model.py
new file mode 100644
index 00000000..e44c08a4
--- /dev/null
+++ b/comfy/ldm/genmo/vae/model.py
@@ -0,0 +1,480 @@
+#original code from https://github.com/genmoai/models under apache 2.0 license
+#adapted to ComfyUI
+
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+
+# import mochi_preview.dit.joint_model.context_parallel as cp
+# from mochi_preview.vae.cp_conv import cp_pass_frames, gather_all_frames
+
+
+def cast_tuple(t, length=1):
+ return t if isinstance(t, tuple) else ((t,) * length)
+
+
+class GroupNormSpatial(ops.GroupNorm):
+ """
+ GroupNorm applied per-frame.
+ """
+
+ def forward(self, x: torch.Tensor, *, chunk_size: int = 8):
+ B, C, T, H, W = x.shape
+ x = rearrange(x, "B C T H W -> (B T) C H W")
+ # Run group norm in chunks.
+ output = torch.empty_like(x)
+ for b in range(0, B * T, chunk_size):
+ output[b : b + chunk_size] = super().forward(x[b : b + chunk_size])
+ return rearrange(output, "(B T) C H W -> B C T H W", B=B, T=T)
+
+class PConv3d(ops.Conv3d):
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ kernel_size: Union[int, Tuple[int, int, int]],
+ stride: Union[int, Tuple[int, int, int]],
+ causal: bool = True,
+ context_parallel: bool = True,
+ **kwargs,
+ ):
+ self.causal = causal
+ self.context_parallel = context_parallel
+ kernel_size = cast_tuple(kernel_size, 3)
+ stride = cast_tuple(stride, 3)
+ height_pad = (kernel_size[1] - 1) // 2
+ width_pad = (kernel_size[2] - 1) // 2
+
+ super().__init__(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ dilation=(1, 1, 1),
+ padding=(0, height_pad, width_pad),
+ **kwargs,
+ )
+
+ def forward(self, x: torch.Tensor):
+ # Compute padding amounts.
+ context_size = self.kernel_size[0] - 1
+ if self.causal:
+ pad_front = context_size
+ pad_back = 0
+ else:
+ pad_front = context_size // 2
+ pad_back = context_size - pad_front
+
+ # Apply padding.
+ assert self.padding_mode == "replicate" # DEBUG
+ mode = "constant" if self.padding_mode == "zeros" else self.padding_mode
+ x = F.pad(x, (0, 0, 0, 0, pad_front, pad_back), mode=mode)
+ return super().forward(x)
+
+
+class Conv1x1(ops.Linear):
+ """*1x1 Conv implemented with a linear layer."""
+
+ def __init__(self, in_features: int, out_features: int, *args, **kwargs):
+ super().__init__(in_features, out_features, *args, **kwargs)
+
+ def forward(self, x: torch.Tensor):
+ """Forward pass.
+
+ Args:
+ x: Input tensor. Shape: [B, C, *] or [B, *, C].
+
+ Returns:
+ x: Output tensor. Shape: [B, C', *] or [B, *, C'].
+ """
+ x = x.movedim(1, -1)
+ x = super().forward(x)
+ x = x.movedim(-1, 1)
+ return x
+
+
+class DepthToSpaceTime(nn.Module):
+ def __init__(
+ self,
+ temporal_expansion: int,
+ spatial_expansion: int,
+ ):
+ super().__init__()
+ self.temporal_expansion = temporal_expansion
+ self.spatial_expansion = spatial_expansion
+
+ # When printed, this module should show the temporal and spatial expansion factors.
+ def extra_repr(self):
+ return f"texp={self.temporal_expansion}, sexp={self.spatial_expansion}"
+
+ def forward(self, x: torch.Tensor):
+ """Forward pass.
+
+ Args:
+ x: Input tensor. Shape: [B, C, T, H, W].
+
+ Returns:
+ x: Rearranged tensor. Shape: [B, C/(st*s*s), T*st, H*s, W*s].
+ """
+ x = rearrange(
+ x,
+ "B (C st sh sw) T H W -> B C (T st) (H sh) (W sw)",
+ st=self.temporal_expansion,
+ sh=self.spatial_expansion,
+ sw=self.spatial_expansion,
+ )
+
+ # cp_rank, _ = cp.get_cp_rank_size()
+ if self.temporal_expansion > 1: # and cp_rank == 0:
+ # Drop the first self.temporal_expansion - 1 frames.
+ # This is because we always want the 3x3x3 conv filter to only apply
+ # to the first frame, and the first frame doesn't need to be repeated.
+ assert all(x.shape)
+ x = x[:, :, self.temporal_expansion - 1 :]
+ assert all(x.shape)
+
+ return x
+
+
+def norm_fn(
+ in_channels: int,
+ affine: bool = True,
+):
+ return GroupNormSpatial(affine=affine, num_groups=32, num_channels=in_channels)
+
+
+class ResBlock(nn.Module):
+ """Residual block that preserves the spatial dimensions."""
+
+ def __init__(
+ self,
+ channels: int,
+ *,
+ affine: bool = True,
+ attn_block: Optional[nn.Module] = None,
+ padding_mode: str = "replicate",
+ causal: bool = True,
+ ):
+ super().__init__()
+ self.channels = channels
+
+ assert causal
+ self.stack = nn.Sequential(
+ norm_fn(channels, affine=affine),
+ nn.SiLU(inplace=True),
+ PConv3d(
+ in_channels=channels,
+ out_channels=channels,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ padding_mode=padding_mode,
+ bias=True,
+ # causal=causal,
+ ),
+ norm_fn(channels, affine=affine),
+ nn.SiLU(inplace=True),
+ PConv3d(
+ in_channels=channels,
+ out_channels=channels,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ padding_mode=padding_mode,
+ bias=True,
+ # causal=causal,
+ ),
+ )
+
+ self.attn_block = attn_block if attn_block else nn.Identity()
+
+ def forward(self, x: torch.Tensor):
+ """Forward pass.
+
+ Args:
+ x: Input tensor. Shape: [B, C, T, H, W].
+ """
+ residual = x
+ x = self.stack(x)
+ x = x + residual
+ del residual
+
+ return self.attn_block(x)
+
+
+class CausalUpsampleBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_res_blocks: int,
+ *,
+ temporal_expansion: int = 2,
+ spatial_expansion: int = 2,
+ **block_kwargs,
+ ):
+ super().__init__()
+
+ blocks = []
+ for _ in range(num_res_blocks):
+ blocks.append(block_fn(in_channels, **block_kwargs))
+ self.blocks = nn.Sequential(*blocks)
+
+ self.temporal_expansion = temporal_expansion
+ self.spatial_expansion = spatial_expansion
+
+ # Change channels in the final convolution layer.
+ self.proj = Conv1x1(
+ in_channels,
+ out_channels * temporal_expansion * (spatial_expansion**2),
+ )
+
+ self.d2st = DepthToSpaceTime(
+ temporal_expansion=temporal_expansion, spatial_expansion=spatial_expansion
+ )
+
+ def forward(self, x):
+ x = self.blocks(x)
+ x = self.proj(x)
+ x = self.d2st(x)
+ return x
+
+
+def block_fn(channels, *, has_attention: bool = False, **block_kwargs):
+ assert has_attention is False #NOTE: if this is ever true add back the attention code.
+
+ attn_block = None #AttentionBlock(channels) if has_attention else None
+
+ return ResBlock(
+ channels, affine=True, attn_block=attn_block, **block_kwargs
+ )
+
+
+class DownsampleBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_res_blocks,
+ *,
+ temporal_reduction=2,
+ spatial_reduction=2,
+ **block_kwargs,
+ ):
+ """
+ Downsample block for the VAE encoder.
+
+ Args:
+ in_channels: Number of input channels.
+ out_channels: Number of output channels.
+ num_res_blocks: Number of residual blocks.
+ temporal_reduction: Temporal reduction factor.
+ spatial_reduction: Spatial reduction factor.
+ """
+ super().__init__()
+ layers = []
+
+ # Change the channel count in the strided convolution.
+ # This lets the ResBlock have uniform channel count,
+ # as in ConvNeXt.
+ assert in_channels != out_channels
+ layers.append(
+ PConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=(temporal_reduction, spatial_reduction, spatial_reduction),
+ stride=(temporal_reduction, spatial_reduction, spatial_reduction),
+ padding_mode="replicate",
+ bias=True,
+ )
+ )
+
+ for _ in range(num_res_blocks):
+ layers.append(block_fn(out_channels, **block_kwargs))
+
+ self.layers = nn.Sequential(*layers)
+
+ def forward(self, x):
+ return self.layers(x)
+
+
+def add_fourier_features(inputs: torch.Tensor, start=6, stop=8, step=1):
+ num_freqs = (stop - start) // step
+ assert inputs.ndim == 5
+ C = inputs.size(1)
+
+ # Create Base 2 Fourier features.
+ freqs = torch.arange(start, stop, step, dtype=inputs.dtype, device=inputs.device)
+ assert num_freqs == len(freqs)
+ w = torch.pow(2.0, freqs) * (2 * torch.pi) # [num_freqs]
+ C = inputs.shape[1]
+ w = w.repeat(C)[None, :, None, None, None] # [1, C * num_freqs, 1, 1, 1]
+
+ # Interleaved repeat of input channels to match w.
+ h = inputs.repeat_interleave(num_freqs, dim=1) # [B, C * num_freqs, T, H, W]
+ # Scale channels by frequency.
+ h = w * h
+
+ return torch.cat(
+ [
+ inputs,
+ torch.sin(h),
+ torch.cos(h),
+ ],
+ dim=1,
+ )
+
+
+class FourierFeatures(nn.Module):
+ def __init__(self, start: int = 6, stop: int = 8, step: int = 1):
+ super().__init__()
+ self.start = start
+ self.stop = stop
+ self.step = step
+
+ def forward(self, inputs):
+ """Add Fourier features to inputs.
+
+ Args:
+ inputs: Input tensor. Shape: [B, C, T, H, W]
+
+ Returns:
+ h: Output tensor. Shape: [B, (1 + 2 * num_freqs) * C, T, H, W]
+ """
+ return add_fourier_features(inputs, self.start, self.stop, self.step)
+
+
+class Decoder(nn.Module):
+ def __init__(
+ self,
+ *,
+ out_channels: int = 3,
+ latent_dim: int,
+ base_channels: int,
+ channel_multipliers: List[int],
+ num_res_blocks: List[int],
+ temporal_expansions: Optional[List[int]] = None,
+ spatial_expansions: Optional[List[int]] = None,
+ has_attention: List[bool],
+ output_norm: bool = True,
+ nonlinearity: str = "silu",
+ output_nonlinearity: str = "silu",
+ causal: bool = True,
+ **block_kwargs,
+ ):
+ super().__init__()
+ self.input_channels = latent_dim
+ self.base_channels = base_channels
+ self.channel_multipliers = channel_multipliers
+ self.num_res_blocks = num_res_blocks
+ self.output_nonlinearity = output_nonlinearity
+ assert nonlinearity == "silu"
+ assert causal
+
+ ch = [mult * base_channels for mult in channel_multipliers]
+ self.num_up_blocks = len(ch) - 1
+ assert len(num_res_blocks) == self.num_up_blocks + 2
+
+ blocks = []
+
+ first_block = [
+ nn.Conv3d(latent_dim, ch[-1], kernel_size=(1, 1, 1))
+ ] # Input layer.
+ # First set of blocks preserve channel count.
+ for _ in range(num_res_blocks[-1]):
+ first_block.append(
+ block_fn(
+ ch[-1],
+ has_attention=has_attention[-1],
+ causal=causal,
+ **block_kwargs,
+ )
+ )
+ blocks.append(nn.Sequential(*first_block))
+
+ assert len(temporal_expansions) == len(spatial_expansions) == self.num_up_blocks
+ assert len(num_res_blocks) == len(has_attention) == self.num_up_blocks + 2
+
+ upsample_block_fn = CausalUpsampleBlock
+
+ for i in range(self.num_up_blocks):
+ block = upsample_block_fn(
+ ch[-i - 1],
+ ch[-i - 2],
+ num_res_blocks=num_res_blocks[-i - 2],
+ has_attention=has_attention[-i - 2],
+ temporal_expansion=temporal_expansions[-i - 1],
+ spatial_expansion=spatial_expansions[-i - 1],
+ causal=causal,
+ **block_kwargs,
+ )
+ blocks.append(block)
+
+ assert not output_norm
+
+ # Last block. Preserve channel count.
+ last_block = []
+ for _ in range(num_res_blocks[0]):
+ last_block.append(
+ block_fn(
+ ch[0], has_attention=has_attention[0], causal=causal, **block_kwargs
+ )
+ )
+ blocks.append(nn.Sequential(*last_block))
+
+ self.blocks = nn.ModuleList(blocks)
+ self.output_proj = Conv1x1(ch[0], out_channels)
+
+ def forward(self, x):
+ """Forward pass.
+
+ Args:
+ x: Latent tensor. Shape: [B, input_channels, t, h, w]. Scaled [-1, 1].
+
+ Returns:
+ x: Reconstructed video tensor. Shape: [B, C, T, H, W]. Scaled to [-1, 1].
+ T + 1 = (t - 1) * 4.
+ H = h * 16, W = w * 16.
+ """
+ for block in self.blocks:
+ x = block(x)
+
+ if self.output_nonlinearity == "silu":
+ x = F.silu(x, inplace=not self.training)
+ else:
+ assert (
+ not self.output_nonlinearity
+ ) # StyleGAN3 omits the to-RGB nonlinearity.
+
+ return self.output_proj(x).contiguous()
+
+
+class VideoVAE(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.encoder = None #TODO once the model releases
+ self.decoder = Decoder(
+ out_channels=3,
+ base_channels=128,
+ channel_multipliers=[1, 2, 4, 6],
+ temporal_expansions=[1, 2, 3],
+ spatial_expansions=[2, 2, 2],
+ num_res_blocks=[3, 3, 4, 6, 3],
+ latent_dim=12,
+ has_attention=[False, False, False, False, False],
+ padding_mode="replicate",
+ output_norm=False,
+ nonlinearity="silu",
+ output_nonlinearity="silu",
+ causal=True,
+ )
+
+ def encode(self, x):
+ return self.encoder(x)
+
+ def decode(self, x):
+ return self.decoder(x)
diff --git a/comfy/ldm/modules/diffusionmodules/mmdit.py b/comfy/ldm/modules/diffusionmodules/mmdit.py
index 759788a9..6f8f506c 100644
--- a/comfy/ldm/modules/diffusionmodules/mmdit.py
+++ b/comfy/ldm/modules/diffusionmodules/mmdit.py
@@ -1,11 +1,11 @@
import logging
import math
-from typing import Dict, Optional
+from typing import Dict, Optional, List
import numpy as np
import torch
import torch.nn as nn
-from .. import attention
+from ..attention import optimized_attention
from einops import rearrange, repeat
from .util import timestep_embedding
import comfy.ops
@@ -97,7 +97,7 @@ class PatchEmbed(nn.Module):
self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
def forward(self, x):
- B, C, H, W = x.shape
+ # B, C, H, W = x.shape
# if self.img_size is not None:
# if self.strict_img_size:
# _assert(H == self.img_size[0], f"Input height ({H}) doesn't match model ({self.img_size[0]}).")
@@ -266,8 +266,6 @@ def split_qkv(qkv, head_dim):
qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, head_dim).movedim(2, 0)
return qkv[0], qkv[1], qkv[2]
-def optimized_attention(qkv, num_heads):
- return attention.optimized_attention(qkv[0], qkv[1], qkv[2], num_heads)
class SelfAttention(nn.Module):
ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
@@ -326,9 +324,9 @@ class SelfAttention(nn.Module):
return x
def forward(self, x: torch.Tensor) -> torch.Tensor:
- qkv = self.pre_attention(x)
+ q, k, v = self.pre_attention(x)
x = optimized_attention(
- qkv, num_heads=self.num_heads
+ q, k, v, heads=self.num_heads
)
x = self.post_attention(x)
return x
@@ -417,6 +415,7 @@ class DismantledBlock(nn.Module):
scale_mod_only: bool = False,
swiglu: bool = False,
qk_norm: Optional[str] = None,
+ x_block_self_attn: bool = False,
dtype=None,
device=None,
operations=None,
@@ -440,6 +439,24 @@ class DismantledBlock(nn.Module):
device=device,
operations=operations
)
+ if x_block_self_attn:
+ assert not pre_only
+ assert not scale_mod_only
+ self.x_block_self_attn = True
+ self.attn2 = SelfAttention(
+ dim=hidden_size,
+ num_heads=num_heads,
+ qkv_bias=qkv_bias,
+ attn_mode=attn_mode,
+ pre_only=False,
+ qk_norm=qk_norm,
+ rmsnorm=rmsnorm,
+ dtype=dtype,
+ device=device,
+ operations=operations
+ )
+ else:
+ self.x_block_self_attn = False
if not pre_only:
if not rmsnorm:
self.norm2 = operations.LayerNorm(
@@ -466,7 +483,11 @@ class DismantledBlock(nn.Module):
multiple_of=256,
)
self.scale_mod_only = scale_mod_only
- if not scale_mod_only:
+ if x_block_self_attn:
+ assert not pre_only
+ assert not scale_mod_only
+ n_mods = 9
+ elif not scale_mod_only:
n_mods = 6 if not pre_only else 2
else:
n_mods = 4 if not pre_only else 1
@@ -527,14 +548,64 @@ class DismantledBlock(nn.Module):
)
return x
+ def pre_attention_x(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+ assert self.x_block_self_attn
+ (
+ shift_msa,
+ scale_msa,
+ gate_msa,
+ shift_mlp,
+ scale_mlp,
+ gate_mlp,
+ shift_msa2,
+ scale_msa2,
+ gate_msa2,
+ ) = self.adaLN_modulation(c).chunk(9, dim=1)
+ x_norm = self.norm1(x)
+ qkv = self.attn.pre_attention(modulate(x_norm, shift_msa, scale_msa))
+ qkv2 = self.attn2.pre_attention(modulate(x_norm, shift_msa2, scale_msa2))
+ return qkv, qkv2, (
+ x,
+ gate_msa,
+ shift_mlp,
+ scale_mlp,
+ gate_mlp,
+ gate_msa2,
+ )
+
+ def post_attention_x(self, attn, attn2, x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2):
+ assert not self.pre_only
+ attn1 = self.attn.post_attention(attn)
+ attn2 = self.attn2.post_attention(attn2)
+ out1 = gate_msa.unsqueeze(1) * attn1
+ out2 = gate_msa2.unsqueeze(1) * attn2
+ x = x + out1
+ x = x + out2
+ x = x + gate_mlp.unsqueeze(1) * self.mlp(
+ modulate(self.norm2(x), shift_mlp, scale_mlp)
+ )
+ return x
+
def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
assert not self.pre_only
- qkv, intermediates = self.pre_attention(x, c)
- attn = optimized_attention(
- qkv,
- num_heads=self.attn.num_heads,
- )
- return self.post_attention(attn, *intermediates)
+ if self.x_block_self_attn:
+ qkv, qkv2, intermediates = self.pre_attention_x(x, c)
+ attn, _ = optimized_attention(
+ qkv[0], qkv[1], qkv[2],
+ num_heads=self.attn.num_heads,
+ )
+ attn2, _ = optimized_attention(
+ qkv2[0], qkv2[1], qkv2[2],
+ num_heads=self.attn2.num_heads,
+ )
+ return self.post_attention_x(attn, attn2, *intermediates)
+ else:
+ qkv, intermediates = self.pre_attention(x, c)
+ attn = optimized_attention(
+ qkv[0], qkv[1], qkv[2],
+ heads=self.attn.num_heads,
+ )
+ return self.post_attention(attn, *intermediates)
def block_mixing(*args, use_checkpoint=True, **kwargs):
@@ -549,7 +620,10 @@ def block_mixing(*args, use_checkpoint=True, **kwargs):
def _block_mixing(context, x, context_block, x_block, c):
context_qkv, context_intermediates = context_block.pre_attention(context, c)
- x_qkv, x_intermediates = x_block.pre_attention(x, c)
+ if x_block.x_block_self_attn:
+ x_qkv, x_qkv2, x_intermediates = x_block.pre_attention_x(x, c)
+ else:
+ x_qkv, x_intermediates = x_block.pre_attention(x, c)
o = []
for t in range(3):
@@ -557,8 +631,8 @@ def _block_mixing(context, x, context_block, x_block, c):
qkv = tuple(o)
attn = optimized_attention(
- qkv,
- num_heads=x_block.attn.num_heads,
+ qkv[0], qkv[1], qkv[2],
+ heads=x_block.attn.num_heads,
)
context_attn, x_attn = (
attn[:, : context_qkv[0].shape[1]],
@@ -570,7 +644,14 @@ def _block_mixing(context, x, context_block, x_block, c):
else:
context = None
- x = x_block.post_attention(x_attn, *x_intermediates)
+ if x_block.x_block_self_attn:
+ attn2 = optimized_attention(
+ x_qkv2[0], x_qkv2[1], x_qkv2[2],
+ heads=x_block.attn2.num_heads,
+ )
+ x = x_block.post_attention_x(x_attn, attn2, *x_intermediates)
+ else:
+ x = x_block.post_attention(x_attn, *x_intermediates)
return context, x
@@ -585,8 +666,13 @@ class JointBlock(nn.Module):
super().__init__()
pre_only = kwargs.pop("pre_only")
qk_norm = kwargs.pop("qk_norm", None)
+ x_block_self_attn = kwargs.pop("x_block_self_attn", False)
self.context_block = DismantledBlock(*args, pre_only=pre_only, qk_norm=qk_norm, **kwargs)
- self.x_block = DismantledBlock(*args, pre_only=False, qk_norm=qk_norm, **kwargs)
+ self.x_block = DismantledBlock(*args,
+ pre_only=False,
+ qk_norm=qk_norm,
+ x_block_self_attn=x_block_self_attn,
+ **kwargs)
def forward(self, *args, **kwargs):
return block_mixing(
@@ -642,7 +728,7 @@ class SelfAttentionContext(nn.Module):
def forward(self, x):
qkv = self.qkv(x)
q, k, v = split_qkv(qkv, self.dim_head)
- x = optimized_attention((q.reshape(q.shape[0], q.shape[1], -1), k, v), self.heads)
+ x = optimized_attention(q.reshape(q.shape[0], q.shape[1], -1), k, v, heads=self.heads)
return self.proj(x)
class ContextProcessorBlock(nn.Module):
@@ -701,9 +787,12 @@ class MMDiT(nn.Module):
qk_norm: Optional[str] = None,
qkv_bias: bool = True,
context_processor_layers = None,
+ x_block_self_attn: bool = False,
+ x_block_self_attn_layers: Optional[List[int]] = [],
context_size = 4096,
num_blocks = None,
final_layer = True,
+ skip_blocks = False,
dtype = None, #TODO
device = None,
operations = None,
@@ -718,6 +807,7 @@ class MMDiT(nn.Module):
self.pos_embed_scaling_factor = pos_embed_scaling_factor
self.pos_embed_offset = pos_embed_offset
self.pos_embed_max_size = pos_embed_max_size
+ self.x_block_self_attn_layers = x_block_self_attn_layers
# hidden_size = default(hidden_size, 64 * depth)
# num_heads = default(num_heads, hidden_size // 64)
@@ -775,26 +865,28 @@ class MMDiT(nn.Module):
self.pos_embed = None
self.use_checkpoint = use_checkpoint
- self.joint_blocks = nn.ModuleList(
- [
- JointBlock(
- self.hidden_size,
- num_heads,
- mlp_ratio=mlp_ratio,
- qkv_bias=qkv_bias,
- attn_mode=attn_mode,
- pre_only=(i == num_blocks - 1) and final_layer,
- rmsnorm=rmsnorm,
- scale_mod_only=scale_mod_only,
- swiglu=swiglu,
- qk_norm=qk_norm,
- dtype=dtype,
- device=device,
- operations=operations
- )
- for i in range(num_blocks)
- ]
- )
+ if not skip_blocks:
+ self.joint_blocks = nn.ModuleList(
+ [
+ JointBlock(
+ self.hidden_size,
+ num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ attn_mode=attn_mode,
+ pre_only=(i == num_blocks - 1) and final_layer,
+ rmsnorm=rmsnorm,
+ scale_mod_only=scale_mod_only,
+ swiglu=swiglu,
+ qk_norm=qk_norm,
+ x_block_self_attn=(i in self.x_block_self_attn_layers) or x_block_self_attn,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+ for i in range(num_blocks)
+ ]
+ )
if final_layer:
self.final_layer = FinalLayer(self.hidden_size, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations)
@@ -857,7 +949,9 @@ class MMDiT(nn.Module):
c_mod: torch.Tensor,
context: Optional[torch.Tensor] = None,
control = None,
+ transformer_options = {},
) -> torch.Tensor:
+ patches_replace = transformer_options.get("patches_replace", {})
if self.register_length > 0:
context = torch.cat(
(
@@ -869,14 +963,25 @@ class MMDiT(nn.Module):
# context is B, L', D
# x is B, L, D
+ blocks_replace = patches_replace.get("dit", {})
blocks = len(self.joint_blocks)
for i in range(blocks):
- context, x = self.joint_blocks[i](
- context,
- x,
- c=c_mod,
- use_checkpoint=self.use_checkpoint,
- )
+ if ("double_block", i) in blocks_replace:
+ def block_wrap(args):
+ out = {}
+ out["txt"], out["img"] = self.joint_blocks[i](args["txt"], args["img"], c=args["vec"])
+ return out
+
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": c_mod}, {"original_block": block_wrap})
+ context = out["txt"]
+ x = out["img"]
+ else:
+ context, x = self.joint_blocks[i](
+ context,
+ x,
+ c=c_mod,
+ use_checkpoint=self.use_checkpoint,
+ )
if control is not None:
control_o = control.get("output")
if i < len(control_o):
@@ -894,6 +999,7 @@ class MMDiT(nn.Module):
y: Optional[torch.Tensor] = None,
context: Optional[torch.Tensor] = None,
control = None,
+ transformer_options = {},
) -> torch.Tensor:
"""
Forward pass of DiT.
@@ -915,7 +1021,7 @@ class MMDiT(nn.Module):
if context is not None:
context = self.context_embedder(context)
- x = self.forward_core_with_concat(x, c, context, control)
+ x = self.forward_core_with_concat(x, c, context, control, transformer_options)
x = self.unpatchify(x, hw=hw) # (N, out_channels, H, W)
return x[:,:,:hw[-2],:hw[-1]]
@@ -929,7 +1035,8 @@ class OpenAISignatureMMDITWrapper(MMDiT):
context: Optional[torch.Tensor] = None,
y: Optional[torch.Tensor] = None,
control = None,
+ transformer_options = {},
**kwargs,
) -> torch.Tensor:
- return super().forward(x, timesteps, context=context, y=y, control=control)
+ return super().forward(x, timesteps, context=context, y=y, control=control, transformer_options=transformer_options)
diff --git a/comfy/lora.py b/comfy/lora.py
index 80057cdd..b745ca4d 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -317,6 +317,10 @@ def model_lora_keys_unet(model, key_map={}):
key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora
key_map[key_lora] = to
+ key_lora = "lycoris_{}".format(k[:-len(".weight")].replace(".", "_")) #simpletuner lycoris format
+ key_map[key_lora] = to
+
+
if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow
diffusers_keys = comfy.utils.auraflow_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")
for k in diffusers_keys:
@@ -415,7 +419,7 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32):
weight *= strength_model
if isinstance(v, list):
- v = (calculate_weight(v[1:], comfy.model_management.cast_to_device(v[0], weight.device, intermediate_dtype, copy=True), key, intermediate_dtype=intermediate_dtype), )
+ v = (calculate_weight(v[1:], v[0][1](comfy.model_management.cast_to_device(v[0][0], weight.device, intermediate_dtype, copy=True), inplace=True), key, intermediate_dtype=intermediate_dtype), )
if len(v) == 1:
patch_type = "diff"
diff --git a/comfy/model_base.py b/comfy/model_base.py
index a98fee1d..f2833168 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -24,6 +24,7 @@ from comfy.ldm.cascade.stage_b import StageB
from comfy.ldm.modules.encoders.noise_aug_modules import CLIPEmbeddingNoiseAugmentation
from comfy.ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
from comfy.ldm.modules.diffusionmodules.mmdit import OpenAISignatureMMDITWrapper
+import comfy.ldm.genmo.joint_model.asymm_models_joint
import comfy.ldm.aura.mmdit
import comfy.ldm.hydit.models
import comfy.ldm.audio.dit
@@ -96,7 +97,8 @@ class BaseModel(torch.nn.Module):
if not unet_config.get("disable_unet_model_creation", False):
if model_config.custom_operations is None:
- operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype, fp8_optimizations=model_config.optimizations.get("fp8", False))
+ fp8 = model_config.optimizations.get("fp8", model_config.scaled_fp8 is not None)
+ operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype, fp8_optimizations=fp8, scaled_fp8=model_config.scaled_fp8)
else:
operations = model_config.custom_operations
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
@@ -244,6 +246,10 @@ class BaseModel(torch.nn.Module):
extra_sds.append(self.model_config.process_clip_vision_state_dict_for_saving(clip_vision_state_dict))
unet_state_dict = self.diffusion_model.state_dict()
+
+ if self.model_config.scaled_fp8 is not None:
+ unet_state_dict["scaled_fp8"] = torch.tensor([], dtype=self.model_config.scaled_fp8)
+
unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
if self.model_type == ModelType.V_PREDICTION:
@@ -713,3 +719,18 @@ class Flux(BaseModel):
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([kwargs.get("guidance", 3.5)]))
return out
+
+class GenmoMochi(BaseModel):
+ def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+ super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.genmo.joint_model.asymm_models_joint.AsymmDiTJoint)
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ attention_mask = kwargs.get("attention_mask", None)
+ if attention_mask is not None:
+ out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
+ out['num_tokens'] = comfy.conds.CONDConstant(max(1, torch.sum(attention_mask).item()))
+ cross_attn = kwargs.get("cross_attn", None)
+ if cross_attn is not None:
+ out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+ return out
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index 1edbcda4..8435de3e 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -70,6 +70,11 @@ def detect_unet_config(state_dict, key_prefix):
context_processor = '{}context_processor.layers.0.attn.qkv.weight'.format(key_prefix)
if context_processor in state_dict_keys:
unet_config["context_processor_layers"] = count_blocks(state_dict_keys, '{}context_processor.layers.'.format(key_prefix) + '{}.')
+ unet_config["x_block_self_attn_layers"] = []
+ for key in state_dict_keys:
+ if key.startswith('{}joint_blocks.'.format(key_prefix)) and key.endswith('.x_block.attn2.qkv.weight'):
+ layer = key[len('{}joint_blocks.'.format(key_prefix)):-len('.x_block.attn2.qkv.weight')]
+ unet_config["x_block_self_attn_layers"].append(int(layer))
return unet_config
if '{}clf.1.weight'.format(key_prefix) in state_dict_keys: #stable cascade
@@ -145,6 +150,34 @@ def detect_unet_config(state_dict, key_prefix):
dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
return dit_config
+ if '{}t5_yproj.weight'.format(key_prefix) in state_dict_keys: #Genmo mochi preview
+ dit_config = {}
+ dit_config["image_model"] = "mochi_preview"
+ dit_config["depth"] = 48
+ dit_config["patch_size"] = 2
+ dit_config["num_heads"] = 24
+ dit_config["hidden_size_x"] = 3072
+ dit_config["hidden_size_y"] = 1536
+ dit_config["mlp_ratio_x"] = 4.0
+ dit_config["mlp_ratio_y"] = 4.0
+ dit_config["learn_sigma"] = False
+ dit_config["in_channels"] = 12
+ dit_config["qk_norm"] = True
+ dit_config["qkv_bias"] = False
+ dit_config["out_bias"] = True
+ dit_config["attn_drop"] = 0.0
+ dit_config["patch_embed_bias"] = True
+ dit_config["posenc_preserve_area"] = True
+ dit_config["timestep_mlp_bias"] = True
+ dit_config["attend_to_padding"] = False
+ dit_config["timestep_scale"] = 1000.0
+ dit_config["use_t5"] = True
+ dit_config["t5_feat_dim"] = 4096
+ dit_config["t5_token_length"] = 256
+ dit_config["rope_theta"] = 10000.0
+ return dit_config
+
+
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
return None
@@ -286,9 +319,15 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
return None
model_config = model_config_from_unet_config(unet_config, state_dict)
if model_config is None and use_base_if_no_match:
- return comfy.supported_models_base.BASE(unet_config)
- else:
- return model_config
+ model_config = comfy.supported_models_base.BASE(unet_config)
+
+ scaled_fp8_weight = state_dict.get("{}scaled_fp8".format(unet_key_prefix), None)
+ if scaled_fp8_weight is not None:
+ model_config.scaled_fp8 = scaled_fp8_weight.dtype
+ if model_config.scaled_fp8 == torch.float32:
+ model_config.scaled_fp8 = torch.float8_e4m3fn
+
+ return model_config
def unet_prefix_from_state_dict(state_dict):
candidates = ["model.diffusion_model.", #ldm/sgm models
diff --git a/comfy/model_management.py b/comfy/model_management.py
index 2346d4ac..855e8911 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -647,6 +647,9 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor
pass
if fp8_dtype is not None:
+ if supports_fp8_compute(device): #if fp8 compute is supported the casting is most likely not expensive
+ return fp8_dtype
+
free_model_memory = maximum_vram_for_weights(device)
if model_params * 2 > free_model_memory:
return fp8_dtype
@@ -840,27 +843,21 @@ def force_channels_last():
#TODO
return False
+def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False):
+ if device is None or weight.device == device:
+ if not copy:
+ if dtype is None or weight.dtype == dtype:
+ return weight
+ return weight.to(dtype=dtype, copy=copy)
+
+ r = torch.empty_like(weight, dtype=dtype, device=device)
+ r.copy_(weight, non_blocking=non_blocking)
+ return r
+
def cast_to_device(tensor, device, dtype, copy=False):
- device_supports_cast = False
- if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
- device_supports_cast = True
- elif tensor.dtype == torch.bfloat16:
- if hasattr(device, 'type') and device.type.startswith("cuda"):
- device_supports_cast = True
- elif is_intel_xpu():
- device_supports_cast = True
+ non_blocking = device_supports_non_blocking(device)
+ return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy)
- non_blocking = device_should_use_non_blocking(device)
-
- if device_supports_cast:
- if copy:
- if tensor.device == device:
- return tensor.to(dtype, copy=copy, non_blocking=non_blocking)
- return tensor.to(device, copy=copy, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
- else:
- return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
- else:
- return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
def xformers_enabled():
global directml_enabled
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index 2ba30633..3bba217a 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -94,6 +94,31 @@ class LowVramPatch:
return comfy.float.stochastic_rounding(comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype), weight.dtype, seed=string_to_seed(self.key))
return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)
+
+def get_key_weight(model, key):
+ set_func = None
+ convert_func = None
+ op_keys = key.rsplit('.', 1)
+ if len(op_keys) < 2:
+ weight = comfy.utils.get_attr(model, key)
+ else:
+ op = comfy.utils.get_attr(model, op_keys[0])
+ try:
+ set_func = getattr(op, "set_{}".format(op_keys[1]))
+ except AttributeError:
+ pass
+
+ try:
+ convert_func = getattr(op, "convert_{}".format(op_keys[1]))
+ except AttributeError:
+ pass
+
+ weight = getattr(op, op_keys[1])
+ if convert_func is not None:
+ weight = comfy.utils.get_attr(model, key)
+
+ return weight, set_func, convert_func
+
class ModelPatcher:
def __init__(self, model, load_device, offload_device, size=0, weight_inplace_update=False):
self.size = size
@@ -294,14 +319,16 @@ class ModelPatcher:
if not k.startswith(filter_prefix):
continue
bk = self.backup.get(k, None)
+ weight, set_func, convert_func = get_key_weight(self.model, k)
if bk is not None:
weight = bk.weight
- else:
- weight = model_sd[k]
+ if convert_func is None:
+ convert_func = lambda a, **kwargs: a
+
if k in self.patches:
- p[k] = [weight] + self.patches[k]
+ p[k] = [(weight, convert_func)] + self.patches[k]
else:
- p[k] = (weight,)
+ p[k] = [(weight, convert_func)]
return p
def model_state_dict(self, filter_prefix=None):
@@ -317,8 +344,7 @@ class ModelPatcher:
if key not in self.patches:
return
- weight = comfy.utils.get_attr(self.model, key)
-
+ weight, set_func, convert_func = get_key_weight(self.model, key)
inplace_update = self.weight_inplace_update or inplace_update
if key not in self.backup:
@@ -328,12 +354,18 @@ class ModelPatcher:
temp_weight = comfy.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
else:
temp_weight = weight.to(torch.float32, copy=True)
+ if convert_func is not None:
+ temp_weight = convert_func(temp_weight, inplace=True)
+
out_weight = comfy.lora.calculate_weight(self.patches[key], temp_weight, key)
- out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))
- if inplace_update:
- comfy.utils.copy_to_param(self.model, key, out_weight)
+ if set_func is None:
+ out_weight = comfy.float.stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key))
+ if inplace_update:
+ comfy.utils.copy_to_param(self.model, key, out_weight)
+ else:
+ comfy.utils.set_attr_param(self.model, key, out_weight)
else:
- comfy.utils.set_attr_param(self.model, key, out_weight)
+ set_func(out_weight, inplace_update=inplace_update, seed=string_to_seed(key))
def load(self, device_to=None, lowvram_model_memory=0, force_patch_weights=False, full_load=False):
mem_counter = 0
diff --git a/comfy/ops.py b/comfy/ops.py
index c90e25ea..3c5ba012 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -19,20 +19,12 @@
import torch
import comfy.model_management
from comfy.cli_args import args
+import comfy.float
-def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False):
- if device is None or weight.device == device:
- if not copy:
- if dtype is None or weight.dtype == dtype:
- return weight
- return weight.to(dtype=dtype, copy=copy)
-
- r = torch.empty_like(weight, dtype=dtype, device=device)
- r.copy_(weight, non_blocking=non_blocking)
- return r
+cast_to = comfy.model_management.cast_to #TODO: remove once no more references
def cast_to_input(weight, input, non_blocking=False, copy=True):
- return cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
+ return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
if input is not None:
@@ -47,12 +39,12 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
non_blocking = comfy.model_management.device_supports_non_blocking(device)
if s.bias is not None:
has_function = s.bias_function is not None
- bias = cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function)
+ bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function)
if has_function:
bias = s.bias_function(bias)
has_function = s.weight_function is not None
- weight = cast_to(s.weight, dtype, device, non_blocking=non_blocking, copy=has_function)
+ weight = comfy.model_management.cast_to(s.weight, dtype, device, non_blocking=non_blocking, copy=has_function)
if has_function:
weight = s.weight_function(weight)
return weight, bias
@@ -258,19 +250,29 @@ def fp8_linear(self, input):
if dtype not in [torch.float8_e4m3fn]:
return None
+ tensor_2d = False
+ if len(input.shape) == 2:
+ tensor_2d = True
+ input = input.unsqueeze(1)
+
+
if len(input.shape) == 3:
- inn = input.reshape(-1, input.shape[2]).to(dtype)
w, bias = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input.dtype)
w = w.t()
scale_weight = self.scale_weight
scale_input = self.scale_input
if scale_weight is None:
- scale_weight = torch.ones((1), device=input.device, dtype=torch.float32)
- if scale_input is None:
- scale_input = scale_weight
+ scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
+ else:
+ scale_weight = scale_weight.to(input.device)
+
if scale_input is None:
- scale_input = torch.ones((1), device=input.device, dtype=torch.float32)
+ scale_input = torch.ones((), device=input.device, dtype=torch.float32)
+ inn = input.reshape(-1, input.shape[2]).to(dtype)
+ else:
+ scale_input = scale_input.to(input.device)
+ inn = (input * (1.0 / scale_input).to(input.dtype)).reshape(-1, input.shape[2]).to(dtype)
if bias is not None:
o = torch._scaled_mm(inn, w, out_dtype=input.dtype, bias=bias, scale_a=scale_input, scale_b=scale_weight)
@@ -280,7 +282,11 @@ def fp8_linear(self, input):
if isinstance(o, tuple):
o = o[0]
+ if tensor_2d:
+ return o.reshape(input.shape[0], -1)
+
return o.reshape((-1, input.shape[1], self.weight.shape[0]))
+
return None
class fp8_ops(manual_cast):
@@ -298,15 +304,63 @@ class fp8_ops(manual_cast):
weight, bias = cast_bias_weight(self, input)
return torch.nn.functional.linear(input, weight, bias)
+def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None):
+ class scaled_fp8_op(manual_cast):
+ class Linear(manual_cast.Linear):
+ def __init__(self, *args, **kwargs):
+ if override_dtype is not None:
+ kwargs['dtype'] = override_dtype
+ super().__init__(*args, **kwargs)
-def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False):
- if comfy.model_management.supports_fp8_compute(load_device):
- if (fp8_optimizations or args.fast) and not disable_fast_fp8:
- return fp8_ops
+ def reset_parameters(self):
+ if not hasattr(self, 'scale_weight'):
+ self.scale_weight = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False)
+
+ if not scale_input:
+ self.scale_input = None
+
+ if not hasattr(self, 'scale_input'):
+ self.scale_input = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False)
+ return None
+
+ def forward_comfy_cast_weights(self, input):
+ if fp8_matrix_mult:
+ out = fp8_linear(self, input)
+ if out is not None:
+ return out
+
+ weight, bias = cast_bias_weight(self, input)
+
+ if weight.numel() < input.numel(): #TODO: optimize
+ return torch.nn.functional.linear(input, weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype), bias)
+ else:
+ return torch.nn.functional.linear(input * self.scale_weight.to(device=weight.device, dtype=weight.dtype), weight, bias)
+
+ def convert_weight(self, weight, inplace=False, **kwargs):
+ if inplace:
+ weight *= self.scale_weight.to(device=weight.device, dtype=weight.dtype)
+ return weight
+ else:
+ return weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype)
+
+ def set_weight(self, weight, inplace_update=False, seed=None, **kwargs):
+ weight = comfy.float.stochastic_rounding(weight / self.scale_weight.to(device=weight.device, dtype=weight.dtype), self.weight.dtype, seed=seed)
+ if inplace_update:
+ self.weight.data.copy_(weight)
+ else:
+ self.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+ return scaled_fp8_op
+
+def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, scaled_fp8=None):
+ fp8_compute = comfy.model_management.supports_fp8_compute(load_device)
+ if scaled_fp8 is not None:
+ return scaled_fp8_ops(fp8_matrix_mult=fp8_compute, scale_input=True, override_dtype=scaled_fp8)
+
+ if fp8_compute and (fp8_optimizations or args.fast) and not disable_fast_fp8:
+ return fp8_ops
if compute_dtype is None or weight_dtype == compute_dtype:
return disable_weight_init
- if args.fast and not disable_fast_fp8:
- if comfy.model_management.supports_fp8_compute(load_device):
- return fp8_ops
+
return manual_cast
diff --git a/comfy/samplers.py b/comfy/samplers.py
index 4a975af7..eb520862 100644
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -358,11 +358,35 @@ def beta_scheduler(model_sampling, steps, alpha=0.6, beta=0.6):
ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
sigs = []
+ last_t = -1
for t in ts:
- sigs += [float(model_sampling.sigmas[int(t)])]
+ if t != last_t:
+ sigs += [float(model_sampling.sigmas[int(t)])]
+ last_t = t
sigs += [0.0]
return torch.FloatTensor(sigs)
+# from: https://github.com/genmoai/models/blob/main/src/mochi_preview/infer.py#L41
+def linear_quadratic_schedule(model_sampling, steps, threshold_noise=0.025, linear_steps=None):
+ if steps == 1:
+ sigma_schedule = [1.0, 0.0]
+ else:
+ if linear_steps is None:
+ linear_steps = steps // 2
+ linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)]
+ threshold_noise_step_diff = linear_steps - threshold_noise * steps
+ quadratic_steps = steps - linear_steps
+ quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps ** 2)
+ linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps ** 2)
+ const = quadratic_coef * (linear_steps ** 2)
+ quadratic_sigma_schedule = [
+ quadratic_coef * (i ** 2) + linear_coef * i + const
+ for i in range(linear_steps, steps)
+ ]
+ sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
+ sigma_schedule = [1.0 - x for x in sigma_schedule]
+ return torch.FloatTensor(sigma_schedule) * model_sampling.sigma_max.cpu()
+
def get_mask_aabb(masks):
if masks.numel() == 0:
return torch.zeros((0, 4), device=masks.device, dtype=torch.int)
@@ -729,7 +753,7 @@ def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model
return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
-SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta"]
+SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta", "linear_quadratic"]
SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
def calculate_sigmas(model_sampling, scheduler_name, steps):
@@ -747,6 +771,8 @@ def calculate_sigmas(model_sampling, scheduler_name, steps):
sigmas = normal_scheduler(model_sampling, steps, sgm=True)
elif scheduler_name == "beta":
sigmas = beta_scheduler(model_sampling, steps)
+ elif scheduler_name == "linear_quadratic":
+ sigmas = linear_quadratic_schedule(model_sampling, steps)
else:
logging.error("error invalid scheduler {}".format(scheduler_name))
return sigmas
diff --git a/comfy/sd.py b/comfy/sd.py
index 67b4ff0c..e9a3579c 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -7,6 +7,7 @@ from .ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine
from .ldm.cascade.stage_a import StageA
from .ldm.cascade.stage_c_coder import StageC_coder
from .ldm.audio.autoencoder import AudioOobleckVAE
+import comfy.ldm.genmo.vae.model
import yaml
import comfy.utils
@@ -25,6 +26,7 @@ import comfy.text_encoders.aura_t5
import comfy.text_encoders.hydit
import comfy.text_encoders.flux
import comfy.text_encoders.long_clipl
+import comfy.text_encoders.genmo
import comfy.model_patcher
import comfy.lora
@@ -241,6 +243,13 @@ class VAE:
self.process_output = lambda audio: audio
self.process_input = lambda audio: audio
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+ elif "blocks.2.blocks.3.stack.5.weight" in sd or "decoder.blocks.2.blocks.3.stack.5.weight" in sd: #genmo mochi vae
+ if "blocks.2.blocks.3.stack.5.weight" in sd:
+ sd = comfy.utils.state_dict_prefix_replace(sd, {"": "decoder."})
+ self.first_stage_model = comfy.ldm.genmo.vae.model.VideoVAE()
+ self.latent_channels = 12
+ self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * shape[4] * (6 * 8 * 8)) * model_management.dtype_size(dtype)
+ self.upscale_ratio = (lambda a: max(0, a * 6 - 5), 8, 8)
else:
logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
self.first_stage_model = None
@@ -296,6 +305,10 @@ class VAE:
decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
return comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_x,), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device)
+ def decode_tiled_3d(self, samples, tile_t=999, tile_x=32, tile_y=32, overlap=(1, 8, 8)):
+ decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
+ return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device))
+
def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
@@ -314,6 +327,7 @@ class VAE:
return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_x,), overlap=overlap, upscale_amount=(1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device)
def decode(self, samples_in):
+ pixel_samples = None
try:
memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
model_management.load_models_gpu([self.patcher], memory_required=memory_used)
@@ -321,16 +335,21 @@ class VAE:
batch_number = int(free_memory / memory_used)
batch_number = max(1, batch_number)
- pixel_samples = torch.empty((samples_in.shape[0], self.output_channels) + tuple(map(lambda a: a * self.upscale_ratio, samples_in.shape[2:])), device=self.output_device)
for x in range(0, samples_in.shape[0], batch_number):
samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
- pixel_samples[x:x+batch_number] = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float())
+ out = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float())
+ if pixel_samples is None:
+ pixel_samples = torch.empty((samples_in.shape[0],) + tuple(out.shape[1:]), device=self.output_device)
+ pixel_samples[x:x+batch_number] = out
except model_management.OOM_EXCEPTION as e:
logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
- if len(samples_in.shape) == 3:
+ dims = samples_in.ndim - 2
+ if dims == 1:
pixel_samples = self.decode_tiled_1d(samples_in)
- else:
+ elif dims == 2:
pixel_samples = self.decode_tiled_(samples_in)
+ elif dims == 3:
+ pixel_samples = self.decode_tiled_3d(samples_in)
pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
return pixel_samples
@@ -398,6 +417,7 @@ class CLIPType(Enum):
STABLE_AUDIO = 4
HUNYUAN_DIT = 5
FLUX = 6
+ MOCHI = 7
def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
clip_data = []
@@ -432,16 +452,14 @@ def detect_te_model(sd):
return None
-def t5xxl_weight_dtype(clip_data):
+def t5xxl_detect(clip_data):
weight_name = "encoder.block.23.layer.1.DenseReluDense.wi_1.weight"
- dtype_t5 = None
for sd in clip_data:
- weight = sd.get(weight_name, None)
- if weight is not None:
- dtype_t5 = weight.dtype
- break
- return dtype_t5
+ if weight_name in sd:
+ return comfy.text_encoders.sd3_clip.t5_xxl_detect(sd)
+
+ return {}
def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
@@ -475,8 +493,12 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
clip_target.clip = comfy.text_encoders.sd2_clip.SD2ClipModel
clip_target.tokenizer = comfy.text_encoders.sd2_clip.SD2Tokenizer
elif te_model == TEModel.T5_XXL:
- clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=False, clip_g=False, t5=True, dtype_t5=t5xxl_weight_dtype(clip_data))
- clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
+ if clip_type == CLIPType.SD3:
+ clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=False, clip_g=False, t5=True, **t5xxl_detect(clip_data))
+ clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
+ else: #CLIPType.MOCHI
+ clip_target.clip = comfy.text_encoders.genmo.mochi_te(**t5xxl_detect(clip_data))
+ clip_target.tokenizer = comfy.text_encoders.genmo.MochiT5Tokenizer
elif te_model == TEModel.T5_XL:
clip_target.clip = comfy.text_encoders.aura_t5.AuraT5Model
clip_target.tokenizer = comfy.text_encoders.aura_t5.AuraT5Tokenizer
@@ -493,19 +515,19 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
elif len(clip_data) == 2:
if clip_type == CLIPType.SD3:
te_models = [detect_te_model(clip_data[0]), detect_te_model(clip_data[1])]
- clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=TEModel.CLIP_L in te_models, clip_g=TEModel.CLIP_G in te_models, t5=TEModel.T5_XXL in te_models, dtype_t5=t5xxl_weight_dtype(clip_data))
+ clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=TEModel.CLIP_L in te_models, clip_g=TEModel.CLIP_G in te_models, t5=TEModel.T5_XXL in te_models, **t5xxl_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
elif clip_type == CLIPType.HUNYUAN_DIT:
clip_target.clip = comfy.text_encoders.hydit.HyditModel
clip_target.tokenizer = comfy.text_encoders.hydit.HyditTokenizer
elif clip_type == CLIPType.FLUX:
- clip_target.clip = comfy.text_encoders.flux.flux_clip(dtype_t5=t5xxl_weight_dtype(clip_data))
+ clip_target.clip = comfy.text_encoders.flux.flux_clip(**t5xxl_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.flux.FluxTokenizer
else:
clip_target.clip = sdxl_clip.SDXLClipModel
clip_target.tokenizer = sdxl_clip.SDXLTokenizer
elif len(clip_data) == 3:
- clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(dtype_t5=t5xxl_weight_dtype(clip_data))
+ clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(**t5xxl_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
parameters = 0
@@ -580,7 +602,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
return None
unet_weight_dtype = list(model_config.supported_inference_dtypes)
- if weight_dtype is not None:
+ if weight_dtype is not None and model_config.scaled_fp8 is None:
unet_weight_dtype.append(weight_dtype)
model_config.custom_operations = model_options.get("custom_operations", None)
@@ -649,6 +671,8 @@ def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffuse
sd = temp_sd
parameters = comfy.utils.calculate_parameters(sd)
+ weight_dtype = comfy.utils.weight_dtype(sd)
+
load_device = model_management.get_torch_device()
model_config = model_detection.model_config_from_unet(sd, "")
@@ -675,8 +699,12 @@ def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffuse
logging.warning("{} {}".format(diffusers_keys[k], k))
offload_device = model_management.unet_offload_device()
+ unet_weight_dtype = list(model_config.supported_inference_dtypes)
+ if weight_dtype is not None and model_config.scaled_fp8 is None:
+ unet_weight_dtype.append(weight_dtype)
+
if dtype is None:
- unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
+ unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=unet_weight_dtype)
else:
unet_dtype = dtype
diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
index bb240526..a454f3bb 100644
--- a/comfy/sd1_clip.py
+++ b/comfy/sd1_clip.py
@@ -94,11 +94,20 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
config = json.load(f)
operations = model_options.get("custom_operations", None)
+ scaled_fp8 = None
+
if operations is None:
- operations = comfy.ops.manual_cast
+ scaled_fp8 = model_options.get("scaled_fp8", None)
+ if scaled_fp8 is not None:
+ operations = comfy.ops.scaled_fp8_ops(fp8_matrix_mult=False, override_dtype=scaled_fp8)
+ else:
+ operations = comfy.ops.manual_cast
self.operations = operations
self.transformer = model_class(config, dtype, device, self.operations)
+ if scaled_fp8 is not None:
+ self.transformer.scaled_fp8 = torch.nn.Parameter(torch.tensor([], dtype=scaled_fp8))
+
self.num_layers = self.transformer.num_layers
self.max_length = max_length
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index 3603313f..9931f4c5 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -10,6 +10,7 @@ import comfy.text_encoders.sa_t5
import comfy.text_encoders.aura_t5
import comfy.text_encoders.hydit
import comfy.text_encoders.flux
+import comfy.text_encoders.genmo
from . import supported_models_base
from . import latent_formats
@@ -529,12 +530,11 @@ class SD3(supported_models_base.BASE):
clip_l = True
if "{}clip_g.transformer.text_model.final_layer_norm.weight".format(pref) in state_dict:
clip_g = True
- t5_key = "{}t5xxl.transformer.encoder.final_layer_norm.weight".format(pref)
- if t5_key in state_dict:
+ t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
+ if "dtype_t5" in t5_detect:
t5 = True
- dtype_t5 = state_dict[t5_key].dtype
- return supported_models_base.ClipTarget(comfy.text_encoders.sd3_clip.SD3Tokenizer, comfy.text_encoders.sd3_clip.sd3_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, dtype_t5=dtype_t5))
+ return supported_models_base.ClipTarget(comfy.text_encoders.sd3_clip.SD3Tokenizer, comfy.text_encoders.sd3_clip.sd3_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, **t5_detect))
class StableAudio(supported_models_base.BASE):
unet_config = {
@@ -653,11 +653,8 @@ class Flux(supported_models_base.BASE):
def clip_target(self, state_dict={}):
pref = self.text_encoder_key_prefix[0]
- t5_key = "{}t5xxl.transformer.encoder.final_layer_norm.weight".format(pref)
- dtype_t5 = None
- if t5_key in state_dict:
- dtype_t5 = state_dict[t5_key].dtype
- return supported_models_base.ClipTarget(comfy.text_encoders.flux.FluxTokenizer, comfy.text_encoders.flux.flux_clip(dtype_t5=dtype_t5))
+ t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
+ return supported_models_base.ClipTarget(comfy.text_encoders.flux.FluxTokenizer, comfy.text_encoders.flux.flux_clip(**t5_detect))
class FluxSchnell(Flux):
unet_config = {
@@ -674,7 +671,36 @@ class FluxSchnell(Flux):
out = model_base.Flux(self, model_type=model_base.ModelType.FLOW, device=device)
return out
+class GenmoMochi(supported_models_base.BASE):
+ unet_config = {
+ "image_model": "mochi_preview",
+ }
-models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, HunyuanDiT, HunyuanDiT1, Flux, FluxSchnell]
+ sampling_settings = {
+ "multiplier": 1.0,
+ "shift": 6.0,
+ }
+
+ unet_extra_config = {}
+ latent_format = latent_formats.Mochi
+
+ memory_usage_factor = 2.0 #TODO
+
+ supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+ vae_key_prefix = ["vae."]
+ text_encoder_key_prefix = ["text_encoders."]
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.GenmoMochi(self, device=device)
+ return out
+
+ def clip_target(self, state_dict={}):
+ pref = self.text_encoder_key_prefix[0]
+ t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
+ return supported_models_base.ClipTarget(comfy.text_encoders.genmo.MochiT5Tokenizer, comfy.text_encoders.genmo.mochi_te(**t5_detect))
+
+
+models = [Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, HunyuanDiT, HunyuanDiT1, Flux, FluxSchnell, GenmoMochi]
models += [SVD_img2vid]
diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py
index e2b70e69..54573abb 100644
--- a/comfy/supported_models_base.py
+++ b/comfy/supported_models_base.py
@@ -49,6 +49,7 @@ class BASE:
manual_cast_dtype = None
custom_operations = None
+ scaled_fp8 = None
optimizations = {"fp8": False}
@classmethod
diff --git a/comfy/text_encoders/flux.py b/comfy/text_encoders/flux.py
index b13fa5b4..b945b1aa 100644
--- a/comfy/text_encoders/flux.py
+++ b/comfy/text_encoders/flux.py
@@ -1,15 +1,11 @@
from comfy import sd1_clip
import comfy.text_encoders.t5
+import comfy.text_encoders.sd3_clip
import comfy.model_management
from transformers import T5TokenizerFast
import torch
import os
-class T5XXLModel(sd1_clip.SDClipModel):
- def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options={}):
- textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_config_xxl.json")
- super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5, model_options=model_options)
-
class T5XXLTokenizer(sd1_clip.SDTokenizer):
def __init__(self, embedding_directory=None, tokenizer_data={}):
tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer")
@@ -41,7 +37,7 @@ class FluxClipModel(torch.nn.Module):
dtype_t5 = comfy.model_management.pick_weight_dtype(dtype_t5, dtype, device)
clip_l_class = model_options.get("clip_l_class", sd1_clip.SDClipModel)
self.clip_l = clip_l_class(device=device, dtype=dtype, return_projected_pooled=False, model_options=model_options)
- self.t5xxl = T5XXLModel(device=device, dtype=dtype_t5, model_options=model_options)
+ self.t5xxl = comfy.text_encoders.sd3_clip.T5XXLModel(device=device, dtype=dtype_t5, model_options=model_options)
self.dtypes = set([dtype, dtype_t5])
def set_clip_options(self, options):
@@ -66,8 +62,11 @@ class FluxClipModel(torch.nn.Module):
else:
return self.t5xxl.load_sd(sd)
-def flux_clip(dtype_t5=None):
+def flux_clip(dtype_t5=None, t5xxl_scaled_fp8=None):
class FluxClipModel_(FluxClipModel):
def __init__(self, device="cpu", dtype=None, model_options={}):
+ if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
+ model_options = model_options.copy()
+ model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
super().__init__(dtype_t5=dtype_t5, device=device, dtype=dtype, model_options=model_options)
return FluxClipModel_
diff --git a/comfy/text_encoders/genmo.py b/comfy/text_encoders/genmo.py
new file mode 100644
index 00000000..5e96cea6
--- /dev/null
+++ b/comfy/text_encoders/genmo.py
@@ -0,0 +1,38 @@
+from comfy import sd1_clip
+import comfy.text_encoders.sd3_clip
+import os
+from transformers import T5TokenizerFast
+
+
+class T5XXLModel(comfy.text_encoders.sd3_clip.T5XXLModel):
+ def __init__(self, **kwargs):
+ kwargs["attention_mask"] = True
+ super().__init__(**kwargs)
+
+
+class MochiT5XXL(sd1_clip.SD1ClipModel):
+ def __init__(self, device="cpu", dtype=None, model_options={}):
+ super().__init__(device=device, dtype=dtype, clip_name="t5xxl", clip_model=T5XXLModel, model_options=model_options)
+
+
+class T5XXLTokenizer(sd1_clip.SDTokenizer):
+ def __init__(self, embedding_directory=None, tokenizer_data={}):
+ tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer")
+ super().__init__(tokenizer_path, embedding_directory=embedding_directory, pad_with_end=False, embedding_size=4096, embedding_key='t5xxl', tokenizer_class=T5TokenizerFast, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=256)
+
+
+class MochiT5Tokenizer(sd1_clip.SD1Tokenizer):
+ def __init__(self, embedding_directory=None, tokenizer_data={}):
+ super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, clip_name="t5xxl", tokenizer=T5XXLTokenizer)
+
+
+def mochi_te(dtype_t5=None, t5xxl_scaled_fp8=None):
+ class MochiTEModel_(MochiT5XXL):
+ def __init__(self, device="cpu", dtype=None, model_options={}):
+ if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
+ model_options = model_options.copy()
+ model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
+ if dtype is None:
+ dtype = dtype_t5
+ super().__init__(device=device, dtype=dtype, model_options=model_options)
+ return MochiTEModel_
diff --git a/comfy/text_encoders/sd3_clip.py b/comfy/text_encoders/sd3_clip.py
index 0340e65b..00d7e31a 100644
--- a/comfy/text_encoders/sd3_clip.py
+++ b/comfy/text_encoders/sd3_clip.py
@@ -10,8 +10,26 @@ import logging
class T5XXLModel(sd1_clip.SDClipModel):
def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=False, model_options={}):
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_config_xxl.json")
+ t5xxl_scaled_fp8 = model_options.get("t5xxl_scaled_fp8", None)
+ if t5xxl_scaled_fp8 is not None:
+ model_options = model_options.copy()
+ model_options["scaled_fp8"] = t5xxl_scaled_fp8
+
super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
+
+def t5_xxl_detect(state_dict, prefix=""):
+ out = {}
+ t5_key = "{}encoder.final_layer_norm.weight".format(prefix)
+ if t5_key in state_dict:
+ out["dtype_t5"] = state_dict[t5_key].dtype
+
+ scaled_fp8_key = "{}scaled_fp8".format(prefix)
+ if scaled_fp8_key in state_dict:
+ out["t5xxl_scaled_fp8"] = state_dict[scaled_fp8_key].dtype
+
+ return out
+
class T5XXLTokenizer(sd1_clip.SDTokenizer):
def __init__(self, embedding_directory=None, tokenizer_data={}):
tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer")
@@ -139,8 +157,11 @@ class SD3ClipModel(torch.nn.Module):
else:
return self.t5xxl.load_sd(sd)
-def sd3_clip(clip_l=True, clip_g=True, t5=True, dtype_t5=None, t5_attention_mask=False):
+def sd3_clip(clip_l=True, clip_g=True, t5=True, dtype_t5=None, t5xxl_scaled_fp8=None, t5_attention_mask=False):
class SD3ClipModel_(SD3ClipModel):
def __init__(self, device="cpu", dtype=None, model_options={}):
+ if t5xxl_scaled_fp8 is not None and "t5xxl_scaled_fp8" not in model_options:
+ model_options = model_options.copy()
+ model_options["t5xxl_scaled_fp8"] = t5xxl_scaled_fp8
super().__init__(clip_l=clip_l, clip_g=clip_g, t5=t5, dtype_t5=dtype_t5, t5_attention_mask=t5_attention_mask, device=device, dtype=dtype, model_options=model_options)
return SD3ClipModel_
diff --git a/comfy/utils.py b/comfy/utils.py
index 78f8314f..06c81107 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -68,7 +68,7 @@ def weight_dtype(sd, prefix=""):
for k in sd.keys():
if k.startswith(prefix):
w = sd[k]
- dtypes[w.dtype] = dtypes.get(w.dtype, 0) + 1
+ dtypes[w.dtype] = dtypes.get(w.dtype, 0) + w.numel()
if len(dtypes) == 0:
return None
@@ -690,9 +690,14 @@ def lanczos(samples, width, height):
return result.to(samples.device, samples.dtype)
def common_upscale(samples, width, height, upscale_method, crop):
+ orig_shape = tuple(samples.shape)
+ if len(orig_shape) > 4:
+ samples = samples.reshape(samples.shape[0], samples.shape[1], -1, samples.shape[-2], samples.shape[-1])
+ samples = samples.movedim(2, 1)
+ samples = samples.reshape(-1, orig_shape[1], orig_shape[-2], orig_shape[-1])
if crop == "center":
- old_width = samples.shape[3]
- old_height = samples.shape[2]
+ old_width = samples.shape[-1]
+ old_height = samples.shape[-2]
old_aspect = old_width / old_height
new_aspect = width / height
x = 0
@@ -701,16 +706,22 @@ def common_upscale(samples, width, height, upscale_method, crop):
x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)
elif old_aspect < new_aspect:
y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)
- s = samples[:,:,y:old_height-y,x:old_width-x]
+ s = samples.narrow(-2, y, old_height - y * 2).narrow(-1, x, old_width - x * 2)
else:
s = samples
if upscale_method == "bislerp":
- return bislerp(s, width, height)
+ out = bislerp(s, width, height)
elif upscale_method == "lanczos":
- return lanczos(s, width, height)
+ out = lanczos(s, width, height)
else:
- return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
+ out = torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
+
+ if len(orig_shape) == 4:
+ return out
+
+ out = out.reshape((orig_shape[0], -1, orig_shape[1]) + (height, width))
+ return out.movedim(2, 1).reshape(orig_shape[:-2] + (height, width))
def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):
rows = 1 if height <= tile_y else math.ceil((height - overlap) / (tile_y - overlap))
@@ -720,7 +731,27 @@ def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):
@torch.inference_mode()
def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
dims = len(tile)
- output = torch.empty([samples.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), samples.shape[2:])), device=output_device)
+
+ if not (isinstance(upscale_amount, (tuple, list))):
+ upscale_amount = [upscale_amount] * dims
+
+ if not (isinstance(overlap, (tuple, list))):
+ overlap = [overlap] * dims
+
+ def get_upscale(dim, val):
+ up = upscale_amount[dim]
+ if callable(up):
+ return up(val)
+ else:
+ return up * val
+
+ def mult_list_upscale(a):
+ out = []
+ for i in range(len(a)):
+ out.append(round(get_upscale(i, a[i])))
+ return out
+
+ output = torch.empty([samples.shape[0], out_channels] + mult_list_upscale(samples.shape[2:]), device=output_device)
for b in range(samples.shape[0]):
s = samples[b:b+1]
@@ -732,27 +763,27 @@ def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_
pbar.update(1)
continue
- out = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device)
- out_div = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device)
+ out = torch.zeros([s.shape[0], out_channels] + mult_list_upscale(s.shape[2:]), device=output_device)
+ out_div = torch.zeros([s.shape[0], out_channels] + mult_list_upscale(s.shape[2:]), device=output_device)
- positions = [range(0, s.shape[d+2], tile[d] - overlap) if s.shape[d+2] > tile[d] else [0] for d in range(dims)]
+ positions = [range(0, s.shape[d+2], tile[d] - overlap[d]) if s.shape[d+2] > tile[d] else [0] for d in range(dims)]
for it in itertools.product(*positions):
s_in = s
upscaled = []
for d in range(dims):
- pos = max(0, min(s.shape[d + 2] - overlap, it[d]))
+ pos = max(0, min(s.shape[d + 2] - (overlap[d] + 1), it[d]))
l = min(tile[d], s.shape[d + 2] - pos)
s_in = s_in.narrow(d + 2, pos, l)
- upscaled.append(round(pos * upscale_amount))
+ upscaled.append(round(get_upscale(d, pos)))
ps = function(s_in).to(output_device)
mask = torch.ones_like(ps)
- feather = round(overlap * upscale_amount)
- for t in range(feather):
- for d in range(2, dims + 2):
+ for d in range(2, dims + 2):
+ feather = round(get_upscale(d - 2, overlap[d - 2]))
+ for t in range(feather):
a = (t + 1) / feather
mask.narrow(d, t, 1).mul_(a)
mask.narrow(d, mask.shape[d] - 1 - t, 1).mul_(a)
diff --git a/comfy_extras/nodes_latent.py b/comfy_extras/nodes_latent.py
index eabae088..af273681 100644
--- a/comfy_extras/nodes_latent.py
+++ b/comfy_extras/nodes_latent.py
@@ -1,4 +1,5 @@
import comfy.utils
+import comfy_extras.nodes_post_processing
import torch
def reshape_latent_to(target_shape, latent):
@@ -145,6 +146,131 @@ class LatentBatchSeedBehavior:
return (samples_out,)
+class LatentApplyOperation:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required": { "samples": ("LATENT",),
+ "operation": ("LATENT_OPERATION",),
+ }}
+
+ RETURN_TYPES = ("LATENT",)
+ FUNCTION = "op"
+
+ CATEGORY = "latent/advanced/operations"
+ EXPERIMENTAL = True
+
+ def op(self, samples, operation):
+ samples_out = samples.copy()
+
+ s1 = samples["samples"]
+ samples_out["samples"] = operation(latent=s1)
+ return (samples_out,)
+
+class LatentApplyOperationCFG:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required": { "model": ("MODEL",),
+ "operation": ("LATENT_OPERATION",),
+ }}
+ RETURN_TYPES = ("MODEL",)
+ FUNCTION = "patch"
+
+ CATEGORY = "latent/advanced/operations"
+ EXPERIMENTAL = True
+
+ def patch(self, model, operation):
+ m = model.clone()
+
+ def pre_cfg_function(args):
+ conds_out = args["conds_out"]
+ if len(conds_out) == 2:
+ conds_out[0] = operation(latent=(conds_out[0] - conds_out[1])) + conds_out[1]
+ else:
+ conds_out[0] = operation(latent=conds_out[0])
+ return conds_out
+
+ m.set_model_sampler_pre_cfg_function(pre_cfg_function)
+ return (m, )
+
+class LatentOperationTonemapReinhard:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required": { "multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+ }}
+
+ RETURN_TYPES = ("LATENT_OPERATION",)
+ FUNCTION = "op"
+
+ CATEGORY = "latent/advanced/operations"
+ EXPERIMENTAL = True
+
+ def op(self, multiplier):
+ def tonemap_reinhard(latent, **kwargs):
+ latent_vector_magnitude = (torch.linalg.vector_norm(latent, dim=(1)) + 0.0000000001)[:,None]
+ normalized_latent = latent / latent_vector_magnitude
+
+ mean = torch.mean(latent_vector_magnitude, dim=(1,2,3), keepdim=True)
+ std = torch.std(latent_vector_magnitude, dim=(1,2,3), keepdim=True)
+
+ top = (std * 5 + mean) * multiplier
+
+ #reinhard
+ latent_vector_magnitude *= (1.0 / top)
+ new_magnitude = latent_vector_magnitude / (latent_vector_magnitude + 1.0)
+ new_magnitude *= top
+
+ return normalized_latent * new_magnitude
+ return (tonemap_reinhard,)
+
+class LatentOperationSharpen:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required": {
+ "sharpen_radius": ("INT", {
+ "default": 9,
+ "min": 1,
+ "max": 31,
+ "step": 1
+ }),
+ "sigma": ("FLOAT", {
+ "default": 1.0,
+ "min": 0.1,
+ "max": 10.0,
+ "step": 0.1
+ }),
+ "alpha": ("FLOAT", {
+ "default": 0.1,
+ "min": 0.0,
+ "max": 5.0,
+ "step": 0.01
+ }),
+ }}
+
+ RETURN_TYPES = ("LATENT_OPERATION",)
+ FUNCTION = "op"
+
+ CATEGORY = "latent/advanced/operations"
+ EXPERIMENTAL = True
+
+ def op(self, sharpen_radius, sigma, alpha):
+ def sharpen(latent, **kwargs):
+ luminance = (torch.linalg.vector_norm(latent, dim=(1)) + 1e-6)[:,None]
+ normalized_latent = latent / luminance
+ channels = latent.shape[1]
+
+ kernel_size = sharpen_radius * 2 + 1
+ kernel = comfy_extras.nodes_post_processing.gaussian_kernel(kernel_size, sigma, device=luminance.device)
+ center = kernel_size // 2
+
+ kernel *= alpha * -10
+ kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0
+
+ padded_image = torch.nn.functional.pad(normalized_latent, (sharpen_radius,sharpen_radius,sharpen_radius,sharpen_radius), 'reflect')
+ sharpened = torch.nn.functional.conv2d(padded_image, kernel.repeat(channels, 1, 1).unsqueeze(1), padding=kernel_size // 2, groups=channels)[:,:,sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius]
+
+ return luminance * sharpened
+ return (sharpen,)
+
NODE_CLASS_MAPPINGS = {
"LatentAdd": LatentAdd,
"LatentSubtract": LatentSubtract,
@@ -152,4 +278,8 @@ NODE_CLASS_MAPPINGS = {
"LatentInterpolate": LatentInterpolate,
"LatentBatch": LatentBatch,
"LatentBatchSeedBehavior": LatentBatchSeedBehavior,
+ "LatentApplyOperation": LatentApplyOperation,
+ "LatentApplyOperationCFG": LatentApplyOperationCFG,
+ "LatentOperationTonemapReinhard": LatentOperationTonemapReinhard,
+ "LatentOperationSharpen": LatentOperationSharpen,
}
diff --git a/comfy_extras/nodes_lora_extract.py b/comfy_extras/nodes_lora_extract.py
index 3c2f179d..dfd4fe9f 100644
--- a/comfy_extras/nodes_lora_extract.py
+++ b/comfy_extras/nodes_lora_extract.py
@@ -82,8 +82,8 @@ class LoraSave:
"lora_type": (tuple(LORA_TYPES.keys()),),
"bias_diff": ("BOOLEAN", {"default": True}),
},
- "optional": {"model_diff": ("MODEL",),
- "text_encoder_diff": ("CLIP",)},
+ "optional": {"model_diff": ("MODEL", {"tooltip": "The ModelSubtract output to be converted to a lora."}),
+ "text_encoder_diff": ("CLIP", {"tooltip": "The CLIPSubtract output to be converted to a lora."})},
}
RETURN_TYPES = ()
FUNCTION = "save"
@@ -113,3 +113,7 @@ class LoraSave:
NODE_CLASS_MAPPINGS = {
"LoraSave": LoraSave
}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "LoraSave": "Extract and Save Lora"
+}
diff --git a/comfy_extras/nodes_mochi.py b/comfy_extras/nodes_mochi.py
new file mode 100644
index 00000000..4cbbea09
--- /dev/null
+++ b/comfy_extras/nodes_mochi.py
@@ -0,0 +1,26 @@
+import nodes
+import torch
+import comfy.model_management
+
+class EmptyMochiLatentVideo:
+ def __init__(self):
+ self.device = comfy.model_management.intermediate_device()
+
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required": { "width": ("INT", {"default": 848, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+ "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}),
+ "length": ("INT", {"default": 25, "min": 7, "max": nodes.MAX_RESOLUTION, "step": 6}),
+ "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
+ RETURN_TYPES = ("LATENT",)
+ FUNCTION = "generate"
+
+ CATEGORY = "latent/mochi"
+
+ def generate(self, width, height, length, batch_size=1):
+ latent = torch.zeros([batch_size, 12, ((length - 1) // 6) + 1, height // 8, width // 8], device=self.device)
+ return ({"samples":latent}, )
+
+NODE_CLASS_MAPPINGS = {
+ "EmptyMochiLatentVideo": EmptyMochiLatentVideo,
+}
diff --git a/comfy_extras/nodes_model_merging_model_specific.py b/comfy_extras/nodes_model_merging_model_specific.py
index 9557d9b1..0e847192 100644
--- a/comfy_extras/nodes_model_merging_model_specific.py
+++ b/comfy_extras/nodes_model_merging_model_specific.py
@@ -101,10 +101,34 @@ class ModelMergeFlux1(comfy_extras.nodes_model_merging.ModelMergeBlocks):
return {"required": arg_dict}
+class ModelMergeSD35_Large(comfy_extras.nodes_model_merging.ModelMergeBlocks):
+ CATEGORY = "advanced/model_merging/model_specific"
+
+ @classmethod
+ def INPUT_TYPES(s):
+ arg_dict = { "model1": ("MODEL",),
+ "model2": ("MODEL",)}
+
+ argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
+
+ arg_dict["pos_embed."] = argument
+ arg_dict["x_embedder."] = argument
+ arg_dict["context_embedder."] = argument
+ arg_dict["y_embedder."] = argument
+ arg_dict["t_embedder."] = argument
+
+ for i in range(38):
+ arg_dict["joint_blocks.{}.".format(i)] = argument
+
+ arg_dict["final_layer."] = argument
+
+ return {"required": arg_dict}
+
NODE_CLASS_MAPPINGS = {
"ModelMergeSD1": ModelMergeSD1,
"ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks
"ModelMergeSDXL": ModelMergeSDXL,
"ModelMergeSD3_2B": ModelMergeSD3_2B,
"ModelMergeFlux1": ModelMergeFlux1,
+ "ModelMergeSD35_Large": ModelMergeSD35_Large,
}
diff --git a/comfy_extras/nodes_sd3.py b/comfy_extras/nodes_sd3.py
index ddf538de..4d664093 100644
--- a/comfy_extras/nodes_sd3.py
+++ b/comfy_extras/nodes_sd3.py
@@ -3,7 +3,7 @@ import comfy.sd
import comfy.model_management
import nodes
import torch
-
+import re
class TripleCLIPLoader:
@classmethod
def INPUT_TYPES(s):
@@ -95,11 +95,70 @@ class ControlNetApplySD3(nodes.ControlNetApplyAdvanced):
CATEGORY = "conditioning/controlnet"
DEPRECATED = True
+class SkipLayerGuidanceSD3:
+ '''
+ Enhance guidance towards detailed dtructure by having another set of CFG negative with skipped layers.
+ Inspired by Perturbed Attention Guidance (https://arxiv.org/abs/2403.17377)
+ Experimental implementation by Dango233@StabilityAI.
+ '''
+ @classmethod
+ def INPUT_TYPES(s):
+ return {"required": {"model": ("MODEL", ),
+ "layers": ("STRING", {"default": "7, 8, 9", "multiline": False}),
+ "scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1}),
+ "start_percent": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.001}),
+ "end_percent": ("FLOAT", {"default": 0.15, "min": 0.0, "max": 1.0, "step": 0.001})
+ }}
+ RETURN_TYPES = ("MODEL",)
+ FUNCTION = "skip_guidance"
+
+ CATEGORY = "advanced/guidance"
+
+
+ def skip_guidance(self, model, layers, scale, start_percent, end_percent):
+ if layers == "" or layers == None:
+ return (model, )
+ # check if layer is comma separated integers
+ def skip(args, extra_args):
+ return args
+
+ model_sampling = model.get_model_object("model_sampling")
+ sigma_start = model_sampling.percent_to_sigma(start_percent)
+ sigma_end = model_sampling.percent_to_sigma(end_percent)
+
+ def post_cfg_function(args):
+ model = args["model"]
+ cond_pred = args["cond_denoised"]
+ cond = args["cond"]
+ cfg_result = args["denoised"]
+ sigma = args["sigma"]
+ x = args["input"]
+ model_options = args["model_options"].copy()
+
+ for layer in layers:
+ model_options = comfy.model_patcher.set_model_options_patch_replace(model_options, skip, "dit", "double_block", layer)
+ model_sampling.percent_to_sigma(start_percent)
+
+ sigma_ = sigma[0].item()
+ if scale > 0 and sigma_ >= sigma_end and sigma_ <= sigma_start:
+ (slg,) = comfy.samplers.calc_cond_batch(model, [cond], x, sigma, model_options)
+ cfg_result = cfg_result + (cond_pred - slg) * scale
+ return cfg_result
+
+ layers = re.findall(r'\d+', layers)
+ layers = [int(i) for i in layers]
+ m = model.clone()
+ m.set_model_sampler_post_cfg_function(post_cfg_function)
+
+ return (m, )
+
+
NODE_CLASS_MAPPINGS = {
"TripleCLIPLoader": TripleCLIPLoader,
"EmptySD3LatentImage": EmptySD3LatentImage,
"CLIPTextEncodeSD3": CLIPTextEncodeSD3,
"ControlNetApplySD3": ControlNetApplySD3,
+ "SkipLayerGuidanceSD3": SkipLayerGuidanceSD3,
}
NODE_DISPLAY_NAME_MAPPINGS = {
diff --git a/nodes.py b/nodes.py
index 15a78352..c81a0af1 100644
--- a/nodes.py
+++ b/nodes.py
@@ -281,7 +281,10 @@ class VAEDecode:
DESCRIPTION = "Decodes latent images back into pixel space images."
def decode(self, vae, samples):
- return (vae.decode(samples["samples"]), )
+ images = vae.decode(samples["samples"])
+ if len(images.shape) == 5: #Combine batches
+ images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
+ return (images, )
class VAEDecodeTiled:
@classmethod
@@ -886,7 +889,7 @@ class CLIPLoader:
@classmethod
def INPUT_TYPES(s):
return {"required": { "clip_name": (folder_paths.get_filename_list("clip"), ),
- "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio"], ),
+ "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi"], ),
}}
RETURN_TYPES = ("CLIP",)
FUNCTION = "load_clip"
@@ -900,6 +903,8 @@ class CLIPLoader:
clip_type = comfy.sd.CLIPType.SD3
elif type == "stable_audio":
clip_type = comfy.sd.CLIPType.STABLE_AUDIO
+ elif type == "mochi":
+ clip_type = comfy.sd.CLIPType.MOCHI
else:
clip_type = comfy.sd.CLIPType.STABLE_DIFFUSION
@@ -1179,10 +1184,10 @@ class LatentUpscale:
if width == 0:
height = max(64, height)
- width = max(64, round(samples["samples"].shape[3] * height / samples["samples"].shape[2]))
+ width = max(64, round(samples["samples"].shape[-1] * height / samples["samples"].shape[-2]))
elif height == 0:
width = max(64, width)
- height = max(64, round(samples["samples"].shape[2] * width / samples["samples"].shape[3]))
+ height = max(64, round(samples["samples"].shape[-2] * width / samples["samples"].shape[-1]))
else:
width = max(64, width)
height = max(64, height)
@@ -1204,8 +1209,8 @@ class LatentUpscaleBy:
def upscale(self, samples, upscale_method, scale_by):
s = samples.copy()
- width = round(samples["samples"].shape[3] * scale_by)
- height = round(samples["samples"].shape[2] * scale_by)
+ width = round(samples["samples"].shape[-1] * scale_by)
+ height = round(samples["samples"].shape[-2] * scale_by)
s["samples"] = comfy.utils.common_upscale(samples["samples"], width, height, upscale_method, "disabled")
return (s,)
@@ -2111,6 +2116,7 @@ def init_builtin_extra_nodes():
"nodes_flux.py",
"nodes_lora_extract.py",
"nodes_torch_compile.py",
+ "nodes_mochi.py",
]
import_failed = []
diff --git a/server.py b/server.py
index c7bf6622..ada6d90c 100644
--- a/server.py
+++ b/server.py
@@ -40,7 +40,7 @@ class BinaryEventTypes:
async def send_socket_catch_exception(function, message):
try:
await function(message)
- except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError) as err:
+ except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError, BrokenPipeError, ConnectionError) as err:
logging.warning("send error: {}".format(err))
def get_comfyui_version():
diff --git a/web/assets/ExtensionPanel-DZLYjWBj.js b/web/assets/ExtensionPanel-DZLYjWBj.js
new file mode 100644
index 00000000..36e944c3
--- /dev/null
+++ b/web/assets/ExtensionPanel-DZLYjWBj.js
@@ -0,0 +1,103 @@
+var __defProp = Object.defineProperty;
+var __name = (target, value) => __defProp(target, "name", { value, configurable: true });
+import { d as defineComponent, bK as useExtensionStore, u as useSettingStore, r as ref, o as onMounted, q as computed, g as openBlock, h as createElementBlock, i as createVNode, y as withCtx, z as unref, bL as script$1, A as createBaseVNode, x as createBlock, M as Fragment, N as renderList, am as toDisplayString, ap as createTextVNode, j as createCommentVNode, D as script$4 } from "./index-CgU1oKZt.js";
+import { s as script, a as script$2, b as script$3 } from "./index-DBWDcZsl.js";
+import "./index-DYEEBf64.js";
+const _hoisted_1 = { class: "extension-panel" };
+const _hoisted_2 = { class: "mt-4" };
+const _sfc_main = /* @__PURE__ */ defineComponent({
+ __name: "ExtensionPanel",
+ setup(__props) {
+ const extensionStore = useExtensionStore();
+ const settingStore = useSettingStore();
+ const editingEnabledExtensions = ref({});
+ onMounted(() => {
+ extensionStore.extensions.forEach((ext) => {
+ editingEnabledExtensions.value[ext.name] = extensionStore.isExtensionEnabled(ext.name);
+ });
+ });
+ const changedExtensions = computed(() => {
+ return extensionStore.extensions.filter(
+ (ext) => editingEnabledExtensions.value[ext.name] !== extensionStore.isExtensionEnabled(ext.name)
+ );
+ });
+ const hasChanges = computed(() => {
+ return changedExtensions.value.length > 0;
+ });
+ const updateExtensionStatus = /* @__PURE__ */ __name(() => {
+ const editingDisabledExtensionNames = Object.entries(
+ editingEnabledExtensions.value
+ ).filter(([_, enabled]) => !enabled).map(([name]) => name);
+ settingStore.set("Comfy.Extension.Disabled", [
+ ...extensionStore.inactiveDisabledExtensionNames,
+ ...editingDisabledExtensionNames
+ ]);
+ }, "updateExtensionStatus");
+ const applyChanges = /* @__PURE__ */ __name(() => {
+ window.location.reload();
+ }, "applyChanges");
+ return (_ctx, _cache) => {
+ return openBlock(), createElementBlock("div", _hoisted_1, [
+ createVNode(unref(script$2), {
+ value: unref(extensionStore).extensions,
+ stripedRows: "",
+ size: "small"
+ }, {
+ default: withCtx(() => [
+ createVNode(unref(script), {
+ field: "name",
+ header: _ctx.$t("extensionName"),
+ sortable: ""
+ }, null, 8, ["header"]),
+ createVNode(unref(script), { pt: {
+ bodyCell: "flex items-center justify-end"
+ } }, {
+ body: withCtx((slotProps) => [
+ createVNode(unref(script$1), {
+ modelValue: editingEnabledExtensions.value[slotProps.data.name],
+ "onUpdate:modelValue": /* @__PURE__ */ __name(($event) => editingEnabledExtensions.value[slotProps.data.name] = $event, "onUpdate:modelValue"),
+ onChange: updateExtensionStatus
+ }, null, 8, ["modelValue", "onUpdate:modelValue"])
+ ]),
+ _: 1
+ })
+ ]),
+ _: 1
+ }, 8, ["value"]),
+ createBaseVNode("div", _hoisted_2, [
+ hasChanges.value ? (openBlock(), createBlock(unref(script$3), {
+ key: 0,
+ severity: "info"
+ }, {
+ default: withCtx(() => [
+ createBaseVNode("ul", null, [
+ (openBlock(true), createElementBlock(Fragment, null, renderList(changedExtensions.value, (ext) => {
+ return openBlock(), createElementBlock("li", {
+ key: ext.name
+ }, [
+ createBaseVNode("span", null, toDisplayString(unref(extensionStore).isExtensionEnabled(ext.name) ? "[-]" : "[+]"), 1),
+ createTextVNode(" " + toDisplayString(ext.name), 1)
+ ]);
+ }), 128))
+ ])
+ ]),
+ _: 1
+ })) : createCommentVNode("", true),
+ createVNode(unref(script$4), {
+ label: _ctx.$t("reloadToApplyChanges"),
+ icon: "pi pi-refresh",
+ onClick: applyChanges,
+ disabled: !hasChanges.value,
+ text: "",
+ fluid: "",
+ severity: "danger"
+ }, null, 8, ["label", "disabled"])
+ ])
+ ]);
+ };
+ }
+});
+export {
+ _sfc_main as default
+};
+//# sourceMappingURL=ExtensionPanel-DZLYjWBj.js.map
diff --git a/web/assets/ExtensionPanel-DZLYjWBj.js.map b/web/assets/ExtensionPanel-DZLYjWBj.js.map
new file mode 100644
index 00000000..eabe9cba
--- /dev/null
+++ b/web/assets/ExtensionPanel-DZLYjWBj.js.map
@@ -0,0 +1 @@
+{"version":3,"file":"ExtensionPanel-DZLYjWBj.js","sources":["../../src/components/dialog/content/setting/ExtensionPanel.vue"],"sourcesContent":["\n \n
\n \n \n \n \n \n \n \n
\n
\n \n - \n \n {{ extensionStore.isExtensionEnabled(ext.name) ? '[-]' : '[+]' }}\n \n {{ ext.name }}\n
\n
\n \n
\n
\n