[Admin maintenance] Support new ZeroGPU hardware (#2)

- [Admin maintenance] Support new ZeroGPU hardware (aeda8bb88a02244d07835bfe924c0fb9cd945e48)

app.py

@@ -1,6 +1,15 @@
 # ruff: noqa: E402
 # Above allows ruff to ignore E402: module level import not at top of file
 
+# `spaces` must be imported before torch (or anything that touches CUDA) so
+# it can patch torch.cuda.* for the ZeroGPU worker hijack.
+try:
+    import spaces
+
+    USING_SPACES = True
+except ImportError:
+    USING_SPACES = False
+
 import gc
 import json
 import re
@@ -17,13 +26,6 @@ import torchaudio
 from cached_path import cached_path
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-try:
-    import spaces
-
-    USING_SPACES = True
-except ImportError:
-    USING_SPACES = False
-
 
 def gpu_decorator(func):
     if USING_SPACES:
@@ -58,16 +60,21 @@ DEFAULT_TTS_MODEL_CFG = [
 vocoder = load_vocoder()
 
 
+# On ZeroGPU, CUDA is not available in the main process (the GPU is only
+# attached inside @spaces.GPU workers). f5_tts.load_model calls
+# safetensors.load_file(..., device="cuda") directly, which cannot be deferred
+# by the `spaces` patcher and raises "No CUDA GPUs are available" at import.
+# Load every model on CPU here, then move it onto the GPU inside the worker.
 def load_f5tts():
     ckpt_path = str(cached_path(DEFAULT_TTS_MODEL_CFG[0]))
     F5TTS_model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])
-    return load_model(DiT, F5TTS_model_cfg, ckpt_path)
+    return load_model(DiT, F5TTS_model_cfg, ckpt_path, device="cpu")
 
 
 def load_e2tts():
     ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors"))
     E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4, text_mask_padding=False, pe_attn_head=1)
-    return load_model(UNetT, E2TTS_model_cfg, ckpt_path)
+    return load_model(UNetT, E2TTS_model_cfg, ckpt_path, device="cpu")
 
 
 def load_custom(ckpt_path: str, vocab_path="", model_cfg=None):
@@ -78,7 +85,7 @@ def load_custom(ckpt_path: str, vocab_path="", model_cfg=None):
         vocab_path = str(cached_path(vocab_path))
     if model_cfg is None:
         model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])
-    return load_model(DiT, model_cfg, ckpt_path, vocab_file=vocab_path)
+    return load_model(DiT, model_cfg, ckpt_path, vocab_file=vocab_path, device="cpu")
 
 
 F5TTS_ema_model = load_f5tts()
@@ -151,6 +158,12 @@ def infer(
             pre_custom_path = model[1]
         ema_model = custom_ema_model
 
+    # Models were loaded on CPU at import (ZeroGPU has no GPU in the main
+    # process); move the selected model onto the GPU now that we're inside the
+    # @spaces.GPU worker. The vocoder was loaded with a deferred .to("cuda") so
+    # `spaces` already materialized it on the GPU.
+    ema_model = ema_model.to("cuda")
+
     final_wave, final_sample_rate, combined_spectrogram = infer_process(
         ref_audio,
         ref_text,

requirements.txt

@@ -1,5 +1,5 @@
-torch==2.4.0
-torchaudio
+torch==2.8.0
+torchaudio==2.8.0
 accelerate>=0.33.0
 bitsandbytes>0.37.0
 cached_path
@@ -23,4 +23,4 @@ vocos
 wandb
 x_transformers>=1.31.14
 f5_tts
-detoxify @ git+https://github.com/unitaryai/detoxify
+detoxify @ git+https://github.com/unitaryai/detoxify