import subprocess import os import sys import time import argparse import whisper # For Whisper model download/load from huggingface_hub import snapshot_download # For HF model downloads # --- Diagnostics at the start --- print("--- PYTHON ENVIRONMENT DIAGNOSTICS ---") print(f"sys.executable: {sys.executable}") print(f"sys.path: {sys.path}") print(f"PYTHONPATH: {os.environ.get('PYTHONPATH', 'Not set')}") print(f"Current working directory: {os.getcwd()}") print("--- END DIAGNOSTICS ---") # Ensure key paths are in sys.path if '/code' not in sys.path: sys.path.insert(0, '/code') # Look for llama_omni2 in various locations llama_omni2_locations = [ '/code/llama_omni2', # Direct copy in /code '/tmp/LLaMA-Omni2/llama_omni2', # Temp location during build '/code/LLaMA-Omni2/llama_omni2', # Old location '/usr/local/lib/python3.10/site-packages/llama_omni2', # Site-packages '/code/llama_omni2_extracted/llama_omni2' # Extracted scripts directory ] LLAMA_OMNI2_MODULE_PATH = None for location in llama_omni2_locations: if os.path.exists(location): print(f"Found llama_omni2 at: {location}") if location not in sys.path: sys.path.insert(0, os.path.dirname(location)) LLAMA_OMNI2_MODULE_PATH = location break if not LLAMA_OMNI2_MODULE_PATH: print("WARNING: Could not find llama_omni2 in any expected location!") # Let's search the entire filesystem for llama_omni2 print("Searching for llama_omni2 in system paths...") for path in sys.path: potential_path = os.path.join(path, 'llama_omni2') if os.path.exists(potential_path): print(f"Found llama_omni2 in sys.path at: {potential_path}") LLAMA_OMNI2_MODULE_PATH = potential_path break # Try to import llama_omni2 to confirm it's available try: import llama_omni2 print(f"Successfully imported llama_omni2 from {llama_omni2.__file__}") except ImportError as e: print(f"WARNING: Failed to import llama_omni2: {e}") print("Will attempt to use direct file paths instead.") # --- Configuration & Model Definitions --- # Main LLaMA-Omni2 model (User can change this to other LLaMA-Omni2 variants) # Options from LLaMA-Omni2 README: LLaMA-Omni2-0.5B, LLaMA-Omni2-1.5B, LLaMA-Omni2-3B, # LLaMA-Omni2-7B, LLaMA-Omni2-14B (English only) # LLaMA-Omni2-0.5B-Bilingual, LLaMA-Omni2-1.5B-Bilingual, LLaMA-Omni2-3B-Bilingual, # LLaMA-Omni2-7B-Bilingual, LLaMA-Omni2-14B-Bilingual, LLaMA-Omni2-32B-Bilingual (English & Chinese) LLAMA_OMNI2_MODEL_NAME = "LLaMA-Omni2-7B-Bilingual" # Example, choose as needed LLAMA_OMNI2_HF_REPO = f"ICTNLP/{LLAMA_OMNI2_MODEL_NAME}" LLAMA_OMNI2_MODEL_PATH_LOCAL = f"models/{LLAMA_OMNI2_MODEL_NAME}" # Local path where the model will be downloaded # Speech Encoder (Whisper) SPEECH_ENCODER_PATH = "models/speech_encoder" # Local path for Whisper # Vocoder (CosyVoice 2) COSYVOICE_HF_REPO = "ICTNLP/cosy2_decoder" COSYVOICE_PATH_LOCAL = "models/cosy2_decoder" # Local path for CosyVoice # Define paths to LLaMA-Omni2 scripts CONTROLLER_SCRIPT = os.path.join(LLAMA_OMNI2_MODULE_PATH, 'serve', 'controller.py') if LLAMA_OMNI2_MODULE_PATH else None MODEL_WORKER_SCRIPT = os.path.join(LLAMA_OMNI2_MODULE_PATH, 'serve', 'model_worker.py') if LLAMA_OMNI2_MODULE_PATH else None GRADIO_SERVER_SCRIPT = os.path.join(LLAMA_OMNI2_MODULE_PATH, 'serve', 'gradio_web_server.py') if LLAMA_OMNI2_MODULE_PATH else None # Print script paths to verify print(f"Controller script path: {CONTROLLER_SCRIPT}") print(f"Model worker script path: {MODEL_WORKER_SCRIPT}") print(f"Gradio server script path: {GRADIO_SERVER_SCRIPT}") # -- Check if scripts exist -- if CONTROLLER_SCRIPT and not os.path.exists(CONTROLLER_SCRIPT): print(f"WARNING: Controller script not found at {CONTROLLER_SCRIPT}") CONTROLLER_SCRIPT = None if MODEL_WORKER_SCRIPT and not os.path.exists(MODEL_WORKER_SCRIPT): print(f"WARNING: Model worker script not found at {MODEL_WORKER_SCRIPT}") MODEL_WORKER_SCRIPT = None if GRADIO_SERVER_SCRIPT and not os.path.exists(GRADIO_SERVER_SCRIPT): print(f"WARNING: Gradio server script not found at {GRADIO_SERVER_SCRIPT}") GRADIO_SERVER_SCRIPT = None # --- Download Models --- def download_all_models(): print("Starting model downloads...") # 1. Download/Load Whisper-large-v3 model print(f"Downloading Whisper-large-v3 to {SPEECH_ENCODER_PATH}...") os.makedirs(SPEECH_ENCODER_PATH, exist_ok=True) try: whisper.load_model("large-v3", download_root=SPEECH_ENCODER_PATH) print("Whisper model downloaded/loaded.") except Exception as e: print(f"Error downloading Whisper model: {e}") # 2. Download CosyVoice 2 decoder print(f"Downloading CosyVoice 2 decoder from {COSYVOICE_HF_REPO} to {COSYVOICE_PATH_LOCAL}...") os.makedirs(COSYVOICE_PATH_LOCAL, exist_ok=True) try: snapshot_download(repo_id=COSYVOICE_HF_REPO, local_dir=COSYVOICE_PATH_LOCAL, local_dir_use_symlinks=False, # Recommended for Spaces/Docker resume_download=True) print("CosyVoice 2 decoder downloaded.") except Exception as e: print(f"Error downloading CosyVoice 2 decoder: {e}") # 3. Download LLaMA-Omni2 main model print(f"Downloading LLaMA-Omni2 model ({LLAMA_OMNI2_MODEL_NAME}) from {LLAMA_OMNI2_HF_REPO} to {LLAMA_OMNI2_MODEL_PATH_LOCAL}...") os.makedirs(LLAMA_OMNI2_MODEL_PATH_LOCAL, exist_ok=True) try: snapshot_download(repo_id=LLAMA_OMNI2_HF_REPO, local_dir=LLAMA_OMNI2_MODEL_PATH_LOCAL, local_dir_use_symlinks=False, resume_download=True) print(f"LLaMA-Omni2 model ({LLAMA_OMNI2_MODEL_NAME}) downloaded.") except Exception as e: print(f"Error downloading LLaMA-Omni2 model ({LLAMA_OMNI2_MODEL_NAME}): {e}") print("All model download attempts finished.") # Search for scripts if they weren't found initially def find_scripts_in_filesystem(): if not any([CONTROLLER_SCRIPT, MODEL_WORKER_SCRIPT, GRADIO_SERVER_SCRIPT]): print("Searching for LLaMA-Omni2 scripts in key locations...") controller_paths = [] model_worker_paths = [] gradio_server_paths = [] # Only search in likely locations to avoid a slow full filesystem scan search_dirs = [ '/code', '/tmp/LLaMA-Omni2', '/usr/local/lib/python3.10/site-packages', '/home/user' ] for search_dir in search_dirs: if not os.path.exists(search_dir): print(f"Skipping {search_dir} - directory doesn't exist") continue print(f"Searching in {search_dir}...") for root, dirs, files in os.walk(search_dir): # Skip .git and node_modules directories to speed up search dirs[:] = [d for d in dirs if d not in ('.git', 'node_modules')] for file in files: if file == 'controller.py': controller_paths.append(os.path.join(root, file)) elif file == 'model_worker.py': model_worker_paths.append(os.path.join(root, file)) elif file == 'gradio_web_server.py': gradio_server_paths.append(os.path.join(root, file)) print(f"Found controller.py at: {controller_paths}") print(f"Found model_worker.py at: {model_worker_paths}") print(f"Found gradio_web_server.py at: {gradio_server_paths}") return controller_paths[0] if controller_paths else None, \ model_worker_paths[0] if model_worker_paths else None, \ gradio_server_paths[0] if gradio_server_paths else None return CONTROLLER_SCRIPT, MODEL_WORKER_SCRIPT, GRADIO_SERVER_SCRIPT # --- Main Execution --- if __name__ == "__main__": # Perform downloads before starting servers # Simple check to avoid re-downloading if key directories exist (can be made more robust) if not os.path.exists(os.path.join(SPEECH_ENCODER_PATH, "large-v3.pt")) or \ not os.path.exists(COSYVOICE_PATH_LOCAL) or \ not os.path.exists(LLAMA_OMNI2_MODEL_PATH_LOCAL): download_all_models() else: print("Model directories seem to exist. Skipping downloads. Delete them to force re-download.") # Use the direct launcher to run LLaMA-Omni2 components print("Starting LLaMA-Omni2 with direct launcher...") launcher_script = os.path.join(os.getcwd(), "launch_llama_omni2.py") if os.path.exists(launcher_script): try: # Make the launcher script executable os.chmod(launcher_script, 0o755) # Run the launcher script subprocess.run([sys.executable, launcher_script], check=True) except subprocess.CalledProcessError as e: print(f"Error running launcher script: {e}") except Exception as e: print(f"Unexpected error: {e}") else: print(f"Error: Launcher script not found at {launcher_script}") print("Please ensure that launch_llama_omni2.py exists in the current directory.")