Update app.py

app.py

@@ -20,7 +20,24 @@ if not os.path.exists(DIFFSYNTH_PATH):
     )
 
 # Import core modules from PusaV1
-from PusaV1.diffsynth import ModelManager, WanVideoPusaPipeline, save_video
+from PusaV1.diffsynth import WanVideoPusaPipeline, save_video
+
+
+from PusaV1.diffsynth import ModelManager as BaseModelManager
+
+class PatchedModelManager(BaseModelManager):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Your custom architecture dict entries to patch or add
+        custom_architecture_dict = {
+            "WanModel": ("diffsynth.models.wan_model", "WanModelPusa", None),
+            # Add more fixes if needed
+        }
+
+        # Update or replace the architecture_dict
+        self.architecture_dict.update(custom_architecture_dict)
+
 
 # Constants
 WAN_SUBFOLDER = "Wan2.1-T2V-14B"
@@ -43,6 +60,8 @@ 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:
@@ -73,12 +92,19 @@ def generate_video(prompt: str):
     ensure_model_downloaded()
 
     # Load model using patched manager
+    # manager = PatchedModelManager(
+    #     file_path_list=[WAN_MODEL_PATH],
+    #     torch_dtype=torch.float16,
+    #     device="cuda"
+    # )
+
     manager = PatchedModelManager(
-        file_path_list=[WAN_MODEL_PATH],
-        torch_dtype=torch.float16,
-        device="cuda"
-    )
-    model = manager.load_model()
+       file_path_list=[WAN_MODEL_PATH],
+       torch_dtype=torch.float16,
+       device="cuda" 
+     )
+   
+    model =  manager.load_model(WAN_MODEL_PATH)
 
     # Set up pipeline
     pipeline = WanVideoPusaPipeline(model=model)