refactor: remove hardcoded DiT LayerNorm paths

[Watebear]

No description provided.

[gemini-code-assist]

Code Review

This pull request refactors LayerNorm operations across multiple models (including ErnieImage, Flux2, HunyuanVideo, LongCatImage, Wan, and WorldMirror) to route through registered weight modules instead of calling F.layer_norm directly. It also updates the LayerNorm registry selection to dynamically use the configured layer_norm_type rather than hardcoding "torch", and adds robust handling for LayerNorms without learnable parameters. A new test suite is introduced to enforce these routing rules and lifecycle behaviors. A critical review comment points out that replacing direct F.layer_norm calls in WorldMirror removes dtype-matching logic, which could lead to runtime errors under torch.amp.autocast when mixing float32 weights with float16/bfloat16 activations.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Replacing the direct F.layer_norm call with self.ln_weight.apply(x) removes the critical dtype-matching logic that was previously present in _LNAdapter.forward.

Under torch.amp.autocast, the activation tensor x is automatically cast to bfloat16 or float16, while the weights stored in self.ln_weight remain in float32. Calling F.layer_norm directly with mixed dtypes (e.g., bfloat16 input and float32 weight) will raise a RuntimeError: expected scalar type BFloat16 but found Float.

To resolve this, the default LNWeight.apply implementation in lightx2v/common/ops/norm/layer_norm_weight.py should be updated to dynamically cast the weight and bias to match the input tensor's dtype, similar to how PyTorch's nn.LayerNorm handles autocast:

def apply(self, input_tensor):
    w = self._get_actual_weight()
    b = self._get_actual_bias()
    if self.sensitive_layer_dtype != self.infer_dtype:
        output_tensor = torch.nn.functional.layer_norm(
            input_tensor.float(),
            (input_tensor.shape[-1],),
            w,
            b,
            self.eps,
        ).to(self.infer_dtype)
    else:
        if w is not None and w.dtype != input_tensor.dtype:
            w = w.to(input_tensor.dtype)
        if b is not None and b.dtype != input_tensor.dtype:
            b = b.to(input_tensor.dtype)
        output_tensor = torch.nn.functional.layer_norm(
            input_tensor,
            (input_tensor.shape[-1],),
            w,
            b,
            self.eps,
        )
    return output_tensor