[pull] master from Comfy-Org:master

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,

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):

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."""

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)
 

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