Defer GPU move to inside process_fn (ZeroGPU can't touch CUDA from a background thread)

app.py

@@ -68,7 +68,7 @@ VARIANTS = {
 # which LLM layers it extracts at construction time). Switching variants evicts the
 # previous one and rebuilds from the local HF cache (no re-download).
 _active_name = None
-_active_entry = None
+_active_entry = None    # built on CPU at first; moved to GPU on first real use
 _loading = True
 _load_error = None
 
@@ -76,11 +76,15 @@ _load_error = None
 # so two requests running together would corrupt each other's results. This lock
 # also covers _active_name/_active_entry below, so no separate lock is needed.
 _gpu_lock = threading.Lock()
+_active_ready = False  # has _active_entry been moved onto the GPU yet?
 
 
 def _build_variant(name):
-    """Download and construct one variant's full stack: text encoder, DiT backbone,
-    and VAE. Mirrors the model-loading steps in unison/pipelines/infer.py's main()."""
+    """Download and construct one variant's full stack on CPU: text encoder, DiT
+    backbone, and VAE. Safe to call from the background thread — nothing here
+    touches CUDA, since ZeroGPU only intercepts CUDA calls made from inside an
+    @spaces.GPU-decorated call. Moving onto the GPU happens later, in
+    _activate_variant()."""
     spec = VARIANTS[name]
     print(f"Building variant: {name} ...")
 
@@ -90,12 +94,12 @@ def _build_variant(name):
     omni_last_layer_idx = model_config.get("omni_last_layer_idx", -1)
 
     extractor = init_text_hidden_extractor(
-        "omni", OMNI_MODEL_ID, None, dit_depth, DEVICE, omni_last_layer_idx=omni_last_layer_idx,
+        "omni", OMNI_MODEL_ID, None, dit_depth, "cpu", omni_last_layer_idx=omni_last_layer_idx,
     )
     sync_omni_dim_with_text_encoder(model_config, extractor)
 
     ckpt_path = hf_hub_download(repo_id=UNISON_REPO, filename=spec["ckpt_file"])
-    model = _load_model(types.SimpleNamespace(model_ckpt=ckpt_path), DEVICE, model_config)
+    model = _load_model(types.SimpleNamespace(model_ckpt=ckpt_path), "cpu", model_config)
 
     vae_ckpt_path = hf_hub_download(repo_id=MMAUDIO_REPO, filename=spec["vae_ckpt_file"])
     vocoder_ckpt_path = (
@@ -106,24 +110,37 @@ def _build_variant(name):
         tod_vae_ckpt=vae_ckpt_path,
         bigvgan_vocoder_ckpt=vocoder_ckpt_path,
         mode=spec["vae_mode"],
-    )
-    audio_vae.to(DEVICE).eval()
+    ).eval()
+
+    scheduler = FlowMatchEulerDiscreteScheduler()
+    print(f"Variant built (CPU): {name}")
+    # gen_target_frames is None until _activate_variant() probes it on the GPU.
+    return (model, extractor, audio_vae, 0.5, spec["sample_rate"], None, scheduler)
+
+
+def _activate_variant(entry):
+    """Move a CPU-built variant onto the GPU and probe its latent frame count.
+    Must only be called from inside process_fn (i.e. inside @spaces.GPU) — this
+    is where it's actually safe to touch CUDA."""
+    model, extractor, audio_vae, vae_scale_factor, sample_rate, _, scheduler = entry
+    extractor.text_backbone = extractor.text_backbone.to(DEVICE)
+    model = model.to(device=DEVICE, dtype=torch.bfloat16).eval()
+    audio_vae = audio_vae.to(DEVICE).eval()
 
     # probe latent length for a MAX_AUDIO_DURATION-long clip (needed for generation mode).
-    sample_rate = spec["sample_rate"]
     dummy_len = int(MAX_AUDIO_DURATION * sample_rate)
     with torch.no_grad():
         dummy_lat = audio_vae.wrapped_encode(torch.zeros(1, dummy_len, device=DEVICE))
     gen_target_frames = int(dummy_lat.shape[-1])
 
-    scheduler = FlowMatchEulerDiscreteScheduler()
-    print(f"Variant ready: {name}")
-    return (model, extractor, audio_vae, 0.5, sample_rate, gen_target_frames, scheduler)
+    print("Variant activated (GPU)")
+    return (model, extractor, audio_vae, vae_scale_factor, sample_rate, gen_target_frames, scheduler)
 
 
 def load_default_variant():
-    """Background-thread target: build DEFAULT_VARIANT at startup and publish it
-    as the active variant, so the Space doesn't block its HTTP server on the load."""
+    """Background-thread target: build DEFAULT_VARIANT (CPU only) at startup and
+    publish it as the active variant, so the Space doesn't block its HTTP server
+    on the download. Actually moving it to the GPU happens on the first request."""
     global _active_name, _active_entry, _loading, _load_error
     try:
         entry = _build_variant(DEFAULT_VARIANT)
@@ -141,21 +158,24 @@ threading.Thread(target=load_default_variant, daemon=True).start()
 
 
 def get_variant(name: str):
-    """Return the active variant's (model, extractor, vae, ...) tuple, building and
-    switching to `name` first if a different variant is currently active.
-
-    Caller must hold _gpu_lock — that's what makes the unlocked reads/writes of
-    _active_name/_active_entry here safe."""
-    global _active_name, _active_entry
-    if name == _active_name:
-        return _active_entry
-
-    print(f"Switching variant: {_active_name!r} -> {name!r}")
-    entry = _build_variant(name)
-    _active_name, _active_entry = name, entry
-    gc.collect()
-    torch.cuda.empty_cache()
-    return entry
+    """Return the active variant's (model, extractor, vae, ...) tuple, ready to
+    use on the GPU — building and/or activating it first if needed.
+
+    Caller must hold _gpu_lock and be inside an @spaces.GPU call — that's both
+    what makes the unlocked _active_* globals safe, and what makes it safe to
+    touch CUDA in _activate_variant()."""
+    global _active_name, _active_entry, _active_ready
+    if name != _active_name:
+        print(f"Switching variant: {_active_name!r} -> {name!r}")
+        _active_name, _active_entry, _active_ready = name, _build_variant(name), False
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    if not _active_ready:
+        _active_entry = _activate_variant(_active_entry)
+        _active_ready = True
+
+    return _active_entry
 
 
 model_card = ModelCard(

unison/models/text_encoders/omni_encoder.py

@@ -108,9 +108,13 @@ class QwenOmniThinkerExtractor(nn.Module):
 
         gc.collect()
 
-        print(f"Moving text backbone to {device}...")
-        self.text_backbone = self.text_backbone.to(device)
-        torch.cuda.empty_cache()
+        # Skip when device is cpu — ZeroGPU only intercepts CUDA calls made inside
+        # an @spaces.GPU-decorated call, so touching CUDA here would crash if this
+        # constructor is ever run from a background thread (see app.py).
+        if str(device) != "cpu":
+            print(f"Moving text backbone to {device}...")
+            self.text_backbone = self.text_backbone.to(device)
+            torch.cuda.empty_cache()
 
         # Load processor; force right-side padding so the system prompt is always
         # at index 0, making the slice offset constant.