Format

optimization.py

@@ -17,16 +17,11 @@ from huggingface_hub import hf_hub_download
 from optimization_utils import capture_component_call
 from optimization_utils import aoti_compile
 from optimization_utils import drain_module_parameters
-from optimization_utils import zerogpu_compiled_from_serializable_dict
 from optimization_utils import ZeroGPUCompiledModel
 
 
 P = ParamSpec('P')
 
-# Expose compiled models so app.py can offer them for download
-COMPILED_TRANSFORMER_1 = None
-COMPILED_TRANSFORMER_2 = None
-
 LATENT_FRAMES_DIM = torch.export.Dim('num_latent_frames', min=8, max=81)
 LATENT_PATCHED_HEIGHT_DIM = torch.export.Dim('latent_patched_height', min=30, max=52)
 LATENT_PATCHED_WIDTH_DIM = torch.export.Dim('latent_patched_width', min=30, max=52)
@@ -55,50 +50,35 @@ def _strtobool(v: str | None, default: bool = True) -> bool:
     return v.strip().lower() in ("1", "true", "yes", "y", "on")
 
 
-def _load_compiled_pt(path: str):
-    """
-    Load either:
-    - a serialized dict produced by to_serializable_dict()  (format zerogpu_aoti_v1), or
-    - an old-style pickled ZeroGPUCompiledModel.
-    """
-    obj = torch.load(path, map_location="cpu", weights_only=False)
-
-    # New format: dict payload
-    if isinstance(obj, dict) and obj.get("format") == "zerogpu_aoti_v1":
-        return zerogpu_compiled_from_serializable_dict(obj)
-
-    # Old format: direct object
-    if isinstance(obj, ZeroGPUCompiledModel):
-        return obj
-
-    raise ValueError(
-        f"Unsupported compiled transformer file format at {path}. "
-        f"Got type={type(obj)} keys={list(obj.keys()) if isinstance(obj, dict) else None}"
-    )
-
-
 def load_compiled_transformers_from_hub(
     repo_id: str,
     filename_1: str = "compiled_transformer_1.pt",
     filename_2: str = "compiled_transformer_2.pt",
+    device: str = "cuda",
 ):
     """
-    Charge les artefacts précompilés depuis le Hub.
-
-    IMPORTANT:
-    Les fichiers attendus sont ceux que tu exportes via to_serializable_dict()
-    (format 'zerogpu_aoti_v1') OU un pickle direct de ZeroGPUCompiledModel.
+    Les fichiers du Hub contiennent un dict (payload) créé via ct.to_serializable_dict()
+    sauvegardé par torch.save(). On doit donc:
+      1) torch.load() -> payload dict
+      2) reconstruire un ZeroGPUCompiledModel avec les constantes sur CUDA
     """
     path_1 = hf_hub_download(repo_id=repo_id, filename=filename_1)
     path_2 = hf_hub_download(repo_id=repo_id, filename=filename_2)
 
-    compiled_1 = _load_compiled_pt(path_1)
-    compiled_2 = _load_compiled_pt(path_2)
+    payload_1 = torch.load(path_1, map_location="cpu", weights_only=False)
+    payload_2 = torch.load(path_2, map_location="cpu", weights_only=False)
+
+    compiled_1 = ZeroGPUCompiledModel.from_serializable_dict(payload_1, device=device)
+    compiled_2 = ZeroGPUCompiledModel.from_serializable_dict(payload_2, device=device)
     return compiled_1, compiled_2
 
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-    global COMPILED_TRANSFORMER_1, COMPILED_TRANSFORMER_2
+    """
+    - Si WAN_USE_PRECOMPILED=1 -> essaie de charger depuis le Hub.
+      Si ça échoue (format incompatible, etc.), fallback compilation GPU.
+    - Sinon -> compile sur GPU.
+    """
 
     @spaces.GPU(duration=1500)
     def compile_transformer():
@@ -145,23 +125,28 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         compiled_2 = aoti_compile(exported_2, INDUCTOR_CONFIGS)
         return compiled_1, compiled_2
 
-    # Text encoder quant (inchangé)
+    # Quantize text encoder
     quantize_(pipeline.text_encoder, Int8WeightOnlyConfig())
 
-    use_precompiled = True
+    use_precompiled = _strtobool(os.getenv("WAN_USE_PRECOMPILED"), default=False)
     precompiled_repo = os.getenv("WAN_PRECOMPILED_REPO", "Fabrice-TIERCELIN/Wan_2.2_compiled")
 
+    compiled_transformer_1 = None
+    compiled_transformer_2 = None
+
     if use_precompiled:
-        compiled_transformer_1, compiled_transformer_2 = load_compiled_transformers_from_hub(
-            repo_id=precompiled_repo
-        )
+        try:
+            compiled_transformer_1, compiled_transformer_2 = load_compiled_transformers_from_hub(
+                repo_id=precompiled_repo,
+                device="cuda",
+            )
+            print(f"[WAN] Loaded precompiled transformers from {precompiled_repo}")
+        except Exception as e:
+            print(f"[WAN] Failed to load precompiled transformers, fallback to GPU compile: {e}")
+            compiled_transformer_1, compiled_transformer_2 = compile_transformer()
     else:
         compiled_transformer_1, compiled_transformer_2 = compile_transformer()
 
-    # expose for downloads
-    COMPILED_TRANSFORMER_1 = compiled_transformer_1
-    COMPILED_TRANSFORMER_2 = compiled_transformer_2
-
     pipeline.transformer.forward = compiled_transformer_1
     drain_module_parameters(pipeline.transformer)