fix(ltx2): support complex RoPE and configurable normalization

lightx2v/common/ops/norm/rms_norm_weight.py

@@ -92,16 +92,25 @@ def __init__(
 
     def _get_base_attrs_mapping(self):
         self.base_attrs = []
-        self.base_attrs.append((self.weight_name, "weight", False))
+        if self.weight_name is not None:
+            self.base_attrs.append((self.weight_name, "weight", False))
+        else:
+            self.weight = None
 
     def _get_lora_attr_mapping(self):
-        _, _, _, self.weight_diff_name, _ = build_lora_and_diff_names(self.weight_name, self.lora_prefix)
-        self.lora_attrs = {
-            "weight_diff": "weight_diff_name",
-        }
-        self.weight_diff = torch.tensor(0.0, dtype=GET_DTYPE(), device=AI_DEVICE)
+        if self.weight_name is not None:
+            _, _, _, self.weight_diff_name, _ = build_lora_and_diff_names(self.weight_name, self.lora_prefix)
+            self.lora_attrs = {
+                "weight_diff": "weight_diff_name",
+            }
+            self.weight_diff = torch.tensor(0.0, dtype=GET_DTYPE(), device=AI_DEVICE)
+        else:
+            self.weight_diff_name = None
+            self.lora_attrs = {}
 
     def _get_actual_weight(self):
+        if self.weight is None:
+            return None
         if not hasattr(self, "weight_diff"):
             return self.weight
         if self.weight_diff.device != self.weight.device or self.weight_diff.dtype != self.weight.dtype:
@@ -110,7 +119,7 @@ def _get_actual_weight(self):
 
     def register_diff(self, weight_dict):
         if not self.lazy_load or self.create_cuda_buffer or self.create_cpu_buffer:
-            if self.weight_diff_name in weight_dict:
+            if self.weight_diff_name is not None and self.weight_diff_name in weight_dict:
                 self.weight_diff = weight_dict[self.weight_diff_name]
                 logger.debug(f"Register Diff to {self.weight_name}")
 
@@ -171,14 +180,17 @@ def load_lora_state_dict_from_disk(self, block_index):
                     )
 
     def load_state_dict_from_disk(self, block_index, adapter_block_index=None):
-        if self.has_lora_branch or self.has_diff:
-            self.load_lora_state_dict_from_disk(block_index)
-        self.weight_name = resolve_block_name(self.weight_name, block_index, adapter_block_index, self.is_post_adapter)
-        lazy_load_file_path = get_lazy_load_file_path(self.lazy_load_file, self.weight_name)
-        with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file:
-            weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype)
-            self.pin_weight = self.pin_weight.copy_(weight_tensor)
-        del weight_tensor
+        if self.weight_name is not None:
+            if self.has_lora_branch or self.has_diff:
+                self.load_lora_state_dict_from_disk(block_index)
+            self.weight_name = resolve_block_name(self.weight_name, block_index, adapter_block_index, self.is_post_adapter)
+            lazy_load_file_path = get_lazy_load_file_path(self.lazy_load_file, self.weight_name)
+            with safe_open(lazy_load_file_path, framework="pt", device="cpu") as lazy_load_file:
+                weight_tensor = lazy_load_file.get_tensor(self.weight_name).to(self.infer_dtype)
+                self.pin_weight = self.pin_weight.copy_(weight_tensor)
+            del weight_tensor
+        else:
+            self.weight = None
 
     @abstractmethod
     def apply(self, input_tensor):
@@ -216,10 +228,11 @@ def _norm(self, x):
 
     def apply(self, input_tensor):
         if GET_SENSITIVE_DTYPE() != GET_DTYPE():
-            input_tensor = self._norm(input_tensor).type_as(input_tensor) * (self._get_actual_weight())
+            output = self._norm(input_tensor).type_as(input_tensor)
         else:
-            input_tensor = self._norm(input_tensor.float()).type_as(input_tensor) * (self._get_actual_weight())
-        return input_tensor
+            output = self._norm(input_tensor.float()).type_as(input_tensor)
+        weight = self._get_actual_weight()
+        return output if weight is None else output * weight
 
 
 @RMS_WEIGHT_REGISTER("torch_native")
