import gradio as gr import subprocess import os import shutil import tempfile """ # Set the PATH and LD_LIBRARY_PATH for CUDA 12.3 cuda_bin_path = "/usr/local/cuda/bin" cuda_lib_path = "/usr/local/cuda/lib64" # Update the environment variables os.environ['PATH'] = f"{cuda_bin_path}:{os.environ.get('PATH', '')}" os.environ['LD_LIBRARY_PATH'] = f"{cuda_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}" """ # Install required package def install_flash_attn(): try: print("Installing flash-attn...") subprocess.run( ["pip", "install", "flash-attn", "--no-build-isolation"], check=True ) print("flash-attn installed successfully!") except subprocess.CalledProcessError as e: print(f"Failed to install flash-attn: {e}") exit(1) # Install flash-attn install_flash_attn() from huggingface_hub import snapshot_download # Create xcodec_mini_infer folder folder_path = './inference/xcodec_mini_infer' # Create the folder if it doesn't exist if not os.path.exists(folder_path): os.mkdir(folder_path) print(f"Folder created at: {folder_path}") else: print(f"Folder already exists at: {folder_path}") snapshot_download( repo_id = "m-a-p/xcodec_mini_infer", local_dir = "./inference/xcodec_mini_infer" ) # Change to the "inference" directory inference_dir = "./inference" try: os.chdir(inference_dir) print(f"Changed working directory to: {os.getcwd()}") except FileNotFoundError: print(f"Directory not found: {inference_dir}") exit(1) def empty_output_folder(output_dir): # List all files in the output directory files = os.listdir(output_dir) # Iterate over the files and remove them for file in files: file_path = os.path.join(output_dir, file) try: if os.path.isdir(file_path): # If it's a directory, remove it recursively shutil.rmtree(file_path) else: # If it's a file, delete it os.remove(file_path) except Exception as e: print(f"Error deleting file {file_path}: {e}") # Function to create a temporary file with string content def create_temp_file(content, prefix, suffix=".txt"): temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix) # Ensure content ends with newline and normalize line endings content = content.strip() + "\n\n" # Add extra newline at end content = content.replace("\r\n", "\n").replace("\r", "\n") temp_file.write(content) temp_file.close() # Debug: Print file contents print(f"\nContent written to {prefix}{suffix}:") print(content) print("---") return temp_file.name def get_last_mp3_file(output_dir): # List all files in the output directory files = os.listdir(output_dir) # Filter only .mp3 files mp3_files = [file for file in files if file.endswith('.mp3')] if not mp3_files: print("No .mp3 files found in the output folder.") return None # Get the full path for the mp3 files mp3_files_with_path = [os.path.join(output_dir, file) for file in mp3_files] # Sort the files based on the modification time (most recent first) mp3_files_with_path.sort(key=lambda x: os.path.getmtime(x), reverse=True) # Return the most recent .mp3 file return mp3_files_with_path[0] def infer(genre_txt_content, lyrics_txt_content): # Create temporary files genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_") lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_") print(f"Genre TXT path: {genre_txt_path}") print(f"Lyrics TXT path: {lyrics_txt_path}") # Ensure the output folder exists output_dir = "./output" os.makedirs(output_dir, exist_ok=True) print(f"Output folder ensured at: {output_dir}") empty_output_folder(output_dir) # Command and arguments with optimized settings command = [ "python", "infer.py", "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot", "--stage2_model", "m-a-p/YuE-s2-1B-general", "--genre_txt", f"{genre_txt_path}", "--lyrics_txt", f"{lyrics_txt_path}", "--run_n_segments", "2", "--stage2_batch_size", "8", # Increased from 4 to 8 "--output_dir", f"{output_dir}", "--cuda_idx", "0", "--max_new_tokens", "3000", "--disable_offload_model" ] # Set up environment variables for CUDA with optimized settings env = os.environ.copy() env.update({ "CUDA_VISIBLE_DEVICES": "0", "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512", "CUDA_HOME": "/usr/local/cuda", "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}", "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}", "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512,garbage_collection_threshold:0.8", # Added garbage collection threshold "TORCH_DISTRIBUTED_DEBUG": "DETAIL", # Added for better debugging "CUDA_LAUNCH_BLOCKING": "0" # Ensure asynchronous CUDA operations }) # Execute the command try: subprocess.run(command, check=True, env=env) print("Command executed successfully!") # Check and print the contents of the output folder output_files = os.listdir(output_dir) if output_files: print("Output folder contents:") for file in output_files: print(f"- {file}") last_mp3 = get_last_mp3_file(output_dir) if last_mp3: print("Last .mp3 file:", last_mp3) return last_mp3 else: return None else: print("Output folder is empty.") return None except subprocess.CalledProcessError as e: print(f"Error occurred: {e}") return None finally: # Clean up temporary files os.remove(genre_txt_path) os.remove(lyrics_txt_path) print("Temporary files deleted.") # Gradio with gr.Blocks() as demo: with gr.Column(): gr.Markdown("# YuE") with gr.Row(): with gr.Column(): genre_txt = gr.Textbox(label="Genre") lyrics_txt = gr.Textbox(label="Lyrics") submit_btn = gr.Button("Submit") with gr.Column(): music_out = gr.Audio(label="Audio Result") submit_btn.click( fn = infer, inputs = [genre_txt, lyrics_txt], outputs = [music_out] ) demo.queue().launch(show_api=False, show_error=True)