Update app.py

app.py

@@ -10,8 +10,8 @@ from safetensors.torch import load_file
 
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_REPO = "SebastianBodza/Kartoffelbox-v0.1" 
-T3_CHECKPOINT_FILE = "t3_kartoffelbox.safetensors"
+MODEL_REPO = "SebastianBodza/Kartoffelbox-v0.2" 
+T3_CHECKPOINT_FILE = " t3_global_step_50.pt"
 print(f"🚀 Running on device: {DEVICE}")
 
 # --- Global Model Initialization ---
@@ -26,8 +26,12 @@ def get_or_load_model():
         try:
             MODEL = ChatterboxTTS.from_pretrained(DEVICE)
             checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename=T3_CHECKPOINT_FILE, token=os.environ["HUGGING_FACE_HUB_TOKEN"])
-            t3_state = load_file(checkpoint_path, device="cpu") 
-            MODEL.t3.load_state_dict(t3_state)
+            t3_state = torch.load(checkpoint_path, device="cpu") 
+            cleaned_state_dict = {
+                key.replace("_orig_mod.", ""): value
+                for key, value in t3_state.items()
+            }
+            MODEL.t3.load_state_dict(cleaned_state_dict)
 
             if hasattr(MODEL, 'to') and str(MODEL.device) != DEVICE:
                 MODEL.to(DEVICE)