@@ -284,12 +297,11 @@ def apply(self, input_tensor):
 
         # Apply normalization with global mean
         if self.sensitive_layer_dtype != self.infer_dtype:
-            input_tensor = input_tensor * torch.rsqrt(global_mean.float() + self.eps).to(self.infer_dtype)
-            input_tensor = (input_tensor * self._get_actual_weight()).to(self.infer_dtype)
+            output = input_tensor * torch.rsqrt(global_mean.float() + self.eps).to(self.infer_dtype)
         else:
-            input_tensor = input_tensor * torch.rsqrt(global_mean + self.eps)
-            input_tensor = input_tensor * self._get_actual_weight()
-        return input_tensor
+            output = input_tensor * torch.rsqrt(global_mean + self.eps)
+        weight = self._get_actual_weight()
+        return output if weight is None else (output * weight).to(self.infer_dtype)
 
 
 @RMS_WEIGHT_REGISTER("sgl-kernel")
@@ -320,11 +332,11 @@ def __init__(
         self.enable_pdl = is_arch_support_pdl() if is_arch_support_pdl is not None else False
 
     def apply(self, input_tensor):
-        if sgl_kernel is not None and self.sensitive_layer_dtype == self.infer_dtype:
+        weight = self._get_actual_weight()
+        if weight is not None and sgl_kernel is not None and self.sensitive_layer_dtype == self.infer_dtype:
             input_tensor = input_tensor.contiguous()
             orig_shape = input_tensor.shape
             input_tensor = input_tensor.view(-1, orig_shape[-1])
-            weight = self._get_actual_weight()
             if torch.compiler.is_compiling() and flashinfer_rmsnorm is not None and input_tensor.dtype in (torch.float16, torch.bfloat16):
                 input_tensor = rmsnorm_flashinfer(input_tensor, weight, self.eps, self.enable_pdl)
             else:
@@ -334,10 +346,10 @@ def apply(self, input_tensor):
             # sgl_kernel is not available or dtype!=torch.bfloat16/float16, fallback to default implementation
             if self.sensitive_layer_dtype != self.infer_dtype:
                 input_tensor = input_tensor * torch.rsqrt(input_tensor.float().pow(2).mean(-1, keepdim=True) + self.eps).to(self.infer_dtype)
-                input_tensor = (input_tensor * (self._get_actual_weight())).to(self.infer_dtype)
             else:
                 input_tensor = input_tensor * torch.rsqrt(input_tensor.pow(2).mean(-1, keepdim=True) + self.eps)
-                input_tensor = input_tensor * (self._get_actual_weight())
+            if weight is not None:
+                input_tensor = (input_tensor * weight).to(self.infer_dtype)
 
         return input_tensor
 
@@ -373,10 +385,11 @@ def apply(self, input_tensor):
         variance = input_tensor.to(torch.float32).pow(2).mean(-1, keepdim=True)
         hidden_states = input_tensor * torch.rsqrt(variance + self.eps)
 
-        if self.weight.dtype in [torch.float16, torch.bfloat16]:
-            hidden_states = hidden_states.to(self.weight.dtype)
-        if self.weight is not None:
-            hidden_states = hidden_states * (self._get_actual_weight())
+        weight = self._get_actual_weight()
+        if weight is not None:
+            if weight.dtype in [torch.float16, torch.bfloat16]:
+                hidden_states = hidden_states.to(weight.dtype)
+            hidden_states = hidden_states * weight
         hidden_states = hidden_states.to(input_dtype)
 
         return hidden_states
@@ -413,7 +426,8 @@ def apply(self, hidden_states):
         hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)
         hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
-        return self.weight * hidden_states.to(input_dtype)
+        hidden_states = hidden_states.to(input_dtype)
+        return hidden_states if self.weight is None else self.weight * hidden_states
 
 
 @RMS_WEIGHT_REGISTER("self_forcing")
@@ -446,7 +460,9 @@ def _norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
 
     def apply(self, x):
-        return self._norm(x.float()).type_as(x) * (self._get_actual_weight())
+        output = self._norm(x.float()).type_as(x)
+        weight = self._get_actual_weight()
+        return output if weight is None else output * weight
 
 
 @RMS_WEIGHT_REGISTER("one-pass")
