refactor: remove hardcoded DiT LayerNorm paths

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

@@ -1,25 +1,15 @@
-import torch.nn.functional as F
-
-
 class ErnieImagePostInfer:
     def __init__(self, config):
         self.config = config
         self.patch_size = config.get("patch_size", 1)
         self.out_channels = config.get("out_channels", config.get("in_channels", 128))
-        self.eps = config.get("eps", 1e-6)
 
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
 
     def infer(self, weights, hidden_states, pre_infer_out):
         scale, shift = weights.final_norm_linear.apply(pre_infer_out.conditioning).chunk(2, dim=-1)
-        hidden_states = F.layer_norm(
-            hidden_states,
-            (hidden_states.shape[-1],),
-            weight=None,
-            bias=None,
-            eps=self.eps,
-        )
+        hidden_states = weights.final_norm.apply(hidden_states)
         hidden_states = hidden_states * (1 + scale) + shift
         patches = weights.final_linear.apply(hidden_states)[: pre_infer_out.image_tokens_len]
 

lightx2v/models/networks/ernie_image/weights/post_weights.py

@@ -1,11 +1,13 @@
 from lightx2v.common.modules.weight_module import WeightModule
-from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER
 
 
 class ErnieImagePostWeights(WeightModule):
     def __init__(self, config):
         super().__init__()
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
+        self.add_module("final_norm", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=config.get("eps", 1e-6)))
         self.add_module(
             "final_norm_linear",
             MM_WEIGHT_REGISTER[self.mm_type](

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

@@ -25,7 +25,7 @@ def infer(self, weights, hidden_states, temb):
         temb_out = weights.norm_out_linear.apply(F.silu(temb))
         scale, shift = torch.chunk(temb_out, 2, dim=-1)
 
-        hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        hidden_states = weights.norm_out.apply(hidden_states)
         hidden_states = hidden_states * (1 + scale) + shift
 
         output = weights.proj_out.apply(hidden_states)

lightx2v/models/networks/flux2/infer/transformer_infer.py

@@ -63,10 +63,10 @@ def infer_double_stream_block(
 
         (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = self._split_double_modulation(temb_mod_img)
         (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = self._split_double_modulation(temb_mod_txt)
-        norm_hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        norm_hidden_states = block_weights.norm1.apply(hidden_states)
         norm_hidden_states = (norm_hidden_states * (1 + scale_msa) + shift_msa).squeeze(0)
 
-        norm_encoder_hidden_states = F.layer_norm(encoder_hidden_states, (encoder_hidden_states.shape[-1],))
+        norm_encoder_hidden_states = block_weights.norm1_context.apply(encoder_hidden_states)
         norm_encoder_hidden_states = (norm_encoder_hidden_states * (1 + c_scale_msa) + c_shift_msa).squeeze(0)
 
         img_query = block_weights.to_q.apply(norm_hidden_states)
@@ -140,15 +140,15 @@ def infer_double_stream_block(
             img_attn_hook(gated_img_attn)
         hidden_states = hidden_states + gated_img_attn
         encoder_hidden_states = encoder_hidden_states + c_gate_msa * txt_attn_output
-        norm_hidden_states2 = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        norm_hidden_states2 = block_weights.norm2.apply(hidden_states)
         norm_hidden_states2 = (norm_hidden_states2 * (1 + scale_mlp) + shift_mlp).squeeze(0)
         ff_output = block_weights.ff_net_0.apply(norm_hidden_states2)
         ff_1, ff_2 = ff_output.chunk(2, dim=-1)
         ff_output = F.silu(ff_1) * ff_2
         ff_output = block_weights.ff_net_2.apply(ff_output)
         hidden_states = hidden_states + gate_mlp * ff_output
 
-        norm_encoder_hidden_states2 = F.layer_norm(encoder_hidden_states, (encoder_hidden_states.shape[-1],))
+        norm_encoder_hidden_states2 = block_weights.norm2_context.apply(encoder_hidden_states)
         norm_encoder_hidden_states2 = (norm_encoder_hidden_states2 * (1 + c_scale_mlp) + c_shift_mlp).squeeze(0)
         context_ff_output = block_weights.ff_context_net_0.apply(norm_encoder_hidden_states2)
         ctx_ff_1, ctx_ff_2 = context_ff_output.chunk(2, dim=-1)
@@ -179,7 +179,7 @@ def infer_single_stream_block(
 
         shift_msa, scale_msa, gate_msa = self._split_single_modulation(temb_mod)
 
-        norm_combined = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        norm_combined = block_weights.norm.apply(hidden_states)
         norm_combined = (norm_combined * (1 + scale_msa) + shift_msa).squeeze(0)
 
         hidden_states_proj = block_weights.to_qkv_mlp_proj.apply(norm_combined)

lightx2v/models/networks/flux2/weights/post_weights.py

@@ -1,5 +1,5 @@
 from lightx2v.common.modules.weight_module import WeightModule
-from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER
 
 
 class Flux2PostWeights(WeightModule):
@@ -12,6 +12,9 @@ def __init__(self, config):
         self.out_channels = config.get("transformer_in_channels", config.get("in_channels", 64))
         self.patch_size = config.get("patch_size", 1)
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
+
+        self.add_module("norm_out", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
 
         self.add_module(
             "norm_out_linear",

lightx2v/models/networks/flux2/weights/transformer_weights.py

@@ -1,5 +1,5 @@
 from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList
-from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER
 
 
 class Flux2DoubleBlockWeights(WeightModule):
@@ -11,11 +11,17 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe
         self.block_idx = block_idx
         self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"]
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
         self.rms_norm_type = config.get("rms_norm_type", "torch")
         self.attn_type = config.get("attn_type", "flash_attn3")
 
         p = f"transformer_blocks.{self.block_idx}"
 
+        self.add_module("norm1", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+        self.add_module("norm1_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+        self.add_module("norm2", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+        self.add_module("norm2_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+
         self.add_module(
             "to_q",
             MM_WEIGHT_REGISTER[self.mm_type](
@@ -189,11 +195,14 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe
         self.block_idx = block_idx
         self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"]
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
         self.rms_norm_type = config.get("rms_norm_type", "torch")
         self.attn_type = config.get("attn_type", "flash_attn3")
 
         p = f"single_transformer_blocks.{self.block_idx}"
 
+        self.add_module("norm", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+
         self.add_module(
             "to_qkv_mlp_proj",
             MM_WEIGHT_REGISTER[self.mm_type](

lightx2v/models/networks/hunyuan_video/infer/feature_caching/transformer_infer.py

@@ -151,7 +151,7 @@ def calculate_should_calc(self, img, vec, block):
 
         img_mod1_shift, img_mod1_scale, _, _, _, _ = img_mod_layer.apply(vec_).chunk(6, dim=-1)
         inp = inp.squeeze(0)
-        normed_inp = torch.nn.functional.layer_norm(inp, (inp.shape[1],), None, None, 1e-6)
+        normed_inp = block.img_branch.img_norm1.apply(inp)
         modulated_inp = normed_inp * (1 + img_mod1_scale) + img_mod1_shift
 
         del normed_inp, inp, vec_

lightx2v/models/networks/hunyuan_video/weights/pre_weights.py

@@ -120,7 +120,7 @@ def __init__(self, block_idx, mm_type, config, block_prefix):
         self.mm_type = mm_type
         self.add_module(
             "norm1",
-            LN_WEIGHT_REGISTER["torch"](f"{block_prefix}.{block_idx}.norm1.weight", f"{block_prefix}.{block_idx}.norm1.bias"),
+            LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](f"{block_prefix}.{block_idx}.norm1.weight", f"{block_prefix}.{block_idx}.norm1.bias"),
         )
         self.add_module(
             "self_attn_qkv",
@@ -134,7 +134,7 @@ def __init__(self, block_idx, mm_type, config, block_prefix):
         self.add_module("self_attention", attention_weights_cls())
         self.add_module(
             "norm2",
-            LN_WEIGHT_REGISTER["torch"](f"{block_prefix}.{block_idx}.norm2.weight", f"{block_prefix}.{block_idx}.norm2.bias"),
+            LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](f"{block_prefix}.{block_idx}.norm2.weight", f"{block_prefix}.{block_idx}.norm2.bias"),
         )
         self.add_module(
             "mlp_fc1",

lightx2v/models/networks/lingbot_video/weights/post_weights.py

@@ -5,6 +5,6 @@
 class LingBotVideoPostWeights(WeightModule):
     def __init__(self, config):
         super().__init__()
-        self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"](eps=config.get("norm_eps", 1e-6)))
+        self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=config.get("norm_eps", 1e-6)))
         self.add_module("norm_out_modulation", MM_WEIGHT_REGISTER["Default-ForceFp32"]("norm_out_modulation.1.weight", "norm_out_modulation.1.bias"))
         self.add_module("proj_out", MM_WEIGHT_REGISTER["Default"]("proj_out.weight", "proj_out.bias"))

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

@@ -30,7 +30,7 @@ def infer(self, weights, hidden_states, temb):
         scale, shift = torch.chunk(temb_out, 2, dim=-1)
 
         # Layer norm (no learnable params) + modulation
-        hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        hidden_states = weights.norm_out.apply(hidden_states)
         hidden_states = hidden_states * (1 + scale) + shift
 
         # Final projection

lightx2v/models/networks/longcat_image/infer/transformer_infer.py

@@ -40,7 +40,7 @@ def __init__(self, config):
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
 
-    def _ada_layer_norm_zero(self, hidden_states, temb, norm_linear):
+    def _ada_layer_norm_zero(self, hidden_states, temb, norm, norm_linear):
         """AdaLayerNormZero: compute shift, scale, gate from temb.
 
         For double-stream blocks: returns (norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp)
@@ -51,12 +51,12 @@ def _ada_layer_norm_zero(self, hidden_states, temb, norm_linear):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1)
 
         # Apply layer norm and modulation
-        norm_hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        norm_hidden_states = norm.apply(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
 
         return norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
-    def _ada_layer_norm_zero_single(self, hidden_states, temb, norm_linear):
+    def _ada_layer_norm_zero_single(self, hidden_states, temb, norm, norm_linear):
         """AdaLayerNormZeroSingle: for single-stream blocks.
 
         Returns (norm_hidden_states, gate)
@@ -67,7 +67,7 @@ def _ada_layer_norm_zero_single(self, hidden_states, temb, norm_linear):
         shift, scale, gate = emb.chunk(3, dim=-1)
 
         # Apply layer norm and modulation
-        norm_hidden_states = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        norm_hidden_states = norm.apply(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale) + shift
 
         return norm_hidden_states, gate
@@ -100,10 +100,15 @@ def infer_double_stream_block(
         head_dim = self.config["attention_head_dim"]
 
         # ===== Image stream: norm1 =====
-        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._ada_layer_norm_zero(hidden_states, temb, block_weights.norm1_linear)
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._ada_layer_norm_zero(hidden_states, temb, block_weights.norm1, block_weights.norm1_linear)
 
         # ===== Text stream: norm1_context =====
-        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self._ada_layer_norm_zero(encoder_hidden_states, temb, block_weights.norm1_context_linear)
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self._ada_layer_norm_zero(
+            encoder_hidden_states,
+            temb,
+            block_weights.norm1_context,
+            block_weights.norm1_context_linear,
+        )
 
         # ===== Attention projections =====
         # Image stream QKV
@@ -186,7 +191,7 @@ def infer_double_stream_block(
 
         # ===== FFN for image stream =====
         # Layer norm without learnable parameters (LongCat/Flux architecture)
-        norm_hidden_states2 = F.layer_norm(hidden_states, (hidden_states.shape[-1],))
+        norm_hidden_states2 = block_weights.norm2.apply(hidden_states)
         norm_hidden_states2 = norm_hidden_states2 * (1 + scale_mlp) + shift_mlp
         ff_output = block_weights.ff_net_0_proj.apply(norm_hidden_states2)
         ff_output = F.gelu(ff_output, approximate="tanh")
@@ -195,7 +200,7 @@ def infer_double_stream_block(
 
         # ===== FFN for text stream =====
         # Layer norm without learnable parameters (LongCat/Flux architecture)
-        norm_encoder_hidden_states2 = F.layer_norm(encoder_hidden_states, (encoder_hidden_states.shape[-1],))
+        norm_encoder_hidden_states2 = block_weights.norm2_context.apply(encoder_hidden_states)
         norm_encoder_hidden_states2 = norm_encoder_hidden_states2 * (1 + c_scale_mlp) + c_shift_mlp
         context_ff_output = block_weights.ff_context_net_0_proj.apply(norm_encoder_hidden_states2)
         context_ff_output = F.gelu(context_ff_output, approximate="tanh")
@@ -242,7 +247,7 @@ def infer_single_stream_block(
         residual = combined
 
         # AdaLayerNormZeroSingle
-        norm_combined, gate = self._ada_layer_norm_zero_single(combined, temb, block_weights.norm_linear)
+        norm_combined, gate = self._ada_layer_norm_zero_single(combined, temb, block_weights.norm, block_weights.norm_linear)
 
         # MLP branch
         mlp_hidden_states = block_weights.proj_mlp.apply(norm_combined)

lightx2v/models/networks/longcat_image/weights/post_weights.py

@@ -1,5 +1,5 @@
 from lightx2v.common.modules.weight_module import WeightModule
-from lightx2v.utils.registry_factory import MM_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER
 
 
 class LongCatImagePostWeights(WeightModule):
@@ -13,6 +13,9 @@ def __init__(self, config):
         self.out_channels = config.get("transformer_in_channels", config.get("in_channels", 64))
         self.patch_size = config.get("patch_size", 1)
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
+
+        self.add_module("norm_out", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
 
         # norm_out (AdaLayerNormContinuous)
         self.add_module(

lightx2v/models/networks/longcat_image/weights/transformer_weights.py

@@ -1,5 +1,5 @@
 from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList
-from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, LN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER
 
 
 class LongCatImageDoubleBlockWeights(WeightModule):
@@ -11,11 +11,17 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe
         self.block_idx = block_idx
         self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"]
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
         self.rms_norm_type = config.get("rms_norm_type", "torch")
         self.attn_type = config.get("attn_type", "flash_attn3")
 
         p = f"transformer_blocks.{self.block_idx}"
 
+        self.add_module("norm1", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+        self.add_module("norm1_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+        self.add_module("norm2", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+        self.add_module("norm2_context", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+
         # Image stream norm1 (AdaLayerNormZero)
         self.add_module(
             "norm1_linear",
@@ -216,11 +222,14 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe
         self.block_idx = block_idx
         self.inner_dim = config["num_attention_heads"] * config["attention_head_dim"]
         self.mm_type = config.get("dit_quant_scheme", "Default")
+        self.layer_norm_type = config.get("layer_norm_type", "torch")
         self.rms_norm_type = config.get("rms_norm_type", "torch")
         self.attn_type = config.get("attn_type", "flash_attn3")
 
         p = f"single_transformer_blocks.{self.block_idx}"
 
+        self.add_module("norm", LN_WEIGHT_REGISTER[self.layer_norm_type](eps=1e-5))
+
         # AdaLayerNormZeroSingle
         self.add_module(
             "norm_linear",

lightx2v/models/networks/ltx2/weights/post_weights.py

@@ -45,7 +45,7 @@ def __init__(self, config):
                 "model.diffusion_model.scale_shift_table",
             ),
         )
-        self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"]())
+        self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]())
         self.add_module(
             "proj_out",
             MM_WEIGHT_REGISTER["Default"](
@@ -60,7 +60,7 @@ def __init__(self, config):
                 "model.diffusion_model.audio_scale_shift_table",
             ),
         )
-        self.add_module("audio_norm_out", LN_WEIGHT_REGISTER["torch"]())
+        self.add_module("audio_norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")]())
         self.add_module(
             "audio_proj_out",
             MM_WEIGHT_REGISTER["Default"](

lightx2v/models/networks/qwen_image/weights/post_weights.py

@@ -18,7 +18,7 @@ def __init__(self, config):
                 "norm_out.linear.bias",
             ),
         )
-        self.add_module("norm_out", LN_WEIGHT_REGISTER["torch"](eps=1e-6))
+        self.add_module("norm_out", LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](eps=1e-6))
 
         # proj_out
         self.add_module(

lightx2v/models/networks/wan/infer/causvid/transformer_infer.py

@@ -96,7 +96,7 @@ def _infer_without_offload(self, weights, grid_sizes, embed, x, embed0, seq_lens
         return x
 
     def infer_self_attn(self, weights, grid_sizes, embed, x, embed0, seq_lens, freqs, context, block_idx, kv_start, kv_end):
-        norm1_out = torch.nn.functional.layer_norm(x, (x.shape[1],), None, None, 1e-6)
+        norm1_out = weights.norm1.apply(x)
         norm1_out = (norm1_out * (1 + embed0[1]) + embed0[0]).squeeze(0)
 
         s, n, d = *norm1_out.shape[:1], self.num_heads, self.head_dim
@@ -197,7 +197,7 @@ def infer_cross_attn(self, weights, x, context, block_idx):
         return x
 
     def infer_ffn(self, weights, x, embed0):
-        norm2_out = torch.nn.functional.layer_norm(x, (x.shape[1],), None, None, 1e-6)
+        norm2_out = weights.norm2.apply(x)
         y = weights.ffn_0.apply(norm2_out * (1 + embed0[4].squeeze(0)) + embed0[3].squeeze(0))
         y = torch.nn.functional.gelu(y, approximate="tanh")
         y = weights.ffn_2.apply(y)

lightx2v/models/networks/wan/infer/matrix_game2/transformer_infer.py

@@ -168,7 +168,7 @@ def _infer_action_model_with_kvcache(
             group_mouse = torch.nn.functional.linear(group_mouse, phase.mouse_mlp_0.weight.T, phase.mouse_mlp_0.bias)
             group_mouse = torch.nn.functional.gelu(group_mouse, approximate="tanh")
             group_mouse = torch.nn.functional.linear(group_mouse, phase.mouse_mlp_2.weight.T, phase.mouse_mlp_2.bias)
-            group_mouse = torch.nn.functional.layer_norm(group_mouse, (group_mouse.shape[-1],), phase.mouse_mlp_3.weight.T, phase.mouse_mlp_3.bias, 1e-5)
+            group_mouse = phase.mouse_mlp_3.apply(group_mouse)
 
             mouse_qkv = torch.nn.functional.linear(group_mouse, phase.t_qkv.weight.T)