Fix Flax/Torch namespace collision in wan_vae_test.py

Author: eltsai

src/maxdiffusion/tests/wan_vae_test.py

@@ -57,7 +57,7 @@
 flax.config.update("flax_always_shard_variable", False)
 
 
-class TorchWanRMS_norm(nn.Module):
+class TorchWanRMS_norm(torch.nn.Module):
   r"""
   A custom RMS normalization layer.
 
@@ -76,14 +76,14 @@ def __init__(self, dim: int, channel_first: bool = True, images: bool = True, bi
 
     self.channel_first = channel_first
     self.scale = dim**0.5
-    self.gamma = nn.Parameter(torch.ones(shape))
-    self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+    self.gamma = torch.nn.Parameter(torch.ones(shape))
+    self.bias = torch.nn.Parameter(torch.zeros(shape)) if bias else 0.0
 
   def forward(self, x):
     return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
 
 
-class TorchWanResample(nn.Module):
+class TorchWanResample(torch.nn.Module):
   r"""
   A custom resampling module for 2D and 3D data.
 
@@ -104,18 +104,18 @@ def __init__(self, dim: int, mode: str) -> None:
 
     # layers
     if mode == "upsample2d":
-      self.resample = nn.Sequential(
-          WanUpsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), nn.Conv2d(dim, dim // 2, 3, padding=1)
+      self.resample = torch.nn.Sequential(
+          torch.nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest"), torch.nn.Conv2d(dim, dim // 2, 3, padding=1)
       )
     elif mode == "upsample3d":
       raise Exception("downsample3d not supported")
 
     elif mode == "downsample2d":
-      self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+      self.resample = torch.nn.Sequential(torch.nn.ZeroPad2d((0, 1, 0, 1)), torch.nn.Conv2d(dim, dim, 3, stride=(2, 2)))
     elif mode == "downsample3d":
       raise Exception("downsample3d not supported")
     else:
-      self.resample = nn.Identity()
+      self.resample = torch.nn.Identity()
 
   def forward(self, x, feat_cache=None, feat_idx=[0]):
     b, c, t, h, w = x.size()
@@ -218,7 +218,7 @@ def test_zero_padded_conv(self):
     dim = 96
     kernel_size = 3
     stride = (2, 2)
-    resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, kernel_size, stride=stride))
+    resample = torch.nn.Sequential(torch.nn.ZeroPad2d((0, 1, 0, 1)), torch.nn.Conv2d(dim, dim, kernel_size, stride=stride))
     input_shape = (1, 96, 480, 720)
     input = torch.ones(input_shape)
     output_torch = resample(input)