Update app.py

app.py

@@ -1,10 +1,10 @@
 import os
-
-# Bỏ qua yêu cầu torchvision (ZeroGPU không cài sẵn).
-os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1")
+import subprocess
+import sys
 
 import spaces
 import torch
+from optimization import optimize_pipeline_
 from diffusers import QwenImageEditPlusPipeline
 from PIL import Image
 import gradio as gr
@@ -24,7 +24,32 @@ FORCE_CPU = bool(int(os.getenv("FORCE_CPU", "1")))
 PREFER_GPU = bool(int(os.getenv("PREFER_GPU", "0")))
 
 
+def ensure_torchvision():
+    """
+    Qwen2VLProcessor yêu cầu torchvision. Thử import, nếu thiếu sẽ cài đặt phiên bản khớp torch.
+    """
+    try:
+        import torchvision  # noqa: F401
+        return
+    except ImportError:
+        torch_version = torch.__version__.split("+")[0]
+        try:
+            subprocess.check_call(
+                [sys.executable, "-m", "pip", "install", f"torchvision=={torch_version}"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            import torchvision  # noqa: F401
+        except Exception as exc:  # pragma: no cover - chỉ chạy trên hạ tầng Spaces
+            raise ImportError(
+                "Torchvision is required for Qwen2VLProcessor. "
+                "Please add a matching torchvision to requirements (e.g. pip install torchvision==torch_version)."
+            ) from exc
+
+
 def _build_pipeline(device: str, dtype: torch.dtype):
+    ensure_torchvision()
+
     pipe = QwenImageEditPlusPipeline.from_pretrained(
         HF_BASE_MODEL,
         torch_dtype=dtype,
@@ -47,6 +72,35 @@ def _build_pipeline(device: str, dtype: torch.dtype):
     return pipe
 
 
+def maybe_optimize_pipeline(pipe):
+    """
+    Áp dụng AOTI tối ưu hóa trên GPU (nếu đang chạy CUDA).
+    Dùng input dummy nhỏ để tránh tốn VRAM, fallback im lặng nếu lỗi.
+    """
+    if EXEC_DEVICE != "cuda":
+        return pipe
+    try:
+        dummy = Image.new("RGB", (256, 256))
+        generator = torch.Generator(device="cuda").manual_seed(0)
+        optimize_pipeline_(
+            pipe,
+            image=[dummy, dummy],
+            prompt="warmup",
+            negative_prompt=" ",
+            num_inference_steps=1,
+            true_cfg_scale=1.0,
+            guidance_scale=1.0,
+            num_images_per_prompt=1,
+            generator=generator,
+            width=256,
+            height=256,
+        )
+    except Exception:
+        # Nếu tối ưu thất bại (thường do bộ nhớ), giữ pipeline gốc để tiếp tục chạy.
+        pass
+    return pipe
+
+
 @spaces.GPU  # bắt buộc cho ZeroGPU
 def load_pipeline():
     global EXEC_DEVICE
@@ -59,6 +113,7 @@ def load_pipeline():
     try:
         pipe = _build_pipeline(device, dtype)
         EXEC_DEVICE = device
+        pipe = maybe_optimize_pipeline(pipe)
         return pipe
     except Exception as exc:
         # GPU worker thường abort vì OOM. Fallback CPU để không crash app.