up

Author: yiyixuxu

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -776,17 +776,16 @@ def __init__(
         attn_scales: List[float] =[],
         temperal_downsample: List[bool] =[False, True, True],
         dropout: float =0.0,
+        latents_mean: List[float] = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 
+        0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921,] ,
+        latents_std: List[float] = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 
+        3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160,] ,
     ) -> None:
         super().__init__()
 
-        # channel-wise mean and std
-        mean = [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 
-        0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921,] 
-        std = [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 
-        3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160,] 
         # Store normalization parameters as tensors
-        self.mean = torch.tensor(mean)
-        self.std = torch.tensor(std)
+        self.mean = torch.tensor(latents_mean)
+        self.std = torch.tensor(latents_std)
         self.scale = torch.stack([self.mean, 1.0 / self.std])  # Shape: [2, C]
 
         self.z_dim = z_dim