From 0db42c65aa428447c6c9d8a28bd79929d4185e13 Mon Sep 17 00:00:00 2001 From: Super User Date: Wed, 5 Aug 2026 20:39:33 +0800 Subject: [PATCH 1/2] fix(rope): support LTX2 complex RoPE inputs --- lightx2v/common/ops/rope/torch_rope.py | 68 ++++++++++++++++++++++---- 1 file changed, 59 insertions(+), 9 deletions(-) diff --git a/lightx2v/common/ops/rope/torch_rope.py b/lightx2v/common/ops/rope/torch_rope.py index f17683c86..0d03b9a8f 100644 --- a/lightx2v/common/ops/rope/torch_rope.py +++ b/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 From 1e868073c7a4c967d03eff8759c86f7ac1be9a84 Mon Sep 17 00:00:00 2001 From: Super User Date: Wed, 5 Aug 2026 20:39:58 +0800 Subject: [PATCH 2/2] refactor(ltx2): register normalization with weights --- lightx2v/common/ops/norm/rms_norm_weight.py | 84 +++++++++++-------- .../models/networks/ltx2/infer/post_infer.py | 13 ++- .../networks/ltx2/infer/transformer_infer.py | 51 ++++++----- .../ltx2/weights/transformer_weights.py | 14 ++++ 4 files changed, 100 insertions(+), 62 deletions(-) diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 7e1877d25..71c35b2c2 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/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) diff --git a/lightx2v/models/networks/ltx2/infer/post_infer.py b/lightx2v/models/networks/ltx2/infer/post_infer.py index 9d2dd4706..a5cbbc16a 100755 --- a/lightx2v/models/networks/ltx2/infer/post_infer.py +++ b/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,6 +69,7 @@ def infer( """ vx = self._process_output( weights.scale_shift_table.tensor, + weights.norm_out, weights.proj_out, vx, video_embedded_timestep, @@ -83,6 +77,7 @@ def infer( 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) diff --git a/lightx2v/models/networks/ltx2/infer/transformer_infer.py b/lightx2v/models/networks/ltx2/infer/transformer_infer.py index 43bf9d331..5c725a0ed 100644 --- a/lightx2v/models/networks/ltx2/infer/transformer_infer.py +++ b/lightx2v/models/networks/ltx2/infer/transformer_infer.py @@ -12,13 +12,11 @@ import torch import torch.distributed as dist -import torch.nn.functional as F from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer from lightx2v.models.networks.ltx2.infer.module_io import LTX2PreInferModuleOutput -from lightx2v.models.networks.ltx2.infer.triton_ops import fuse_scale_shift_kernel, fused_rmsnorm_modulate -from lightx2v.models.networks.ltx2.infer.utils import modulate_torch_naive, modulate_with_rmsnorm_torch_naive, rmsnorm_torch_naive -from lightx2v.models.networks.wan.infer.triton_ops import norm_infer +from lightx2v.models.networks.ltx2.infer.triton_ops import fuse_scale_shift_kernel +from lightx2v.models.networks.ltx2.infer.utils import modulate_torch_naive class LTX2TransformerInfer(BaseTransformerInfer): @@ -65,14 +63,13 @@ def __init__(self, config): self.tp_rank = 0 self.tp_size = 1 - if config.get("norm_modulate_backend", "triton") == "triton": - self.norm_infer_func = norm_infer + modulate_type = config.get("modulate_type", config.get("norm_modulate_backend", "triton")) + if modulate_type == "triton": self.modulate_func = fuse_scale_shift_kernel - self.modulate_with_rmsnorm_func = fused_rmsnorm_modulate - else: - self.norm_infer_func = rmsnorm_torch_naive + elif modulate_type == "torch": self.modulate_func = modulate_torch_naive - self.modulate_with_rmsnorm_func = modulate_with_rmsnorm_torch_naive + else: + raise ValueError(f"Unsupported modulate_type={modulate_type!r}; expected 'torch' or 'triton'.") self.init_compile(config) self.reset_infer_states() self.reset_guidance_perturbation() @@ -177,11 +174,11 @@ def _apply_text_cross_attention_adaln( prompt_scale_shift_table: torch.Tensor, prompt_timestep: torch.Tensor, is_audio: bool, + norm, ) -> torch.Tensor: """Text cross-attention with per-block prompt AdaLN (ltx_core apply_cross_attention_adaln).""" q_shift, q_scale, q_gate = self._get_ada_values(scale_shift_table, timesteps, slice(6, 9)) - # ltx_core.utils.rms_norm -> F.rms_norm(last-dim); do not use Triton norm_infer here (bf16 drift vs reference) - norm_x = F.rms_norm(x.unsqueeze(0), (x.shape[-1],), eps=1e-6).squeeze(0) + norm_x = norm.apply(x.unsqueeze(0)).squeeze(0) attn_input = norm_x * (1 + q_scale) + q_shift d = prompt_scale_shift_table.shape[-1] @@ -212,6 +209,7 @@ def _apply_text_cross_attention_adaln( k_pe=None, is_audio=is_audio, need_gather_video_context=False, + norm=norm, ) return attn_out * q_gate @@ -225,6 +223,7 @@ def _infer_attn( is_audio=False, need_gather_video_context=False, # Only True for video-to-audio cross-attention bypass_attention: bool = False, + norm=None, ) -> torch.Tensor: """ Unified attention inference method supporting both TP and non-TP modes. @@ -248,7 +247,9 @@ def _infer_attn( if is_self_attn or self.apply_gated_attention: q_in = x else: - q_in = self.norm_infer_func(x, weight=None, bias=None, eps=1e-6) + if norm is None: + raise ValueError("A registered RMSNorm module is required for cross-attention.") + q_in = norm.apply(x) num_heads = self.v_num_heads if not is_audio else self.a_num_heads head_dim = self.v_head_dim if not is_audio else self.a_head_dim @@ -396,7 +397,7 @@ def infer_block( slice(0, 3), ) - norm_vx = self.modulate_with_rmsnorm_func(vx, vscale_msa, vshift_msa, weight=None, bias=None, eps=1e-6) + norm_vx = self.modulate_func(block.norm.apply(vx), vscale_msa, vshift_msa) # Video self-attention vx = ( vx @@ -406,6 +407,7 @@ def infer_block( pe=pre_infer_out.video_args.positional_embeddings, is_audio=False, bypass_attention=skip_video_self, + norm=block.norm, ) * vgate_msa ) @@ -420,13 +422,15 @@ def infer_block( block.prompt_scale_shift_table.tensor, pre_infer_out.video_args.prompt_timestep, is_audio=False, + norm=block.norm, ) else: vx = vx + self._infer_attn( attn_phase=block.compute_phases[1], - x=self.norm_infer_func(vx, weight=None, bias=None, eps=1e-6), + x=block.norm.apply(vx), context=pre_infer_out.video_args.context, is_audio=False, + norm=block.norm, ) del vshift_msa, vscale_msa, vgate_msa @@ -438,7 +442,7 @@ def infer_block( slice(0, 3), ) - norm_ax = self.modulate_with_rmsnorm_func(ax, ascale_msa, ashift_msa, weight=None, bias=None, eps=1e-6) + norm_ax = self.modulate_func(block.norm.apply(ax), ascale_msa, ashift_msa) # Audio self-attention ax = ( @@ -449,6 +453,7 @@ def infer_block( pe=pre_infer_out.audio_args.positional_embeddings, is_audio=True, bypass_attention=skip_audio_self, + norm=block.norm, ) * agate_msa ) @@ -463,20 +468,22 @@ def infer_block( block.audio_prompt_scale_shift_table.tensor, pre_infer_out.audio_args.prompt_timestep, is_audio=True, + norm=block.norm, ) else: ax = ax + self._infer_attn( attn_phase=block.compute_phases[3], - x=self.norm_infer_func(ax, weight=None, bias=None, eps=1e-6), + x=block.norm.apply(ax), context=pre_infer_out.audio_args.context, is_audio=True, + norm=block.norm, ) del ashift_msa, ascale_msa, agate_msa # Audio-video cross-attention - vx_norm3 = self.norm_infer_func(vx, weight=None, bias=None, eps=1e-6) - ax_norm3 = self.norm_infer_func(ax, weight=None, bias=None, eps=1e-6) + vx_norm3 = block.norm.apply(vx) + ax_norm3 = block.norm.apply(ax) # Get audio scale-shift values ( @@ -522,6 +529,7 @@ def infer_block( k_pe=pre_infer_out.audio_args.cross_positional_embeddings, is_audio=True, need_gather_video_context=False, # Audio is global, no gather needed + norm=block.norm, ) * gate_out_a2v ) @@ -543,6 +551,7 @@ def infer_block( k_pe=pre_infer_out.video_args.cross_positional_embeddings, is_audio=True, need_gather_video_context=not (self.tp_size > 1), # Need gather for SP, not for TP + norm=block.norm, ) * gate_out_v2a ) @@ -565,7 +574,7 @@ def infer_block( pre_infer_out.video_args.timesteps, slice(3, 6), ) - vx_scaled = self.modulate_with_rmsnorm_func(vx, vscale_mlp, vshift_mlp, weight=None, bias=None, eps=1e-6) + vx_scaled = self.modulate_func(block.norm.apply(vx), vscale_mlp, vshift_mlp) vx = vx + self._infer_ffn(block.compute_phases[6], vx_scaled) * vgate_mlp del vshift_mlp, vscale_mlp, vgate_mlp @@ -575,7 +584,7 @@ def infer_block( pre_infer_out.audio_args.timesteps, slice(3, 6), ) - ax_scaled = self.modulate_with_rmsnorm_func(ax, ascale_mlp, ashift_mlp, weight=None, bias=None, eps=1e-6) + ax_scaled = self.modulate_func(block.norm.apply(ax), ascale_mlp, ashift_mlp) ax = ax + self._infer_ffn(block.compute_phases[7], ax_scaled) * agate_mlp del ashift_mlp, ascale_mlp, agate_mlp diff --git a/lightx2v/models/networks/ltx2/weights/transformer_weights.py b/lightx2v/models/networks/ltx2/weights/transformer_weights.py index 3c2f1a92d..1f6b39903 100755 --- a/lightx2v/models/networks/ltx2/weights/transformer_weights.py +++ b/lightx2v/models/networks/ltx2/weights/transformer_weights.py @@ -91,6 +91,20 @@ def __init__( block_prefix = "transformer_blocks" model_prefix = "model.diffusion_model" + # LTX2 block pre-normalization is an affine-free RMSNorm. Register it + # with the weight tree (like Qwen Image's affine-free norms) so the + # implementation is selected consistently by rms_norm_type. + self.add_module( + "norm", + RMS_WEIGHT_REGISTER[config.get("rms_norm_type", "sgl-kernel")]( + weight_name=None, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, + ), + ) + # Video scale-shift table self.scale_shift_table = TENSOR_REGISTER["Default"]( tensor_name=f"{model_prefix}.{block_prefix}.{self.block_index}.scale_shift_table",