Update app.py

app.py

@@ -28,15 +28,15 @@ def get_or_load_model():
             checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename=T3_CHECKPOINT_FILE, token=os.environ["HUGGING_FACE_HUB_TOKEN"])
             t3_state = torch.load(checkpoint_path) 
             
-            if "cond_enc.spkr_enc.0.weight" in t3_state:
-                t3_state["cond_enc.spkr_enc.weight"] = t3_state.pop("cond_enc.spkr_enc.0.weight")
-            if "cond_enc.spkr_enc.0.bias" in t3_state:
-                t3_state["cond_enc.spkr_enc.bias"] = t3_state.pop("cond_enc.spkr_enc.0.bias")
-
             cleaned_state_dict = {
                 key.replace("_orig_mod.", ""): value
                 for key, value in t3_state.items()
             }
+            
+            if "cond_enc.spkr_enc.0.weight" in cleaned_state_dict:
+                cleaned_state_dict["cond_enc.spkr_enc.weight"] = cleaned_state_dict.pop("cond_enc.spkr_enc.0.weight")
+            if "cond_enc.spkr_enc.0.bias" in t3_state:
+                cleaned_state_dict["cond_enc.spkr_enc.bias"] = cleaned_state_dict.pop("cond_enc.spkr_enc.0.bias")
 
             MODEL.t3.load_state_dict(cleaned_state_dict)