Fit FLUX hot swap within 16 GiB VRAM

platform/docker/flux-worker/flux_worker.py

@@ -30,7 +30,15 @@ def load_pipeline() -> None:
     from diffusers import DiffusionPipeline
     started = time.monotonic()
     PIPE = DiffusionPipeline.from_pretrained(
-        MODEL_DIR, torch_dtype=torch.bfloat16, device_map="cuda")
+        MODEL_DIR, torch_dtype=torch.bfloat16)
+    # The full pipeline leaves too little activation headroom on a 16 GiB
+    # RTX 5080. Model CPU offload keeps each active component on CUDA while
+    # parking inactive components in system RAM between the four steps.
+    PIPE.enable_model_cpu_offload()
+    if hasattr(PIPE, "enable_vae_slicing"):
+        PIPE.enable_vae_slicing()
+    if hasattr(PIPE, "enable_vae_tiling"):
+        PIPE.enable_vae_tiling()
     LOAD_SECONDS = time.monotonic() - started
 
 
@@ -98,6 +106,8 @@ class Handler(BaseHTTPRequestHandler):
                 raise ValueError("invalid request size")
             self.reply(200, generate(json.loads(self.rfile.read(length))))
         except Exception as exc:
+            print(f"[flux-worker] generation failed: "
+                  f"{type(exc).__name__}: {str(exc)[:1000]}", flush=True)
             self.reply(400, {"status": "error", "message": str(exc)})