Upload chatterbox_utils.py with huggingface_hub

chatterbox_utils.py

@@ -116,18 +116,33 @@ def prepare_language(txt, lang_id):
 def load_chatterbox(device="cuda"):
     """Pre-load ONNX sessions"""
     if SESSIONS["speech_encoder"]: return
-    print("🚀 Loading Chatterbox ONNX...")
+    print(f"🚀 Loading Chatterbox ONNX into {device}...")
     opts = onnxruntime.SessionOptions()
-    provs = ["CUDAExecutionProvider"] if device == "cuda" else ["CPUExecutionProvider"]
+    provs = ["CUDAExecutionProvider"] if device == "cuda" and onnxruntime.get_device() == "GPU" else ["CPUExecutionProvider"]
     
     for sess_name in ["speech_encoder", "embed_tokens", "conditional_decoder", "language_model"]:
         fname = "onnx/" + (sess_name + ".onnx" if sess_name != "language_model" else "language_model.onnx")
         path = hf_hub_download(repo_id=MODEL_ID, filename=fname)
-        hf_hub_download(repo_id=MODEL_ID, filename=fname + "_data") # Ensure sidecar data is present
+        hf_hub_download(repo_id=MODEL_ID, filename=fname + "_data", local_files_only=False) # Ensure sidecar data is present
         SESSIONS[sess_name] = onnxruntime.InferenceSession(path, providers=provs)
     
     SESSIONS["tokenizer"] = AutoTokenizer.from_pretrained(MODEL_ID)
 
+def warmup_chatterbox():
+    """CPU-only download of model files for caching"""
+    print("🔥 Warming up Chatterbox (Downloading files)...")
+    try:
+        AutoTokenizer.from_pretrained(MODEL_ID)
+        hf_hub_download(repo_id=MODEL_ID, filename="default_voice.wav")
+        hf_hub_download(repo_id=MODEL_ID, filename="Cangjie5_TC.json")
+        for sess_name in ["speech_encoder", "embed_tokens", "conditional_decoder", "language_model"]:
+            fname = "onnx/" + (sess_name + ".onnx" if sess_name != "language_model" else "language_model.onnx")
+            hf_hub_download(repo_id=MODEL_ID, filename=fname)
+            hf_hub_download(repo_id=MODEL_ID, filename=fname + "_data")
+        print("✅ Chatterbox warmup complete")
+    except Exception as e:
+        print(f"⚠️ Chatterbox warmup warning: {e}")
+
 def run_chatterbox_inference(text, lang_id, speaker_wav_path=None):
     """Ported logic from model card with session reuse"""
     load_chatterbox() # Ensure sessions ready