Load on the CPU in the background and move to the GPU inside process_fn.

app.py

@@ -46,17 +46,13 @@ def load_model():
         vocoder_checkpoint_path = os.path.join(model_dir, "pretrained/vocoder")
 
         print("Initializing AnyAccomp InferencePipeline...", flush=True)
-        # accelerate loads a checkpoint into a model already on CUDA through calls
-        # that ZeroGPU's emulation does not intercept, so the weights load on the
-        # CPU and the modules move afterwards
+        # ZeroGPU does not intercept CUDA calls made from a background thread, so
+        # the pipeline loads on the CPU here and process_fn moves it to the GPU
         pipeline = Sing2SongInferencePipeline(
             checkpoint_path, cfg_path,
             vocoder_checkpoint_path, vocoder_cfg_path,
             device="cpu",
         )
-        for module in (pipeline.model, pipeline.coco_model, pipeline.vocoder):
-            module.to(DEVICE)
-        pipeline.device = DEVICE
         pipeline.sample_rate = 24000
         inference_pipeline = pipeline
         print("Model loaded successfully.", flush=True)
@@ -104,6 +100,11 @@ def process_fn(vocal_filepath, n_timesteps):  # add , cfg_scale to re-enable
     if inference_pipeline is None:
         raise gr.Error(f"Model failed to load: {model_error}")
 
+    pipeline = inference_pipeline
+    for module in (pipeline.model, pipeline.coco_model, pipeline.vocoder):
+        module.to(DEVICE)
+    inference_pipeline.device = DEVICE
+
     if vocal_filepath is None:
         raise gr.Error("Please upload a vocal audio file.")