Move Flux2 download to startup (CPU) to avoid burning GPU time

from_pretrained + load_lora_weights now run at module level on CPU —
no GPU seconds consumed. Inside @spaces.GPU, only pipe.to("cuda")
runs (fast memory transfer). GPU duration reduced from 300s to 120s.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -17,7 +17,7 @@ from image_utils import _b64_to_pil, _collect_source_pils
 from validators import _validate_target, _validate_selection, _validate_sources
 from logging_utils import spawn_log
 from pipeline import detect as _pipeline_detect, run_swap, MAX_FACES
-from pipeline_flux import run_swap_flux, load_flux_pipeline, MAX_FLUX_FACES
+from pipeline_flux import run_swap_flux, load_flux_pipeline_cpu, move_flux_to_device, MAX_FLUX_FACES
 
 # basicsr (gfpgan dep) imports this submodule which was removed in torchvision>=0.17
 import sys
@@ -103,8 +103,13 @@ _cpu_face_analyzer = FaceAnalysis(name="buffalo_l",
 _cpu_face_analyzer.prepare(ctx_id=0, det_size=(640, 640))
 print("[startup] CPU FaceAnalysis ready.", flush=True)
 
-# ── Flux2 pipeline (lazy — loaded on first Quality request inside @spaces.GPU) ─
-_flux_pipe   = None
+# ── Flux2 pipeline: weights downloaded on CPU at startup (no GPU cost),
+#    then moved to CUDA inside @spaces.GPU on the first Quality request. ──────
+print("[startup] loading Flux2 pipeline on CPU (download may take a while)…", flush=True)
+_flux_pipe_cpu = load_flux_pipeline_cpu()
+print("[startup] Flux2 CPU pipeline ready.", flush=True)
+
+_flux_pipe   = None   # set to the CUDA pipe on first Quality call
 _flux_device = None
 
 
@@ -182,7 +187,7 @@ def swap_faces(target_b64, selected_json, sources_json):
         print("[swap] ===== END =====", flush=True)
 
 
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=120)
 def swap_faces_quality(target_b64, selected_json, sources_json):
     """Flux2 Quality swap: per-face head crop → Flux2 → paste back. Generator: yields (image, status)."""
     global _flux_pipe, _flux_device
@@ -197,10 +202,10 @@ def swap_faces_quality(target_b64, selected_json, sources_json):
 
     target_pil = _b64_to_pil(target_b64)
 
-    yield gr.update(), "Loading Quality model…"
+    yield gr.update(), "Moving Quality model to GPU…"
     if _flux_pipe is None:
-        _flux_pipe, _flux_device = load_flux_pipeline()
-        print("[flux] model loaded", flush=True)
+        _flux_pipe, _flux_device = move_flux_to_device(_flux_pipe_cpu)
+        print("[flux] pipeline on GPU ready", flush=True)
 
     yield gr.update(), "Swapping faces (Quality)…"
     try:

pipeline_flux.py

@@ -43,25 +43,31 @@ def _resize_for_flux(pil_img):
     return pil_img.resize((new_w, new_h), Image.LANCZOS), (w, h)
 
 
-def load_flux_pipeline():
-    """Download and initialise the FLUX.2 Klein pipeline with the face-swap LoRA."""
+def load_flux_pipeline_cpu():
+    """Download weights and load onto CPU — safe to call at startup outside @spaces.GPU."""
     import os
     from diffusers import Flux2KleinPipeline
-    token  = os.environ.get("HF_TOKEN") or None
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"[flux] loading pipeline on {device}…", flush=True)
+    token = os.environ.get("HF_TOKEN") or None
+    print("[flux] downloading / loading pipeline on CPU…", flush=True)
     pipe = Flux2KleinPipeline.from_pretrained(
         "black-forest-labs/FLUX.2-klein-9B",
         torch_dtype=torch.bfloat16,
         token=token,
     )
-    pipe.to(device)
     pipe.load_lora_weights(
         "Alissonerdx/BFS-Best-Face-Swap",
         weight_name="bfs_head_v1_flux-klein_9b_step3750_rank64.safetensors",
         token=token,
     )
-    print("[flux] pipeline ready.", flush=True)
+    print("[flux] pipeline on CPU ready.", flush=True)
+    return pipe
+
+
+def move_flux_to_device(pipe):
+    """Move the CPU-loaded pipeline to the available device. Call inside @spaces.GPU."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"[flux] moving pipeline to {device}…", flush=True)
+    pipe.to(device)
     return pipe, device