Update utils/wan_wrapper.py

utils/wan_wrapper.py

@@ -22,12 +22,17 @@ class WanTextEncoder(torch.nn.Module):
             device=torch.device('cpu')
         ).eval().requires_grad_(False)
         self.text_encoder.load_state_dict(
-            torch.load("/repository/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
+            # I should have called the folder "Wan2.1-T2V-1.3B" instead of "wan2.1"
+            # torch.load("/repository/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",
+            torch.load("/repository/wan2.1/models_t5_umt5-xxl-enc-bf16.pth",
                        map_location='cpu', weights_only=False)
         )
 
         self.tokenizer = HuggingfaceTokenizer(
-            name="/repository/Wan2.1-T2V-1.3B/google/umt5-xxl/", seq_len=512, clean='whitespace')
+            # I should have called the folder "Wan2.1-T2V-1.3B" instead of "wan2.1"
+            #name="/repository/Wan2.1-T2V-1.3B/google/umt5-xxl/",
+            name="/repository/wan2.1/google/umt5-xxl/",
+            seq_len=512, clean='whitespace')
 
     @property
     def device(self):
@@ -66,7 +71,9 @@ class WanVAEWrapper(torch.nn.Module):
 
         # init model
         self.model = _video_vae(
-            pretrained_path="/repository/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",
+            # I should have called the folder "Wan2.1-T2V-1.3B" instead of "wan2.1"
+            #pretrained_path="/repository/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",
+            pretrained_path="/repository/wan2.1/Wan2.1_VAE.pth",
             z_dim=16,
         ).eval().requires_grad_(False)
 
@@ -115,7 +122,9 @@ class WanVAEWrapper(torch.nn.Module):
 class WanDiffusionWrapper(torch.nn.Module):
     def __init__(
             self,
-            model_name="Wan2.1-T2V-1.3B",
+            # I should have called the folder "Wan2.1-T2V-1.3B" instead of "wan2.1"
+            #model_name="Wan2.1-T2V-1.3B",
+            model_name="wan2.1",
             timestep_shift=8.0,
             is_causal=False,
             local_attn_size=-1,