Update download_model.py

download_model.py

@@ -1,147 +1,135 @@
 from huggingface_hub import list_repo_files, hf_hub_download
 import os
 import shutil
+import torch
+from itertools import combinations
+import platform
 
 # Repository ID
 repo_id = "hexgrad/Kokoro-82M"
 
 # Set up the cache directory
-cache_dir = "./cache"  # Customize this path if needed
+cache_dir = "./cache"
 os.makedirs(cache_dir, exist_ok=True)
 
+# Set up the base model paths
+KOKORO_DIR = "./KOKORO"
+VOICES_DIR = os.path.join(KOKORO_DIR, "voices")
+FP16_DIR = os.path.join(KOKORO_DIR, "fp16")
+KOKORO_FILE = "kokoro-v0_19.pth"
+FP16_FILE = "fp16/kokoro-v0_19-half.pth"
+
+def download_files(repo_id, filenames, destination_dir, cache_dir):
+     # Ensure directories exist
+     os.makedirs(destination_dir, exist_ok=True)
+     
+     for filename in filenames:
+        destination = os.path.join(destination_dir, os.path.basename(filename))
+        if not os.path.exists(destination):
+            file_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
+            shutil.copy(file_path, destination)
+            print(f"Downloaded and saved: {destination}")
+        else:
+              print(f"File already exist in: {destination}")
+        
+
 def get_voice_models():
-    # Ensure the 'voices' directory exists
-    voices_dir = './KOKORO/voices'
+    """Downloads missing voice models from the Hugging Face repository."""
+
+    # Create or empty the 'voices' directory
+    if os.path.exists(VOICES_DIR):
+       shutil.rmtree(VOICES_DIR)
     os.makedirs(VOICES_DIR, exist_ok=True)
-    # Get the list of all files
-    files = list_repo_files(repo_id)
 
-    # Filter files for the 'voices/' folder
+    # Get list of files from the repository
+    files = list_repo_files(repo_id)
+    
+    # Filter for voice files
     voice_files = [file.replace("voices/", "") for file in files if file.startswith("voices/")]
 
-    # Get current files in the 'voices' folder
-    current_voice = os.listdir(voices_dir)
+    # Get current voice files
+    current_voice = os.listdir(VOICES_DIR)
 
-    # Identify files that need to be downloaded
+    # Download new voices
     download_voice = [file for file in voice_files if file not in current_voice]
     if download_voice:
-        print(f"Files to download: {download_voice}")
-
-    # Download each missing file
-    for file in download_voice:
-        file_path = hf_hub_download(repo_id=repo_id, filename=f"voices/{file}", cache_dir=cache_dir)
-        target_path = os.path.join(voices_dir, file)
-        shutil.copy(file_path, target_path)
-        print(f"Downloaded: {file} to {target_path}")
-
-# Call the function to execute the code
-get_voice_models()
-
-# Check and download additional required files with caching
-kokoro_file = "kokoro-v0_19.pth"
-fp16_file = "fp16/kokoro-v0_19-half.pth"
-
-if kokoro_file not in os.listdir("./KOKORO/"):
-    file_path = hf_hub_download(repo_id=repo_id, filename=kokoro_file, cache_dir=cache_dir)
-    shutil.copy(file_path, os.path.join("./KOKORO/", kokoro_file))
-    print(f"Downloaded: {kokoro_file} to ./KOKORO/")
-
-if "fp16" not in os.listdir("./KOKORO/"):
-    os.makedirs("./KOKORO/fp16", exist_ok=True)
-
-if os.path.basename(fp16_file) not in os.listdir("./KOKORO/fp16/"):
-    file_path = hf_hub_download(repo_id=repo_id, filename=fp16_file, cache_dir=cache_dir)
-    shutil.copy(file_path, os.path.join("./KOKORO/fp16/", os.path.basename(fp16_file)))
-    print(f"Downloaded: {os.path.basename(fp16_file)} to ./KOKORO/fp16/")
-
-
+        #  print(f"Files to download: {download_voice}")
+         pass
+    
+    download_files(repo_id, [f"voices/{file}" for file in download_voice], VOICES_DIR, cache_dir)
 
+def download_base_models():
+    """Downloads Kokoro base model and fp16 version if missing."""
 
-#For Windows one click run
-import os
-import platform
+    download_files(repo_id, [KOKORO_FILE], KOKORO_DIR, cache_dir)
+    download_files(repo_id, [FP16_FILE], FP16_DIR, cache_dir)
 
 def setup_batch_file():
-    # Check if the system is Windows
+    """Creates a 'run_app.bat' file for Windows if it doesn't exist."""
+
     if platform.system() == "Windows":
-        # Check if 'run.bat' exists in the current folder
-        if os.path.exists("run.bat"):
-            print("'run.bat' already exists in the current folder.")
-        else:
-            # Content for run_app.bat
+        bat_file_name = 'run_app.bat'
+        if not os.path.exists(bat_file_name):
             bat_content_app = '''@echo off
 call myenv\\Scripts\\activate
 @python.exe app.py %*
 @pause
 '''
-            # Save the content to run_app.bat
-            with open('run_app.bat', 'w') as bat_file:
+            with open(bat_file_name, 'w') as bat_file:
                 bat_file.write(bat_content_app)
-            print("The 'run_app.bat' file has been created.")
+            print(f"Created '{bat_file_name}'.")
+        else:
+            print(f"'{bat_file_name}' already exists.")
     else:
