diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py
index 15575a1c5907..6df55d96d2dc 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py
@@ -31,6 +31,7 @@
 CACHE_T = 2
 
 
+# Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AvgDown3D with AvgDown3D->QwenImage21AvgDown3D
 class QwenImage21AvgDown3D(nn.Module):
     def __init__(
         self,
@@ -46,7 +47,6 @@ def __init__(
                 f"`in_channels` ({in_channels}) times the downsampling factor ({factor}) must be divisible by "
                 f"`out_channels` ({out_channels})."
             )
-
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.factor_t = factor_t
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
index de8a56edc20e..a66dca647439 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
@@ -40,14 +40,18 @@ def __init__(
         factor_s=1,
     ):
         super().__init__()
+        factor = factor_t * factor_s * factor_s
+        if in_channels * factor % out_channels != 0:
+            raise ValueError(
+                f"`in_channels` ({in_channels}) times the downsampling factor ({factor}) must be divisible by "
+                f"`out_channels` ({out_channels})."
+            )
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.factor_t = factor_t
         self.factor_s = factor_s
-        self.factor = self.factor_t * self.factor_s * self.factor_s
-
-        assert in_channels * self.factor % out_channels == 0
-        self.group_size = in_channels * self.factor // out_channels
+        self.factor = factor
+        self.group_size = in_channels * factor // out_channels
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t