ZeroGPU: load model on CPU at startup, move to GPU inside @spaces.GPU

app.py

@@ -46,13 +46,16 @@ MODEL_REPO = "MenaVoice/KasbahTTS-V0"
 CKPT_FILE = "ALGERIA.safetensors"
 V1_CFG = dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)
 
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# On ZeroGPU, CUDA exists only *inside* a @spaces.GPU call — never at import time.
+# So load everything on CPU here and move to GPU inside generate().
+ON_ZEROGPU = USING_ZEROGPU and bool(os.environ.get("SPACES_ZERO_GPU"))
+GPU_DEVICE = "cuda" if (ON_ZEROGPU or torch.cuda.is_available()) else "cpu"
 
 ckpt_path = hf_hub_download(MODEL_REPO, CKPT_FILE)
 vocab_path = hf_hub_download(MODEL_REPO, "vocab.txt")
 
-vocoder = load_vocoder(vocoder_name="vocos", is_local=False, device=DEVICE)
-model = load_model(DiT, V1_CFG, ckpt_path, vocab_file=vocab_path, device=DEVICE)
+vocoder = load_vocoder(vocoder_name="vocos", is_local=False, device="cpu")
+model = load_model(DiT, V1_CFG, ckpt_path, vocab_file=vocab_path, device="cpu")
 
 EXAMPLE_REF = os.path.join(os.path.dirname(__file__), "examples", "ALG.wav")
 EXAMPLE_REF_TEXT = "أنيا هكا باغية ناكل هكا أني ن نشوف فيها الحاجة هذيكا."
@@ -94,21 +97,33 @@ def generate(
     # Leaving ref_text blank triggers Whisper auto-transcription inside f5-tts.
     ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text or "", show_info=gr.Info)
 
-    # KasbahTTS is a *specialized* single-dialect checkpoint, so no dialect token
-    # is prepended — infer_process defaults dialect_id to None.
-    wave, sample_rate, _ = infer_process(
-        ref_audio,
-        ref_text,
-        gen_text,
-        model,
-        vocoder,
-        cross_fade_duration=cross_fade,
-        nfe_step=int(nfe_step),
-        cfg_strength=cfg_strength,
-        speed=speed,
-        show_info=gr.Info,
-        progress=progress,
-    )
+    # Move to GPU now that we're inside the @spaces.GPU context (loaded on CPU at startup).
+    model.to(GPU_DEVICE)
+    vocoder.to(GPU_DEVICE)
+    try:
+        # KasbahTTS is a *specialized* single-dialect checkpoint, so no dialect
+        # token is prepended — infer_process defaults dialect_id to None.
+        wave, sample_rate, _ = infer_process(
+            ref_audio,
+            ref_text,
+            gen_text,
+            model,
+            vocoder,
+            cross_fade_duration=cross_fade,
+            nfe_step=int(nfe_step),
+            cfg_strength=cfg_strength,
+            speed=speed,
+            show_info=gr.Info,
+            progress=progress,
+            device=GPU_DEVICE,
+        )
+    finally:
+        # ZeroGPU tears down CUDA after the call — return the model to CPU so the
+        # next call starts from a valid copy.
+        if ON_ZEROGPU:
+            model.to("cpu")
+            vocoder.to("cpu")
+            torch.cuda.empty_cache()
 
     if remove_sil:
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: