Update app.py

app.py

@@ -2,16 +2,11 @@ import gradio as gr
 import os
 import torch
 import tempfile
+import sys
 from huggingface_hub import snapshot_download
-import sys, os
-sys.path.insert(0, os.path.abspath("./PusaV1"))
-
-
 import spaces
 
-import sys, os
-
-# Add PusaV1 to sys.path if not already
+# Setup paths
 PUSA_PATH = os.path.abspath("./PusaV1")
 if PUSA_PATH not in sys.path:
     sys.path.insert(0, PUSA_PATH)
@@ -24,15 +19,9 @@ if not os.path.exists(DIFFSYNTH_PATH):
         f"Ensure PusaV1 is correctly cloned and folder structure is intact."
     )
 
-
-
-
-
-
-
-
+# Import core modules from PusaV1
 from PusaV1.diffsynth import ModelManager, WanVideoPusaPipeline, save_video
-    
+
 # Constants
 WAN_SUBFOLDER = "Wan2.1-T2V-14B"
 MODEL_REPO_ID = "RaphaelLiu/PusaV1"
@@ -53,19 +42,42 @@ def ensure_model_downloaded():
         )
         print("Model downloaded.")
 
+# Subclass ModelManager to force WanModelPusa
+class PatchedModelManager(ModelManager):
+    def load_model(self, file_path=None, model_names=None, device=None, torch_dtype=None):
+        if file_path is None:
+            file_path = self.file_path_list[0]
+        print(f"[app.py] Forcing architecture: WanModelPusa for {file_path}")
+        for detector in self.model_detector:
+            if detector.match(file_path, {}):
+                model_names, models = detector.load(
+                    file_path,
+                    state_dict={},
+                    device=device or self.device,
+                    torch_dtype=torch_dtype or self.torch_dtype,
+                    allowed_model_names=model_names,
+                    model_manager=self,
+                    forced_architecture="WanModelPusa"
+                )
+                for name, model in zip(model_names, models):
+                    self.model.append(model)
+                    self.model_path.append(file_path)
+                    self.model_name.append(name)
+                return models[0] if models else None
+        print("No suitable model detector matched.")
+        return None
+
 # Video generation logic
-@spaces.GPU(duration = 200)
+@spaces.GPU(duration=200)
 def generate_video(prompt: str):
     ensure_model_downloaded()
 
-    # Load model
-    manager = ModelManager(
-      file_path_list=[WAN_MODEL_PATH],
-      torch_dtype=torch.float16,
-      device="cuda"
-     )
-
-   
+    # Load model using patched manager
+    manager = PatchedModelManager(
+        file_path_list=[WAN_MODEL_PATH],
+        torch_dtype=torch.float16,
+        device="cuda"
+    )
     model = manager.load_model()
 
     # Set up pipeline