Merge pull request #77 from ModelTC/dev_fix

Fix.

lightx2v/common/ops/norm/layer_norm_weight.py

@@ -94,9 +94,7 @@ class LNWeight(LNWeightTemplate):
             self.bias = None
 
     def apply(self, input_tensor):
-        if self.weight is None or self.weight.dtype == torch.bfloat16:
-            input_tensor = torch.nn.functional.layer_norm(input_tensor, (input_tensor.shape[-1],), self.weight, self.bias, self.eps)
-        else:
+        if GET_DTYPE() != "BF16":
             input_tensor = torch.nn.functional.layer_norm(
                 input_tensor.float(),
                 (input_tensor.shape[-1],),
@@ -104,4 +102,6 @@ class LNWeight(LNWeightTemplate):
                 self.bias,
                 self.eps,
             ).to(torch.bfloat16)
+        else:
+            input_tensor = torch.nn.functional.layer_norm(input_tensor, (input_tensor.shape[-1],), self.weight, self.bias, self.eps)
         return input_tensor
\ No newline at end of file

lightx2v/models/networks/wan/model.py

@@ -143,9 +143,8 @@ class WanModel:
 
     def _init_weights(self, weight_dict=None):
         use_bf16 = GET_DTYPE() == "BF16"
-
         # Some layers run with float32 to achieve high accuracy
-        skip_bf16 = {"norm", "embedding", "modulation", "time"}
+        skip_bf16 = {"norm", "embedding", "modulation", "time", "img_emb.proj.0", "img_emb.proj.4"}
         if weight_dict is None:
             if not self.dit_quantized or self.weight_auto_quant:
                 self.original_weight_dict = self._load_ckpt(use_bf16, skip_bf16)