Update app.py

app.py

@@ -8,7 +8,7 @@ from huggingface_hub import hf_hub_download
 # from modules.v2.vc_wrapper import VoiceConversionWrapper 
 
 # --- CONFIGURATION (UPDATE YOUR_USERNAME HERE) ---
-# Your correct model repository ID
+# Your correct model repository ID for automatic download in the Space
 MODEL_REPO_ID = "Bajiyo/dhanush_seedvc"
 CFM_FILE = "CFM_epoch_00651_step_21500.pth"
 AR_FILE = "AR_epoch_00651_step_21500.pth"
@@ -25,32 +25,44 @@ dtype = torch.float16
 
 def load_models(args):
     """
-    Loads models, handling checkpoint download from Hugging Face Hub.
+    Loads models, prioritizing command-line arguments for local paths,
+    and falling back to Hugging Face Hub download for the Space environment.
     """
-    # 1. Setup local directory and download checkpoints
-    LOCAL_CHECKPOINTS_DIR = "downloaded_checkpoints"
-    os.makedirs(LOCAL_CHECKPOINTS_DIR, exist_ok=True)
-    print(f"Downloading checkpoints from {MODEL_REPO_ID}...")
-
-    # Download CFM
-    cfm_local_path = hf_hub_download(
-        repo_id=MODEL_REPO_ID, 
-        filename=CFM_FILE, 
-        local_dir=LOCAL_CHECKPOINTS_DIR,
-        local_dir_use_symlinks=False
-    )
-    print(f"CFM checkpoint downloaded to: {cfm_local_path}")
     
-    # Download AR
-    ar_local_path = hf_hub_download(
-        repo_id=MODEL_REPO_ID, 
-        filename=AR_FILE, 
-        local_dir=LOCAL_CHECKPOINTS_DIR,
-        local_dir_use_symlinks=False
-    )
-    print(f"AR checkpoint downloaded to: {ar_local_path}")
+    # --- 1. Determine Checkpoint Paths ---
+    if args.cfm_checkpoint_path:
+        cfm_local_path = args.cfm_checkpoint_path
+        print(f"Using local CFM checkpoint path from arguments: {cfm_local_path}")
+    else:
+        # Default behavior for Space: download from HF
+        LOCAL_CHECKPOINTS_DIR = "downloaded_checkpoints"
+        os.makedirs(LOCAL_CHECKPOINTS_DIR, exist_ok=True)
+        print(f"Arguments not provided. Downloading CFM checkpoint from {MODEL_REPO_ID}...")
+        cfm_local_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID, 
+            filename=CFM_FILE, 
+            local_dir=LOCAL_CHECKPOINTS_DIR,
+            local_dir_use_symlinks=False
+        )
+        print(f"CFM checkpoint downloaded to: {cfm_local_path}")
+
+    if args.ar_checkpoint_path:
+        ar_local_path = args.ar_checkpoint_path
+        print(f"Using local AR checkpoint path from arguments: {ar_local_path}")
+    else:
+        # Default behavior for Space: download from HF
+        LOCAL_CHECKPOINTS_DIR = "downloaded_checkpoints"
+        os.makedirs(LOCAL_CHECKPOINTS_DIR, exist_ok=True) # Ensure dir exists
+        print(f"Arguments not provided. Downloading AR checkpoint from {MODEL_REPO_ID}...")
+        ar_local_path = hf_hub_download(
+            repo_id=MODEL_REPO_ID, 
+            filename=AR_FILE, 
+            local_dir=LOCAL_CHECKPOINTS_DIR,
+            local_dir_use_symlinks=False
+        )
+        print(f"AR checkpoint downloaded to: {ar_local_path}")
 
-    # 2. Instantiate and load models
+    # --- 2. Instantiate and load models ---
     from hydra.utils import instantiate
     from omegaconf import DictConfig
     
@@ -58,7 +70,7 @@ def load_models(args):
     cfg = DictConfig(yaml.safe_load(open("configs/v2/vc_wrapper.yaml", "r")))
     vc_wrapper = instantiate(cfg)
     
-    # Load the downloaded checkpoints
+    # Load the determined checkpoints (either local paths or downloaded HF paths)
     vc_wrapper.load_checkpoints(
         ar_checkpoint_path=ar_local_path,
         cfm_checkpoint_path=cfm_local_path
@@ -85,6 +97,7 @@ def main(args):
     vc_wrapper = load_models(args)
     
     # Define wrapper function for Gradio. NO DECORATORS HERE.
+    # This wrapper ensures the streaming output works correctly in the Gradio Interface.
     def convert_voice_wrapper(source_audio_path, target_audio_path, diffusion_steps,
                               length_adjust, intelligibility_cfg_rate, similarity_cfg_rate,
                               top_p, temperature, repetition_penalty, convert_style, 
@@ -149,7 +162,7 @@ def main(args):
     
     # Launch the Gradio interface
     gr.Interface(
-        fn=convert_voice_wrapper,
+        fn=convert_voice_wrapper, # Using the wrapper for reliable streaming
         description=description,
         inputs=inputs,
         outputs=outputs,
@@ -162,10 +175,10 @@ if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser()
     parser.add_argument("--compile", action="store_true", help="Compile the model using torch.compile")
-    # These arguments are now effectively ignored/not needed since we download the models
+    # These are the arguments that allow you to run the script locally with specific paths
     parser.add_argument("--ar-checkpoint-path", type=str, default=None,
-                        help="Path to custom checkpoint file (overridden by HF download in Space)")
+                        help="Path to custom AR checkpoint file. Defaults to HF download in Space.")
     parser.add_argument("--cfm-checkpoint-path", type=str, default=None,
-                        help="Path to custom checkpoint file (overridden by HF download in Space)")
+                        help="Path to custom CFM checkpoint file. Defaults to HF download in Space.")
     args = parser.parse_args()
     main(args)