Update app.py

app.py

@@ -22,15 +22,7 @@ EXEC_DEVICE = "cpu"
 FORCE_CPU = bool(int(os.getenv("FORCE_CPU", "0")))
 
 
-@spaces.GPU  # bắt buộc cho ZeroGPU
-def load_pipeline():
-    global EXEC_DEVICE
-
-    # Ưu tiên GPU của ZeroGPU; cho phép ép CPU để tránh abort/OOM.
-    device = "cuda" if torch.cuda.is_available() and not FORCE_CPU else "cpu"
-    dtype = torch.float16 if device == "cuda" else torch.float32
-    EXEC_DEVICE = device
-
+def _build_pipeline(device: str, dtype: torch.dtype):
     pipe = QwenImageEditPlusPipeline.from_pretrained(
         HF_BASE_MODEL,
         torch_dtype=dtype,
@@ -53,6 +45,30 @@ def load_pipeline():
     return pipe
 
 
+@spaces.GPU  # bắt buộc cho ZeroGPU
+def load_pipeline():
+    global EXEC_DEVICE
+
+    # Ưu tiên GPU của ZeroGPU; cho phép ép CPU để tránh abort/OOM.
+    prefer_cuda = torch.cuda.is_available() and not FORCE_CPU
+    device = "cuda" if prefer_cuda else "cpu"
+    dtype = torch.float16 if device == "cuda" else torch.float32
+
+    try:
+        pipe = _build_pipeline(device, dtype)
+        EXEC_DEVICE = device
+        return pipe
+    except Exception as exc:
+        # GPU worker thường abort vì OOM. Fallback CPU để không crash app.
+        if device == "cuda":
+            device = "cpu"
+            dtype = torch.float32
+            pipe = _build_pipeline(device, dtype)
+            EXEC_DEVICE = device
+            return pipe
+        raise exc
+
+
 pipe = load_pipeline()