Update app.py

app.py

@@ -164,47 +164,72 @@ except ImportError:
     except ImportError:
         QwenDoubleStreamAttnProcessorFA3 = None
 
+# ── dtype: bfloat16 on GPU (fast), float32 on CPU (fallback) ──────────────
 dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
+# ── Lazy pipeline: only load when a GPU is actually present ───────────────
+# Loading a 10 GB model on CPU OOMs free-tier Spaces (16 GB RAM limit).
+# We defer loading to first inference; if no GPU is present the UI still
+# starts and shows a clear warning instead of crashing at startup.
+pipe = None
+_pipe_loaded_version: str | None = None
+
 
 def _load_pipe_with_version(version: str) -> QwenImageEditPlusPipeline:
     sub = f"{version}/transformer"
     print(f"📦 Loading AIO transformer: {AIO_REPO_ID} / {sub}")
-    device_map = "cuda" if torch.cuda.is_available() else "cpu"
     p = QwenImageEditPlusPipeline.from_pretrained(
         "Qwen/Qwen-Image-Edit-2511",
         transformer=QwenImageTransformer2DModel.from_pretrained(
             AIO_REPO_ID,
             subfolder=sub,
             torch_dtype=dtype,
-            device_map=device_map,
+            device_map="cuda",          # only called when CUDA is confirmed
         ),
         torch_dtype=dtype,
     ).to(device)
     return p
 
 
-# Forgiving load: try env/default version, fallback to v19 if it fails
-try:
-    pipe = _load_pipe_with_version(AIO_VERSION)
-except Exception as e:
-    print("❌ Failed to load requested AIO_VERSION. Falling back to v19.")
-    print("---- exception ----")
-    print(traceback.format_exc())
-    print("-------------------")
-    AIO_VERSION = DEFAULT_AIO_VERSION
-    AIO_VERSION_SOURCE = "fallback_to_v19"
-    pipe = _load_pipe_with_version(AIO_VERSION)
-
-# Apply FA3 Optimization
-try:
-    if QwenDoubleStreamAttnProcessorFA3 is not None:
-        pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-        print("Flash Attention 3 Processor set successfully.")
-    else:
-        print("Warning: QwenDoubleStreamAttnProcessorFA3 not available in this diffusers version; skipping.")
-except Exception as e:
-    print(f"Warning: Could not set FA3 processor: {e}")
+def _ensure_pipe_loaded() -> None:
+    """Load the pipeline on first call (requires CUDA). Raises if no GPU."""
+    global pipe, _pipe_loaded_version, AIO_VERSION, AIO_VERSION_SOURCE
+    if pipe is not None:
+        return
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "No GPU detected. This Space requires a CUDA GPU to run. "
+            "Go to Space Settings > Hardware and select a GPU tier (e.g. T4 Medium)."
+        )
+    try:
+        pipe = _load_pipe_with_version(AIO_VERSION)
+        _pipe_loaded_version = AIO_VERSION
+    except Exception as e:
+        print("❌ Failed to load requested AIO_VERSION. Falling back to v19.")
+        print(traceback.format_exc())
+        AIO_VERSION = DEFAULT_AIO_VERSION
+        AIO_VERSION_SOURCE = "fallback_to_v19"
+        pipe = _load_pipe_with_version(AIO_VERSION)
+        _pipe_loaded_version = AIO_VERSION
+
+    # Apply FA3 Optimization
+    try:
+        if QwenDoubleStreamAttnProcessorFA3 is not None:
+            pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+            print("Flash Attention 3 Processor set successfully.")
+        else:
+            print("Warning: QwenDoubleStreamAttnProcessorFA3 not available; skipping.")
+    except Exception as e:
+        print(f"Warning: Could not set FA3 processor: {e}")
+
+
+if torch.cuda.is_available():
+    # GPU available at startup — load immediately so first request is fast
+    _ensure_pipe_loaded()
+else:
+    print("⚠️  No GPU at startup — pipeline will NOT be loaded.")
+    print("    The UI will start but inference will fail until a GPU is assigned.")
+    print("    Fix: Space Settings → Hardware → select a GPU tier.")
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -754,6 +779,12 @@ def infer(
     steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    # Ensure pipeline is loaded (raises friendly error if no GPU)
+    try:
+        _ensure_pipe_loaded()
+    except RuntimeError as e:
+        raise gr.Error(str(e))
+
     gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -860,9 +891,20 @@ aio_status_line = (
     f"({AIO_VERSION_SOURCE}; env `AIO_VERSION`={_AIO_ENV_RAW!r})"
 )
 
+NO_GPU_WARNING = not torch.cuda.is_available()
+
 with gr.Blocks() as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# **Qwen-Image-Edit-2511-LoRAs-Fast**", elem_id="main-title")
+        if NO_GPU_WARNING:
+            gr.Markdown(
+                "## ⚠️ No GPU detected — inference is disabled\n\n"
+                "This model requires a **CUDA GPU** to run. The app has started, but clicking "
+                "**Edit Image** will show an error.\n\n"
+                "**Fix:** Go to your Space → ⚙️ **Settings** → **Hardware** → "
+                "select **T4 Medium** (or better) → **Save**. The Space will rebuild with a GPU.",
+                visible=True,
+            )
         gr.Markdown(
             "Perform diverse image edits using specialized "
             "[LoRA](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2511) adapters for the "