From 7c59a078d60c85baded8789f10c841221bff80a8 Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:17:31 -0700 Subject: [PATCH 1/2] Use comfy kitchen rope functions in ltx models. (#15056) --- comfy/ldm/lightricks/embeddings_connector.py | 16 ++- comfy/ldm/lightricks/model.py | 134 +++++++++---------- 2 files changed, 76 insertions(+), 74 deletions(-) diff --git a/comfy/ldm/lightricks/embeddings_connector.py b/comfy/ldm/lightricks/embeddings_connector.py index 2811080be0a..1a6ddcc8de9 100644 --- a/comfy/ldm/lightricks/embeddings_connector.py +++ b/comfy/ldm/lightricks/embeddings_connector.py @@ -6,9 +6,8 @@ from comfy.ldm.lightricks.model import ( CrossAttention, FeedForward, + freqs_cis_matrix, generate_freq_grid_np, - interleaved_freqs_cis, - split_freqs_cis, ) from torch import nn @@ -244,12 +243,15 @@ def precompute_freqs_cis(self, indices_grid, spacing="exp", out_dtype=None): expected_freqs = dim // 2 current_freqs = freqs.shape[-1] pad_size = expected_freqs - current_freqs - cos_freq, sin_freq = split_freqs_cis( - freqs, pad_size, self.num_attention_heads - ) else: - cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem) - return cos_freq.to(dtype=out_dtype), sin_freq.to(dtype=out_dtype), self.split_rope + pad_size = dim % n_elem + return freqs_cis_matrix( + freqs, + pad_size, + self.split_rope, + self.num_attention_heads, + out_dtype, + ) def forward( self, diff --git a/comfy/ldm/lightricks/model.py b/comfy/ldm/lightricks/model.py index 9953b667965..92bb8118c24 100644 --- a/comfy/ldm/lightricks/model.py +++ b/comfy/ldm/lightricks/model.py @@ -12,6 +12,8 @@ import comfy.patcher_extension import comfy.ldm.modules.attention import comfy.ldm.common_dit +import comfy.model_management +import comfy.quant_ops from .symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords @@ -322,40 +324,42 @@ def forward(self, x): return self.net(x) def apply_rotary_emb(input_tensor, freqs_cis): - cos_freqs, sin_freqs = freqs_cis[0], freqs_cis[1] - split_pe = freqs_cis[2] if len(freqs_cis) > 2 else False - return ( - apply_split_rotary_emb(input_tensor, cos_freqs, sin_freqs) - if split_pe else - apply_interleaved_rotary_emb(input_tensor, cos_freqs, sin_freqs) + rotation_matrix, split_pe = freqs_cis + original_shape = input_tensor.shape + input_tensor = input_tensor.reshape( + input_tensor.shape[0], input_tensor.shape[1], rotation_matrix.shape[2], -1 ) -def apply_interleaved_rotary_emb(input_tensor, cos_freqs, sin_freqs): # TODO: remove duplicate funcs and pick the best/fastest one - t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2) - t1, t2 = t_dup.unbind(dim=-1) - t_dup = torch.stack((-t2, t1), dim=-1) - input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)") - - out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs - - return out - -def apply_split_rotary_emb(input_tensor, cos, sin): - needs_reshape = False - if input_tensor.ndim != 4 and cos.ndim == 4: - B, H, T, _ = cos.shape - input_tensor = input_tensor.reshape(B, T, H, -1).swapaxes(1, 2) - needs_reshape = True - split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2) - first_half_input = split_input[..., :1, :] - second_half_input = split_input[..., 1:, :] - output = split_input * cos.unsqueeze(-2) - first_half_output = output[..., :1, :] - second_half_output = output[..., 1:, :] - first_half_output.addcmul_(-sin.unsqueeze(-2), second_half_input) - second_half_output.addcmul_(sin.unsqueeze(-2), first_half_input) - output = rearrange(output, "... d r -> ... (d r)") - return output.swapaxes(1, 2).reshape(B, T, -1) if needs_reshape else output + if comfy.model_management.in_training: + if split_pe: + t = input_tensor.reshape(*input_tensor.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2) + else: + t = input_tensor.reshape(*input_tensor.shape[:-1], -1, 1, 2) + t = t.to(rotation_matrix.dtype) + output = rotation_matrix[..., 0] * t[..., 0] + rotation_matrix[..., 1] * t[..., 1] + if split_pe: + output = output.movedim(-1, -2) + output = output.reshape(input_tensor.shape).type_as(input_tensor) + elif split_pe: + output = comfy.quant_ops.ck.apply_rope_split_half1(input_tensor, rotation_matrix) + else: + output = comfy.quant_ops.ck.apply_rope1(input_tensor, rotation_matrix) + return output.reshape(original_shape) + +def apply_rotary_emb_qk(q, k, freqs_cis): + if comfy.model_management.in_training: + return apply_rotary_emb(q, freqs_cis), apply_rotary_emb(k, freqs_cis) + + rotation_matrix, split_pe = freqs_cis + q_shape = q.shape + k_shape = k.shape + q = q.reshape(q.shape[0], q.shape[1], rotation_matrix.shape[2], -1) + k = k.reshape(k.shape[0], k.shape[1], rotation_matrix.shape[2], -1) + if split_pe: + q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rotation_matrix) + else: + q, k = comfy.quant_ops.ck.apply_rope(q, k, rotation_matrix) + return q.reshape(q_shape), k.reshape(k_shape) class GuideAttentionMask: @@ -461,9 +465,13 @@ def forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_op q = self.q_norm(q) k = self.k_norm(k) + # These norms span all heads, so the per-head RMS+RoPE kernel is not equivalent. if pe is not None: - q = apply_rotary_emb(q, pe) - k = apply_rotary_emb(k, pe if k_pe is None else k_pe) + if k_pe is None and q.shape == k.shape: + q, k = apply_rotary_emb_qk(q, k, pe) + else: + q = apply_rotary_emb(q, pe) + k = apply_rotary_emb(k, pe if k_pe is None else k_pe) if mask is None: out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options) @@ -653,36 +661,23 @@ def generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid): ) return freqs -def interleaved_freqs_cis(freqs, pad_size): - cos_freq = freqs.cos().repeat_interleave(2, dim=-1) - sin_freq = freqs.sin().repeat_interleave(2, dim=-1) - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, : pad_size]) - sin_padding = torch.zeros_like(cos_freq[:, :, : pad_size]) - cos_freq = torch.cat([cos_padding, cos_freq], dim=-1) - sin_freq = torch.cat([sin_padding, sin_freq], dim=-1) - return cos_freq, sin_freq - -def split_freqs_cis(freqs, pad_size, num_attention_heads): - cos_freq = freqs.cos() - sin_freq = freqs.sin() - - if pad_size != 0: - cos_padding = torch.ones_like(cos_freq[:, :, :pad_size]) - sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size]) - - cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1) - sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1) - - # Reshape freqs to be compatible with multi-head attention - B , T, half_HD = cos_freq.shape - - cos_freq = cos_freq.reshape(B, T, num_attention_heads, half_HD // num_attention_heads) - sin_freq = sin_freq.reshape(B, T, num_attention_heads, half_HD // num_attention_heads) - - cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2) - sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2) - return cos_freq, sin_freq +def freqs_cis_matrix(freqs, pad_size, split_mode, num_attention_heads, out_dtype): + cos_freq = freqs.cos().to(out_dtype) + sin_freq = freqs.sin().to(out_dtype) + if pad_size: + matrix_pad_size = pad_size if split_mode else pad_size // 2 + cos_padding = torch.ones_like(cos_freq[:, :, :matrix_pad_size]) + sin_padding = torch.zeros_like(sin_freq[:, :, :matrix_pad_size]) + cos_freq = torch.cat((cos_padding, cos_freq), dim=-1) + sin_freq = torch.cat((sin_padding, sin_freq), dim=-1) + + B, T, _ = cos_freq.shape + cos_freq = cos_freq.reshape(B, T, num_attention_heads, -1) + sin_freq = sin_freq.reshape(B, T, num_attention_heads, -1) + rotation_matrix = torch.stack( + (cos_freq, -sin_freq, sin_freq, cos_freq), dim=-1 + ) + return rotation_matrix.reshape(*rotation_matrix.shape[:-1], 2, 2), split_mode class LTXBaseModel(torch.nn.Module, ABC): """ @@ -885,12 +880,17 @@ def _precompute_freqs_cis( expected_freqs = dim // 2 current_freqs = freqs.shape[-1] pad_size = expected_freqs - current_freqs - cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads) else: # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only n_elem = 2 * indices_grid.shape[1] - cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem) - return cos_freq.to(out_dtype), sin_freq.to(out_dtype), split_mode + pad_size = dim % n_elem + return freqs_cis_matrix( + freqs, + pad_size, + split_mode, + num_attention_heads, + out_dtype, + ) def _prepare_positional_embeddings(self, pixel_coords, frame_rate, x_dtype): """Prepare positional embeddings.""" From f8a3fd9d79837bd377d4e15e634271b488d9ee26 Mon Sep 17 00:00:00 2001 From: rattus <46076784+rattus128@users.noreply.github.com> Date: Sat, 25 Jul 2026 06:34:40 +1000 Subject: [PATCH 2/2] upscalers: convert latent_upsampler model to DynamicVram (#15063) These were alll non-dynamic (some non-ModelPatcher) code path calling FreeMemory for management requiring up-front memory freeing. Convert it to dynamic to avoid legacy free behaviour mixing into otherwise dynamic workflows. --- comfy/ldm/lightricks/latent_upsampler.py | 35 ++++++++++++++---------- comfy_extras/nodes_hunyuan.py | 9 ++++-- comfy_extras/nodes_lt_upsampler.py | 24 ++++++---------- comfy_extras/nodes_upscale_model.py | 35 +++++++++++------------- 4 files changed, 52 insertions(+), 51 deletions(-) diff --git a/comfy/ldm/lightricks/latent_upsampler.py b/comfy/ldm/lightricks/latent_upsampler.py index 78ed7653f74..6a4beb1bf29 100644 --- a/comfy/ldm/lightricks/latent_upsampler.py +++ b/comfy/ldm/lightricks/latent_upsampler.py @@ -97,11 +97,11 @@ class SpatialRationalResampler(nn.Module): For dims==3, work per-frame for spatial scaling (temporal axis untouched). """ - def __init__(self, mid_channels: int, scale: float): + def __init__(self, mid_channels: int, scale: float, operations): super().__init__() self.scale = float(scale) self.num, self.den = _rational_for_scale(self.scale) - self.conv = nn.Conv2d( + self.conv = operations.Conv2d( mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1 ) self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) @@ -119,18 +119,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ResBlock(nn.Module): def __init__( - self, channels: int, mid_channels: Optional[int] = None, dims: int = 3 + self, channels: int, operations, mid_channels: Optional[int] = None, dims: int = 3 ): super().__init__() if mid_channels is None: mid_channels = channels - Conv = nn.Conv2d if dims == 2 else nn.Conv3d + Conv = operations.Conv2d if dims == 2 else operations.Conv3d self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1) - self.norm1 = nn.GroupNorm(32, mid_channels) + self.norm1 = operations.GroupNorm(32, mid_channels) self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1) - self.norm2 = nn.GroupNorm(32, channels) + self.norm2 = operations.GroupNorm(32, channels) self.activation = nn.SiLU() def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -159,6 +159,7 @@ class LatentUpsampler(nn.Module): def __init__( self, + operations, in_channels: int = 128, mid_channels: int = 512, num_blocks_per_stage: int = 4, @@ -179,34 +180,34 @@ def __init__( self.spatial_scale = float(spatial_scale) self.rational_resampler = rational_resampler - Conv = nn.Conv2d if dims == 2 else nn.Conv3d + Conv = operations.Conv2d if dims == 2 else operations.Conv3d self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1) - self.initial_norm = nn.GroupNorm(32, mid_channels) + self.initial_norm = operations.GroupNorm(32, mid_channels) self.initial_activation = nn.SiLU() self.res_blocks = nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + [ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)] ) if spatial_upsample and temporal_upsample: self.upsampler = nn.Sequential( - nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), + operations.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1), PixelShuffleND(3), ) elif spatial_upsample: if rational_resampler: self.upsampler = SpatialRationalResampler( - mid_channels=mid_channels, scale=self.spatial_scale + mid_channels=mid_channels, scale=self.spatial_scale, operations=operations ) else: self.upsampler = nn.Sequential( - nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), + operations.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1), PixelShuffleND(2), ) elif temporal_upsample: self.upsampler = nn.Sequential( - nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), + operations.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1), PixelShuffleND(1), ) else: @@ -215,11 +216,14 @@ def __init__( ) self.post_upsample_res_blocks = nn.ModuleList( - [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + [ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)] ) self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1) + def get_dtype(self): + return getattr(self.initial_conv, "weight_comfy_model_dtype", self.initial_conv.weight.dtype) + def forward(self, latent: torch.Tensor) -> torch.Tensor: b, c, f, h, w = latent.shape @@ -266,7 +270,7 @@ def forward(self, latent: torch.Tensor) -> torch.Tensor: return x @classmethod - def from_config(cls, config): + def from_config(cls, config, operations): return cls( in_channels=config.get("in_channels", 4), mid_channels=config.get("mid_channels", 128), @@ -276,6 +280,7 @@ def from_config(cls, config): temporal_upsample=config.get("temporal_upsample", False), spatial_scale=config.get("spatial_scale", 2.0), rational_resampler=config.get("rational_resampler", False), + operations=operations, ) def config(self): diff --git a/comfy_extras/nodes_hunyuan.py b/comfy_extras/nodes_hunyuan.py index 8df2c8908b6..ce299724584 100644 --- a/comfy_extras/nodes_hunyuan.py +++ b/comfy_extras/nodes_hunyuan.py @@ -2,6 +2,8 @@ import node_helpers import torch import comfy.model_management +import comfy.model_patcher +import comfy.ops from typing_extensions import override from comfy_api.latest import ComfyExtension, io from comfy.ldm.hunyuan_video.upsampler import HunyuanVideo15SRModel @@ -217,8 +219,11 @@ def execute(cls, model_name) -> io.NodeOutput: model.load_sd(sd) elif "post_upsample_res_blocks.0.conv2.bias" in sd: config = json.loads(metadata["config"]) - model = LatentUpsampler.from_config(config).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32])) - model.load_state_dict(sd) + model = LatentUpsampler.from_config(config, operations=comfy.ops.disable_weight_init).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32])) + comfy.model_management.archive_model_dtypes(model) + model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device()) + model.load_state_dict(sd, assign=model_patcher.is_dynamic()) + model = model_patcher return io.NodeOutput(model) diff --git a/comfy_extras/nodes_lt_upsampler.py b/comfy_extras/nodes_lt_upsampler.py index ef36109d14d..7e7975495cb 100644 --- a/comfy_extras/nodes_lt_upsampler.py +++ b/comfy_extras/nodes_lt_upsampler.py @@ -38,26 +38,20 @@ def execute(cls, samples, upscale_model, vae) -> IO.NodeOutput: Returns: tuple: Tuple containing the upsampled latent """ - device = model_management.get_torch_device() - memory_required = model_management.module_size(upscale_model) - - model_dtype = next(upscale_model.parameters()).dtype + device = upscale_model.load_device + model = upscale_model.model + model_dtype = upscale_model.model_dtype() latents = samples["samples"] input_dtype = latents.dtype - memory_required += math.prod(latents.shape) * 3000.0 # TODO: more accurate - model_management.free_memory(memory_required, device) - - try: - upscale_model.to(device) # TODO: use the comfy model management system. + memory_required = math.prod(latents.shape) * 3000.0 # TODO: more accurate + model_management.load_models_gpu([upscale_model], memory_required=memory_required) - latents = latents.to(dtype=model_dtype, device=device) + latents = latents.to(dtype=model_dtype, device=device) - """Upsample latents without tiling.""" - latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents) - upsampled_latents = upscale_model(latents) - finally: - upscale_model.cpu() + """Upsample latents without tiling.""" + latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents) + upsampled_latents = model(latents) upsampled_latents = vae.first_stage_model.per_channel_statistics.normalize( upsampled_latents diff --git a/comfy_extras/nodes_upscale_model.py b/comfy_extras/nodes_upscale_model.py index 1cf5a5d0180..a4d692955fa 100644 --- a/comfy_extras/nodes_upscale_model.py +++ b/comfy_extras/nodes_upscale_model.py @@ -7,6 +7,7 @@ from typing_extensions import override from comfy_api.latest import ComfyExtension, io import comfy.model_management +import comfy.model_patcher try: from spandrel_extra_arches import EXTRA_REGISTRY @@ -42,6 +43,7 @@ def execute(cls, model_name) -> io.NodeOutput: if not isinstance(out, ImageModelDescriptor): raise Exception("Upscale model must be a single-image model.") + out.patcher = comfy.model_patcher.CoreModelPatcher(out.model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device()) return io.NodeOutput(out) load_model = execute # TODO: remove @@ -66,14 +68,12 @@ def define_schema(cls): @classmethod def execute(cls, upscale_model, image) -> io.NodeOutput: - device = model_management.get_torch_device() + device = upscale_model.patcher.load_device - memory_required = model_management.module_size(upscale_model.model) - memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate + memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate memory_required += image.nelement() * image.element_size() - model_management.free_memory(memory_required, device) + model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required) - upscale_model.to(device) in_img = image.movedim(-1,-3).to(device) tile = 512 @@ -82,20 +82,17 @@ def execute(cls, upscale_model, image) -> io.NodeOutput: output_device = comfy.model_management.intermediate_device() oom = True - try: - while oom: - try: - steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) - pbar = comfy.utils.ProgressBar(steps) - s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a.float()), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=output_device) - oom = False - except Exception as e: - model_management.raise_non_oom(e) - tile //= 2 - if tile < 128: - raise e - finally: - upscale_model.to("cpu") + while oom: + try: + steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap) + pbar = comfy.utils.ProgressBar(steps) + s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a.float()), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar, output_device=output_device) + oom = False + except Exception as e: + model_management.raise_non_oom(e) + tile //= 2 + if tile < 128: + raise e s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0).to(comfy.model_management.intermediate_dtype()) return io.NodeOutput(s)