-        print("This system is not Windows. Batch file creation skipped.")
-
-# Run the setup function
-setup_batch_file()
-
-
-
-
-import torch
-import os
-from itertools import combinations
+        print("Not a Windows system, skipping batch file creation.")
+
+def download_ffmpeg():
+    """Downloads ffmpeg and ffprobe executables from Hugging Face."""
+    print("For Kokoro TTS we don't need ffmpeg, But for Subtitle Dubbing we need ffmpeg")
+    os_name=platform.system()
+    if os_name == "Windows":
+        repo_id = "fishaudio/fish-speech-1"
+        filenames = ["ffmpeg.exe", "ffprobe.exe"]
+        ffmpeg_dir = "./ffmpeg"
+        download_files(repo_id, filenames, ffmpeg_dir, cache_dir)
+    elif os_name == "Linux":
+         print("Please install ffmpeg using the package manager for your system.")
+         print("'sudo apt install ffmpeg' on Debian/Ubuntu")
+    else:
+        print(f"Manually install ffmpeg for {os_name} from https://ffmpeg.org/download.html")
 
-def mix_all_voices(folder_path="./KOKORO/voices"):
-    """Mix all pairs of voice models and save the new models."""
+def mix_all_voices(folder_path=VOICES_DIR):
+     """Mix all pairs of voice models and save the new models."""
     # Get the list of available voice packs
-    available_voice_pack = [
+     available_voice_pack = [
         os.path.splitext(filename)[0]
         for filename in os.listdir(folder_path)
         if filename.endswith('.pt')
     ]
 
-    # Generate all unique pairs of voices
-    voice_combinations = combinations(available_voice_pack, 2)
-
-    # def mix_model(voice_1, voice_2, weight_1=0.6, weight_2=0.4):
-    #     """Mix two voice models with a weighted average and save the new model."""
-    #     new_name = f"{voice_1}_mix_{voice_2}"
-    #     voice_id_1 = torch.load(f'{folder_path}/{voice_1}.pt', weights_only=True)
-    #     voice_id_2 = torch.load(f'{folder_path}/{voice_2}.pt', weights_only=True)
-        
-    #     # Create the mixed model using a weighted average
-    #     mixed_voice = (weight_1 * voice_id_1) + (weight_2 * voice_id_2)
-        
-    #     # Save the mixed model
-    #     torch.save(mixed_voice, f'{folder_path}/{new_name}.pt')
-    #     print(f"Created new voice model: {new_name}")
-    
-
-    
-    # Function to mix two voices
-    def mix_model(voice_1, voice_2):
-        """Mix two voice models and save the new model."""
-        new_name = f"{voice_1}_mix_{voice_2}"
-        voice_id_1 = torch.load(f'{folder_path}/{voice_1}.pt', weights_only=True)
-        voice_id_2 = torch.load(f'{folder_path}/{voice_2}.pt', weights_only=True)
+     # Generate all unique pairs of voices
+     voice_combinations = combinations(available_voice_pack, 2)
+   
+     # Function to mix two voices
+     def mix_model(voice_1, voice_2):
+          """Mix two voice models and save the new model."""
+          new_name = f"{voice_1}_mix_{voice_2}"
+          voice_id_1 = torch.load(f'{folder_path}/{voice_1}.pt', weights_only=True)
+          voice_id_2 = torch.load(f'{folder_path}/{voice_2}.pt', weights_only=True)
 
-        # Create the mixed model by averaging the weights
-        mixed_voice = torch.mean(torch.stack([voice_id_1, voice_id_2]), dim=0)
+          # Create the mixed model by averaging the weights
+          mixed_voice = torch.mean(torch.stack([voice_id_1, voice_id_2]), dim=0)
 
-        # Save the mixed model
-        torch.save(mixed_voice, f'{folder_path}/{new_name}.pt')
-        print(f"Created new voice model: {new_name}")
+          # Save the mixed model
+          torch.save(mixed_voice, f'{folder_path}/{new_name}.pt')
+          print(f"Created new voice model: {new_name}")
 
     # Create mixed voices for each pair
-    for voice_1, voice_2 in voice_combinations:
-        print(f"Mixing {voice_1} ❤️ {voice_2}")
-        mix_model(voice_1, voice_2)
-
-# Call the function to mix all voices
-mix_all_voices("./KOKORO/voices")
+     for voice_1, voice_2 in voice_combinations:
+          print(f"Mixing {voice_1} ❤️ {voice_2}")
+          mix_model(voice_1, voice_2)
 
-
-def save_voice_names(directory="./KOKORO/voices", output_file="./voice_names.txt"):
+def save_voice_names(directory=VOICES_DIR, output_file="./voice_names.txt"):
     """
     Retrieves voice names from a directory, sorts them by length, and saves to a file.
 
@@ -168,4 +156,13 @@ def save_voice_names(directory="./KOKORO/voices", output_file="./voice_names.txt
             f.write(f"{voice_name}\n")
 
     print(f"Voice names saved to {output_file}")
-save_voice_names()
+
+# --- Main Execution ---
+if __name__ == "__main__":
+    get_voice_models()
+    download_base_models()
+    setup_batch_file()
+    # mix_all_voices()
+    save_voice_names()
+    download_ffmpeg()
+    print("Setup complete!")