Load all model variants at module scope for ZeroGPU

app.py

@@ -3,15 +3,12 @@
 Gradio Space demo with a single unified interface: image presence selects
 editing vs. generation, while the model control selects fast vs. quality.
 """
-import gc
 import os
-import threading
 
 # Use flash_attention_2 for the HF text encoder (flash_attn is installed via wheel)
 os.environ.setdefault("VF_HF_ATTN_IMPL", "flash_attention_2")
 
 import spaces  # MUST be first (after env setup)
-import torch
 import gradio as gr
 from PIL import Image
 
@@ -28,27 +25,16 @@ MODEL_VARIANTS = {
     },
 }
 
-_pipe_slots = {
-    "t2i": {"variant": "turbo", "pipe": MageFlowPipeline.from_pretrained(MODEL_VARIANTS["turbo"]["t2i"], device="cuda")},
-    "edit": {"variant": "turbo", "pipe": MageFlowPipeline.from_pretrained(MODEL_VARIANTS["turbo"]["edit"], device="cuda")},
+# ZeroGPU requires every model to be placed on CUDA at module scope: the backend
+# registers the weights, offloads them to disk at startup, and streams them into
+# VRAM for each @spaces.GPU call. Loading inside the GPU function instead would
+# charge every switch to the caller's GPU quota and is not carried across the
+# forked GPU workers.
+PIPES = {
+    (task, variant): MageFlowPipeline.from_pretrained(spec[task], device="cuda")
+    for variant, spec in MODEL_VARIANTS.items()
+    for task in ("t2i", "edit")
 }
-_pipe_lock = threading.Lock()
-
-
-def _get_pipe(task: str, variant: str):
-    """Keep one loaded variant per task, matching the original two-pipeline footprint."""
-    with _pipe_lock:
-        slot = _pipe_slots.get(task)
-        if slot and slot["variant"] == variant:
-            return slot["pipe"]
-        if slot:
-            del _pipe_slots[task]
-            del slot
-            gc.collect()
-            torch.cuda.empty_cache()
-        pipe = MageFlowPipeline.from_pretrained(MODEL_VARIANTS[variant][task], device="cuda")
-        _pipe_slots[task] = {"variant": variant, "pipe": pipe}
-        return pipe
 
 
 def _recommended(variant: str, image):
@@ -91,7 +77,7 @@ def generate(
 
     if image is not None:
         # Route to the edit model when an image is provided.
-        pipe_edit = _get_pipe("edit", model_variant)
+        pipe_edit = PIPES[("edit", model_variant)]
         if isinstance(image, str):
             image = Image.open(image)
         refs = [image.convert("RGB")]
@@ -115,7 +101,7 @@ def generate(
 
     # No image: route to the text-to-image model.
     # Content-safety gate: blocked requests return a blank image.
-    pipe_t2i = _get_pipe("t2i", model_variant)
+    pipe_t2i = PIPES[("t2i", model_variant)]
     verdict = pipe_t2i.model.txt_enc.screen_text(prompt)
     if verdict.violates:
         return Image.new("RGB", (int(width), int(height)), (255, 255, 255))