Optimize model preloading on ZeroGPU by initializing pipeline pools at module scope

app.py

@@ -3302,6 +3302,7 @@ class PipelinePool:
 
 ACTIVE_PIPELINE_POOL: Optional[PipelinePool] = None
 ACTIVE_POOL_LOCK = threading.Lock()
+GLOBAL_POOLS: dict[str, PipelinePool] = {}
 QUEUE_MAX_SIZE = DEFAULT_QUEUE_SIZE
 QUEUE_CONCURRENCY_LIMIT = DEFAULT_CONCURRENCY_LIMIT
 
@@ -3383,7 +3384,6 @@ def is_pipeline_pool_ready_for_task(task: str) -> bool:
 
 
 def get_pipeline_pool(task: str) -> PipelinePool:
-    global ACTIVE_PIPELINE_POOL
     if not torch.cuda.is_available():
         raise RuntimeError(
             "Lance inference requires a GPU. The Gradio UI can start on CPU, but generation is disabled "
@@ -3391,6 +3391,14 @@ def get_pipeline_pool(task: str) -> PipelinePool:
         )
     model_variant = get_task_model_variant(task)
     gpu_ids = parse_gpu_ids(os.getenv("LANCE_GPUS", DEFAULT_GPUS))
+    
+    if model_variant in GLOBAL_POOLS:
+        pool = GLOBAL_POOLS[model_variant]
+        if not pool.is_initialized:
+            pool.initialize_all()
+        return pool
+
+    global ACTIVE_PIPELINE_POOL
     with ACTIVE_POOL_LOCK:
         if ACTIVE_PIPELINE_POOL is not None and ACTIVE_PIPELINE_POOL.model_variant == model_variant:
             if not ACTIVE_PIPELINE_POOL.is_initialized:
@@ -4214,6 +4222,38 @@ def extract_path(file_data):
     return str(file_data)
 
 
+def initialize_global_pools() -> None:
+    if not torch.cuda.is_available():
+        print("[startup] CUDA is not available, skipping global preloading.", flush=True)
+        return
+
+    variants_text = os.getenv("LANCE_PREFETCH_MODEL_VARIANTS", f"{MODEL_VARIANT_VIDEO},{MODEL_VARIANT_IMAGE}")
+    variants: list[str] = []
+    for raw_variant in variants_text.split(","):
+        raw_variant = raw_variant.strip()
+        if not raw_variant:
+            continue
+        variant = normalize_model_variant(raw_variant)
+        if variant not in variants:
+            variants.append(variant)
+
+    gpu_ids = parse_gpu_ids(os.getenv("LANCE_GPUS", DEFAULT_GPUS))
+    print(f"[startup] Initializing global pools at startup for variants {variants} on GPUs {gpu_ids}", flush=True)
+    
+    for variant in variants:
+        try:
+            print(f"[startup] Instantiating and initializing global pool for: {variant}", flush=True)
+            pool = PipelinePool(gpu_ids, model_variant=variant)
+            pool.initialize_all()
+            GLOBAL_POOLS[variant] = pool
+        except Exception as exc:
+            print(f"[startup] Global pool preloading failed for {variant}: {exc}", flush=True)
+            traceback.print_exc()
+
+
+initialize_global_pools()
+
+
 if __name__ == "__main__":
     args = parse_args()
     os.environ["LANCE_GPUS"] = args.gpus