Upload app.py with huggingface_hub

app.py

@@ -73,6 +73,26 @@ def _distributed_addr(self):
 
 EngineConfig.distributed_addr = property(_distributed_addr)
 
+import zonos2.engine.engine as zonos2_engine
+from zonos2.models.weight import _normalize_zonos2_state_dict
+
+# Deserialize the 15.3 GB checkpoint once in the main process (mmap keeps it
+# page-cache backed); forked GPU workers inherit it copy-on-write, so cold
+# engine init skips the ~17s torch.load and only pays the host->device copy.
+_STATE_DICT = torch.load(
+    f"{MODEL_PATH}/model.pth", map_location="cpu", weights_only=False, mmap=True
+)
+if "model" in _STATE_DICT:
+    _STATE_DICT = _STATE_DICT["model"]
+_STATE_DICT = _normalize_zonos2_state_dict(_STATE_DICT)
+
+
+def _preloaded_checkpoint_weight(model_path, device):
+    return {k: v.to(device) for k, v in _STATE_DICT.items()}
+
+
+zonos2_engine.load_checkpoint_weight = _preloaded_checkpoint_weight
+
 LANGUAGES = {
     "English (US)": "en_us",
     "English (UK)": "en_gb",
@@ -146,6 +166,10 @@ def _embed_speaker(models, speaker_audio):
         wav = wav.astype(np.float32)
     if wav.ndim == 2:
         wav = wav.T  # (samples, channels) -> (channels, samples)
+    else:
+        # The embedder's reflect-pad requires a 2D (channels, samples) input;
+        # mono uploads arrive 1D.
+        wav = wav[None, :]
     wav_t = torch.from_numpy(wav)
 
     embedder = models["embedder"]