@@ -477,6 +493,8 @@ def __init__(
 
     def apply(self, input_tensor):
         w = self._get_actual_weight()
+        if w is None:
+            return torch.nn.functional.rms_norm(input_tensor, (input_tensor.shape[-1],), eps=self.eps)
         if use_magi_custom_ops() and magi_register_custom_op is not None:
             return torch.ops.lightx2v.rms_norm(input_tensor, w, self.eps)
         return rms_norm_kernel(input_tensor, w, self.eps)

lightx2v/common/ops/rope/torch_rope.py

@@ -16,25 +16,75 @@
 @ROPE_REGISTER("torch_complex_rope")
 class TorchComplexRope(RopeTemplate):
     def __init__(self, layout="interleaved", compute_dtype: torch.dtype = torch.float32):
-        if layout != "interleaved":
-            raise ValueError("TorchComplexRope only supports interleaved layout.")
         super().__init__(layout=layout, compute_dtype=compute_dtype)
 
+    def prepare_freqs(self, freqs, rotary_dim: int | None = None):
+        if torch.is_tensor(freqs):
+            if torch.is_complex(freqs):
+                return freqs
+            if rotary_dim is None:
+                raise ValueError("rotary_dim is required for real RoPE frequencies.")
+            if freqs.shape[-1] != rotary_dim:
+                raise ValueError(f"Concatenated cos-sin cache must have last dim {rotary_dim}, got {freqs.shape[-1]}.")
+            cos, sin = freqs.chunk(2, dim=-1)
+        elif isinstance(freqs, tuple):
+            if len(freqs) != 2:
+                raise ValueError(f"Expected a (cos, sin) tuple, got {len(freqs)} entries.")
+            if rotary_dim is None:
+                raise ValueError("rotary_dim is required for tuple RoPE frequencies.")
+            cos, sin = freqs
+            if cos.shape != sin.shape:
+                raise ValueError(f"RoPE cos/sin shapes must match, got cos={cos.shape}, sin={sin.shape}.")
+            if cos.shape[-1] == rotary_dim:
+                if self.layout == "interleaved":
+                    cos, sin = cos[..., 0::2], sin[..., 0::2]
+                else:
+                    cos, sin = cos[..., : rotary_dim // 2], sin[..., : rotary_dim // 2]
+            elif cos.shape[-1] != rotary_dim // 2:
+                raise ValueError(f"RoPE frequency width must be {rotary_dim // 2} or {rotary_dim}, got {cos.shape[-1]}.")
+        else:
+            raise TypeError(f"Unsupported RoPE frequency type: {type(freqs)!r}")
+        return torch.complex(cos.to(self.compute_dtype), sin.to(self.compute_dtype)).contiguous()
+
     def apply(self, q: torch.Tensor, k: torch.Tensor, freqs, **kwargs):
-        if q.ndim == 3 and k.ndim == 3 and self.compute_dtype == torch.float32 and torch.is_complex(freqs) and use_magi_custom_ops() and magi_register_custom_op is not None and not kwargs:
+        rotary_dim = kwargs.get("rotary_dim", q.shape[-1])
+        freqs = self.prepare_freqs(freqs, rotary_dim=rotary_dim)
+        if (
+            q.ndim == 3
+            and k.ndim == 3
+            and self.compute_dtype == torch.float32
+            and self.layout == "interleaved"
+            and torch.is_tensor(freqs)
+            and torch.is_complex(freqs)
+            and use_magi_custom_ops()
+            and magi_register_custom_op is not None
+            and not kwargs
+        ):
             return torch.ops.lightx2v.rope_torch_complex(q, k, freqs)
         return super().apply(q, k, freqs, **kwargs)
 
-    def apply_single(self, x: torch.Tensor, freqs: torch.Tensor, rotary_dim: int | None = None, unsqueeze_dim: int = -2, **kwargs):
-        if not torch.is_complex(freqs):
-            raise TypeError("TorchComplexRope expects a complex frequency tensor.")
-        rotary_dim = rotary_dim or x.shape[-1]
+    def apply_single(self, x: torch.Tensor, freqs, rotary_dim: int | None = None, unsqueeze_dim: int = -2, **kwargs):
+        rotary_dim = x.shape[-1] if rotary_dim is None else rotary_dim
         if rotary_dim % 2:
             raise ValueError(f"rotary_dim must be even, got {rotary_dim}.")
+        freqs = self.prepare_freqs(freqs, rotary_dim=rotary_dim)
+        if not torch.is_tensor(freqs) or not torch.is_complex(freqs):
+            raise TypeError("TorchComplexRope expects a complex frequency tensor.")
         x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:]
-        x_complex = torch.view_as_complex(x_rot.to(self.compute_dtype).reshape(*x_rot.shape[:-1], -1, 2).contiguous())
+        x_float = x_rot.to(self.compute_dtype)
+        if self.layout == "interleaved":
+            x_pairs = x_float.reshape(*x_rot.shape[:-1], -1, 2)
+        else:
+            first, second = x_float.chunk(2, dim=-1)
+            x_pairs = torch.stack((first, second), dim=-1)
+        x_complex = torch.view_as_complex(x_pairs.contiguous())
         freqs = broadcast_freqs(freqs, x_complex, unsqueeze_dim)
-        output = torch.view_as_real(x_complex * freqs).flatten(-2).to(x.dtype)
+        output_pairs = torch.view_as_real(x_complex * freqs)
+        if self.layout == "interleaved":
+            output = output_pairs.flatten(-2)
+        else:
+            output = torch.cat((output_pairs[..., 0], output_pairs[..., 1]), dim=-1)
+        output = output.to(x.dtype)
         return torch.cat((output, x_pass), dim=-1) if x_pass.shape[-1] else output
 
 

lightx2v/models/networks/ltx2/infer/post_infer.py

@@ -9,9 +9,6 @@
 
 import torch
 
-from lightx2v.models.networks.ltx2.infer.triton_ops import fused_rmsnorm_modulate
-from lightx2v.models.networks.ltx2.infer.utils import modulate_with_rmsnorm_torch_naive
-
 
 def to_denoised(
     sample: torch.Tensor,
@@ -45,10 +42,6 @@ def __init__(self, config):
         """
         self.config = config
         self.clean_cuda_cache = config.get("clean_cuda_cache", False)
-        if config.get("modulate_with_rmsnorm", "triton") == "triton":
-            self.modulate_with_rmsnorm_func = fused_rmsnorm_modulate
-        else:
-            self.modulate_with_rmsnorm_func = modulate_with_rmsnorm_torch_naive
 
     def set_scheduler(self, scheduler):
         """Set the scheduler for inference."""
@@ -76,13 +69,15 @@ def infer(
         """
         vx = self._process_output(
             weights.scale_shift_table.tensor,
+            weights.norm_out,
             weights.proj_out,
             vx,
             video_embedded_timestep,
         )
 
         ax = self._process_output(
             weights.audio_scale_shift_table.tensor,
+            weights.audio_norm_out,
             weights.audio_proj_out,
             ax,
             audio_embedded_timestep,
@@ -103,6 +98,7 @@ def infer(
     def _process_output(
         self,
         scale_shift_table: torch.Tensor,
+        norm_out,
         proj_out,
         x: torch.Tensor,
         embedded_timestep: torch.Tensor,
@@ -112,6 +108,7 @@ def _process_output(
 
         Args:
             scale_shift_table: Scale-shift table, shape [2, hidden_dim]
+            norm_out: Registered output LayerNorm module
             proj_out: Output projection layer
             x: Input tensor, shape [seq_len, hidden_dim]
             embedded_timestep: Embedded timestep, shape [seq_len, hidden_dim]
@@ -125,7 +122,7 @@ def _process_output(
         # Result shape: [seq_len, 2, hidden_dim]
         scale_shift_values = scale_shift_table[None, :, :].to(device=x.device, dtype=x.dtype) + embedded_timestep[:, None, :]
         shift, scale = scale_shift_values[:, 0], scale_shift_values[:, 1]
-        x = torch.nn.functional.layer_norm(x, (x.shape[-1],), eps=1e-6)
+        x = norm_out.apply(x)
         x = x * (1 + scale) + shift
         x = proj_out.apply(x)