fix model path

app.py

@@ -5,6 +5,7 @@ import torch
 import argparse
 import librosa
 import soundfile as sf
+from huggingface_hub import snapshot_download
 
 from gpa_inference import GPAInference
 
@@ -61,8 +62,8 @@ def process_tts_a(text, ref_audio):
     # Preprocess audio
     ref_audio = preprocess_audio(ref_audio)
 
-    # Direct inference call
-    return inference.run_tts(
+    # Direct inference call - returns (sample_rate, audio_array)
+    result = inference.run_tts(
         task="tts-a",
         output_filename="tts_output.wav",
         text=text,
@@ -70,6 +71,8 @@ def process_tts_a(text, ref_audio):
         temperature=0.8,
         do_sample=True,
     )
+    # Return tuple format for Gradio Audio component
+    return result
 
 def process_vc(src_audio, ref_audio):
     global inference
@@ -83,12 +86,14 @@ def process_vc(src_audio, ref_audio):
     src_audio = preprocess_audio(src_audio)
     ref_audio = preprocess_audio(ref_audio)
 
-    # Direct inference call
-    return inference.run_vc(
+    # Direct inference call - returns (sample_rate, audio_array)
+    result = inference.run_vc(
         source_audio_path=src_audio,
         ref_audio_path=ref_audio,
         output_filename="vc_output.wav",
     )
+    # Return tuple format for Gradio Audio component
+    return result
 
 # ======================== Gradio UI Layout ========================
 
@@ -139,42 +144,40 @@ def parse_args():
 
     # Model Paths
     parser.add_argument(
-        "--tokenizer_path",
+        "--hf_model_id",
         type=str,
-        default="/data3/gpa_ckpt/gpa_final/glm-4-voice-tokenizer",
-        help="Path to GLM4 tokenizer",
+        default="AutoArk-AI/GPA",
+        help="Hugging Face model ID to download",
     )
     parser.add_argument(
-        "--text_tokenizer_path",
+        "--cache_dir",
         type=str,
-        default="/data3/gpa_ckpt/gpa_final",
-        help="Path to text tokenizer",
+        default="./models",
+        help="Directory to cache downloaded models",
     )
     parser.add_argument(
-        "--bicodec_tokenizer_path",
+        "--tokenizer_path",
         type=str,
-        default="/data3/gpa_ckpt/gpa_final/BiCodec/",
-        help="Path to BiCodec tokenizer",
+        default=None,
+        help="Path to GLM4 tokenizer (if None, will use downloaded model)",
     )
     parser.add_argument(
-        "--gpa_model_path",
+        "--text_tokenizer_path",
         type=str,
-        default="/data3/gpa_ckpt/gpa_final",
-        help="Path to GPA model",
+        default=None,
+        help="Path to text tokenizer (if None, will use downloaded model)",
     )
-
-    # System Config
     parser.add_argument(
-        "--output_dir",
+        "--bicodec_tokenizer_path",
         type=str,
-        default="./output_gui",
-        help="Directory to save output files",
+        default=None,
+        help="Path to BiCodec tokenizer (if None, will use downloaded model)",
     )
     parser.add_argument(
-        "--device",
+        "--gpa_model_path",
         type=str,
-        default="cuda" if torch.cuda.is_available() else "cpu",
-        help="Device to use",
+        default=None,
+        help="Path to GPA model (if None, will use downloaded model)",
     )
 
     # Server Config
@@ -189,18 +192,41 @@ def parse_args():
 
 args = parse_args()
 
+# Download model from Hugging Face Hub
+print(f"Downloading model from {args.hf_model_id}...")
+model_base_path = snapshot_download(
+    repo_id=args.hf_model_id,
+    cache_dir=args.cache_dir,
+    resume_download=True,
+)
+print(f"Model downloaded to: {model_base_path}")
+
+# Construct actual paths from downloaded model
+tokenizer_path = args.tokenizer_path or os.path.join(
+    model_base_path, "glm-4-voice-tokenizer"
+)
+text_tokenizer_path = args.text_tokenizer_path or model_base_path
+bicodec_tokenizer_path = args.bicodec_tokenizer_path or os.path.join(
+    model_base_path, "BiCodec"
+)
+gpa_model_path = args.gpa_model_path or model_base_path
+
 # Instantiate Model
 print(f"Initializing GPA Inference System on {args.device}...")
-os.makedirs(args.output_dir, exist_ok=True)
+print(f"Tokenizer path: {tokenizer_path}")
+print(f"Text tokenizer path: {text_tokenizer_path}")
+print(f"BiCodec tokenizer path: {bicodec_tokenizer_path}")
+print(f"GPA model path: {gpa_model_path}")
 
+# Use None for output_dir to enable temporary directory in HF Spaces
 inference = GPAInference(
-    tokenizer_path=args.tokenizer_path,
-    text_tokenizer_path=args.text_tokenizer_path,
-    bicodec_tokenizer_path=args.bicodec_tokenizer_path,
-    gpa_model_path=args.gpa_model_path,
-    output_dir=args.output_dir,
-    device=args.device,
+    tokenizer_path=tokenizer_path,
+    text_tokenizer_path=text_tokenizer_path,
+    bicodec_tokenizer_path=bicodec_tokenizer_path,
+    gpa_model_path=gpa_model_path,
+    output_dir=None,  # Will use temporary directory
+    device="cuda" if torch.cuda.is_available() else "cpu",
 )
 
 # Launch Gradio Demo
-demo.queue().launch()
+demo.queue().launch(server_name=args.server_name, server_port=args.server_port)

gpa_inference.py

@@ -3,6 +3,7 @@ import argparse
 import torch
 import soundfile as sf
 import re
+import tempfile
 from transformers import AutoTokenizer, AutoModelForCausalLM, WhisperFeatureExtractor
 import numpy as np
 
@@ -14,13 +15,31 @@ from models.glm_speech_tokenizer.modeling_whisper import WhisperVQEncoder
 
 from data_utils.audio_dataset_ark_audio import ark_infer_processor
 
+
 class GPAInference:
-    def __init__(self, tokenizer_path, text_tokenizer_path, bicodec_tokenizer_path, gpa_model_path, output_dir, device):
+
+    def __init__(
+        self,
+        tokenizer_path,
+        text_tokenizer_path,
+        bicodec_tokenizer_path,
+        gpa_model_path,
+        output_dir=None,
+        device=None,
+    ):
         self.tokenizer_path = tokenizer_path
         self.text_tokenizer_path = text_tokenizer_path
         self.bicodec_tokenizer_path = bicodec_tokenizer_path
         self.gpa_model_path = gpa_model_path
-        self.output_dir = output_dir
+        
+        # Use temporary directory if output_dir is None
+        if output_dir is None:
+            self.output_dir = tempfile.mkdtemp()
+            print(f"Using temporary output directory: {self.output_dir}")
+        else:
+            self.output_dir = output_dir
+            os.makedirs(self.output_dir, exist_ok=True)
+        
         self.device = device
 
         print(f"Using device: {self.device}")
@@ -29,15 +48,22 @@ class GPAInference:
     def _load_models(self):
         print("Loading tokenizers...")
         feature_extractor = WhisperFeatureExtractor.from_pretrained(self.tokenizer_path)
-        audio_model = WhisperVQEncoder.from_pretrained(self.tokenizer_path).eval().to(self.device)
-        self.glm_tokenizer = SpeechTokenExtractor(model=audio_model, feature_extractor=feature_extractor, device=self.device)
+        audio_model = (
+            WhisperVQEncoder.from_pretrained(self.tokenizer_path).eval().to(self.device)
+        )
+        self.glm_tokenizer = SpeechTokenExtractor(
+            model=audio_model, feature_extractor=feature_extractor, device=self.device
+        )
         self.text_tokenizer = AutoTokenizer.from_pretrained(
-            self.text_tokenizer_path, 
-            trust_remote_code=True
+            self.text_tokenizer_path, trust_remote_code=True
         )
 
-        self.bicodec_tokenizer = SparkTokenizer(model_path=self.bicodec_tokenizer_path, device=self.device)
-        self.bicodec_detokenizer = SparkDeTokenizer(model_path=self.bicodec_tokenizer_path, device=self.device)
+        self.bicodec_tokenizer = SparkTokenizer(
+            model_path=self.bicodec_tokenizer_path, device=self.device
+        )
+        self.bicodec_detokenizer = SparkDeTokenizer(
+            model_path=self.bicodec_tokenizer_path, device=self.device
+        )
         self.processor = ark_infer_processor(
             glm_tokenizer=self.glm_tokenizer,
             bicodec_tokenizer=self.bicodec_tokenizer,
@@ -48,8 +74,7 @@ class GPAInference:
 
         print("Loading model...")
         self.model = AutoModelForCausalLM.from_pretrained(
-            self.gpa_model_path, 
-            trust_remote_code=True
+            self.gpa_model_path, trust_remote_code=True
         ).to(self.device)
 
     def generate(self, inputs, **kwargs):
@@ -73,13 +98,15 @@ class GPAInference:
         generation_config.update(kwargs)
 
         # Remove keys that might be None if passed from args mistakenly
-        generation_config = {k: v for k, v in generation_config.items() if v is not None}
+        generation_config = {
+            k: v for k, v in generation_config.items() if v is not None
+        }
         print(f"Generation config: {generation_config}")
 
         outputs = self.model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
-            **generation_config
+            **generation_config,
         )
         return outputs
 
@@ -105,9 +132,13 @@ class GPAInference:
         text = self.text_tokenizer.decode(outputs[0].tolist())
 
         if "<|start_content|>" in text:
-            return text.split("<|start_content|>")[1].replace("<|im_end|>","").replace("<|end_content|>","")
+            return (
+                text.split("<|start_content|>")[1]
+                .replace("<|im_end|>", "")
+                .replace("<|end_content|>", "")
+            )
         else:
-            return text.replace("<|im_end|>","")
+            return text.replace("<|im_end|>", "")
 
     def run_tts(self, task, output_filename, text, ref_audio_path, **kwargs):
         """
@@ -129,12 +160,11 @@ class GPAInference:
         }
 
         print(f"\n--- {task.upper()} ---")
-        output_path = os.path.join(self.output_dir, output_filename)
 
         # Pass processor specific args (e.g. emotion, pitch) here
         inputs = self.processor.process_input(
-            task=task, 
-            ref_audio_path=ref_audio_path, 
+            task=task,
+            ref_audio_path=ref_audio_path,
             text=text,
         )
 
@@ -154,7 +184,9 @@ class GPAInference:
         audio_list = [int(x) for x in audio_ids]
 
         if ref_audio_path:
-            global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])['global_tokens']
+            global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])[
+                "global_tokens"
+            ]
         else:
             global_tokens = torch.zeros((1, 32), dtype=torch.long).to(self.device)
 
@@ -168,6 +200,7 @@ class GPAInference:
         if reconstructed_wav.size > 0:
             reconstructed_wav -= reconstructed_wav.mean()
 
+        output_path = os.path.join(self.output_dir, output_filename)
         sf.write(output_path, reconstructed_wav, 16000)
         print(f"Saved output to {output_path}")
         return 16000, reconstructed_wav
@@ -204,7 +237,9 @@ class GPAInference:
         audio_ids = re.findall(r"<\|bicodec_semantic_(\d+)\|>", content)
         audio_list = [int(x) for x in audio_ids]
 
-        global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])['global_tokens']
+        global_tokens = self.bicodec_tokenizer.tokenize([ref_audio_path])[
+            "global_tokens"
+        ]
 
         req = {
             "global_tokens": global_tokens,
@@ -224,10 +259,30 @@ def parse_args():
     parser = argparse.ArgumentParser(description="GPA Inference Script")
 
     # Paths
-    parser.add_argument("--tokenizer_path", type=str, default="/nasdata/model/gpa/glm-4-voice-tokenizer", help="Path to GLM4 tokenizer")
-    parser.add_argument("--text_tokenizer_path", type=str, default="/nasdata/model/gpa", help="Path to text tokenizer")
-    parser.add_argument("--bicodec_tokenizer_path", type=str, default="/nasdata/model/gpa/BiCodec/", help="Path to BiCodec tokenizer")
-    parser.add_argument("--gpa_model_path", type=str, default="/nasdata/model/gpa", help="Path to GPA model")
+    parser.add_argument(
+        "--tokenizer_path",
+        type=str,
+        default="/nasdata/model/gpa/glm-4-voice-tokenizer",
+        help="Path to GLM4 tokenizer",
+    )
+    parser.add_argument(
+        "--text_tokenizer_path",
+        type=str,
+        default="/nasdata/model/gpa",
+        help="Path to text tokenizer",
+    )
+    parser.add_argument(
+        "--bicodec_tokenizer_path",
+        type=str,
+        default="/nasdata/model/gpa/BiCodec/",
+        help="Path to BiCodec tokenizer",
+    )
+    parser.add_argument(
+        "--gpa_model_path",
+        type=str,
+        default="/nasdata/model/gpa",
+        help="Path to GPA model",
+    )
 
     # Audio inputs
     parser.add_argument(
@@ -238,20 +293,34 @@ def parse_args():
     )
 
     # Output
-    parser.add_argument("--output_dir", type=str, default=".", help="Directory to save output files")
+    parser.add_argument(
+        "--output_dir", type=str, default=".", help="Directory to save output files"
+    )
 
     # Device
     default_device = "cuda" if torch.cuda.is_available() else "cpu"
-    parser.add_argument("--device", type=str, default=default_device, help="Device to use (e.g., cuda:0, cpu)")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default=default_device,
+        help="Device to use (e.g., cuda:0, cpu)",
+    )
 
     # Task
-    parser.add_argument("--task", type=str, required=True, choices=["stt", "tts-a", "vc"], help="Task to run")
+    parser.add_argument(
+        "--task",
+        type=str,
+        required=True,
+        choices=["stt", "tts-a", "vc"],
+        help="Task to run",
+    )
 
     # TTS Inputs (Processor Arguments)
     parser.add_argument("--text", type=str, default=None, help="Text for TTS")
 
     return parser.parse_args()
 
+
 def main():
     args = parse_args()
 
@@ -289,5 +358,6 @@ def main():
             output_filename="output_gpa_vc.wav",
         )
 
+
 if __name__ == "__main__":
     main()