Update app.py

app.py

@@ -2,14 +2,31 @@ import gradio as gr
 import torch
 import torchaudio
 import os
+import re
 import subprocess
 from transformers import AutoModelForCausalLM
-from outetts.wav_tokenizer.decoder import WavTokenizer
-from yarngpt.audiotokenizer import AudioTokenizer
+from yarngpt_utils import AudioTokenizer
+
+# Download model files if they don't exist
+def download_if_not_exists(url, filename):
+    if not os.path.exists(filename):
+        print(f"Downloading {filename}...")
+        subprocess.run(["wget", url, "-O", filename])
+        print(f"Downloaded {filename}")
+
+# Download necessary files
+download_if_not_exists(
+    "https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml",
+    "wavtokenizer_config.yaml"
+)
+download_if_not_exists(
+    "https://huggingface.co/novateur/WavTokenizer-large-speech-75token/blob/main/wavtokenizer_large_speech_320_v2.ckpt",
+    "wavtokenizer_model.ckpt"
+)
 
 # Initialize the model (this runs when the app starts)
 def initialize_model():
-    # Download model and tokenizer
+    # Set paths
     hf_path = "saheedniyi/YarnGPT"
     wav_tokenizer_config_path = "wavtokenizer_config.yaml"
     wav_tokenizer_model_path = "wavtokenizer_model.ckpt"
@@ -24,22 +41,6 @@ def initialize_model():
     
     return model, audio_tokenizer
 
-    def download_if_not_exists(url, filename):
-    if not os.path.exists(filename):
-        print(f"Downloading {filename}...")
-        subprocess.run(["wget", url, "-O", filename])
-        print(f"Downloaded {filename}")
-
-# Download the model files if they're not already present
-download_if_not_exists(
-    "https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml",
-    "wavtokenizer_config.yaml"
-)
-download_if_not_exists(
-    "https://huggingface.co/novateur/WavTokenizer-large-speech-75token/blob/main/wavtokenizer_large_speech_320_v2.ckpt",
-    "wavtokenizer_model.ckpt"
-)
-
 # Generate audio from text
 def generate_speech(text, speaker_name):
     # Create prompt