MiloMusic / app.py
futurespyhi
add debugging for generation failure
ef411bc
raw
history blame
31.4 kB
#!/usr/bin/env python3
"""
MiloMusic - Hugging Face Spaces Version
AI-powered music generation platform optimized for cloud deployment with high-performance configuration.
"""
import os
import sys
import subprocess
import tempfile
import gradio as gr
import soundfile as sf
from dataclasses import dataclass, field
from typing import Any
import xxhash
import numpy as np
import spaces
import groq
# Import environment setup for Spaces
def setup_spaces_environment():
"""Setup environment variables and paths for Hugging Face Spaces"""
# Set HuggingFace cache directory
os.environ["HF_HOME"] = "/tmp/hf_cache"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/transformers_cache"
os.environ["HF_HUB_CACHE"] = "/tmp/hf_hub_cache"
# PyTorch CUDA memory optimization
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
# Set temp directory for audio files
os.environ["TMPDIR"] = "/tmp"
print("🚀 Environment setup complete for Spaces")
# Install flash-attn if not already installed
def install_flash_attn():
"""Install flash-attn from source with proper compilation flags"""
try:
import flash_attn
print("✅ flash-attn already installed")
return True
except ImportError:
print("📦 Installing flash-attn from source...")
try:
# Install with optimized settings for Spaces
cmd = [
sys.executable, "-m", "pip", "install",
"--no-build-isolation",
"--no-cache-dir",
"flash-attn",
"--verbose"
]
# Use more parallel jobs for faster compilation in Spaces
env = os.environ.copy()
env["MAX_JOBS"] = "4" # Utilize more CPU cores
env["NVCC_PREPEND_FLAGS"] = "-ccbin /usr/bin/gcc"
result = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=1800) # 30 min timeout
if result.returncode == 0:
print("✅ flash-attn installed successfully")
return True
else:
print(f"❌ flash-attn installation failed: {result.stderr}")
return False
except subprocess.TimeoutExpired:
print("⏰ flash-attn installation timed out")
return False
except Exception as e:
print(f"❌ Error installing flash-attn: {e}")
return False
# Setup environment first
setup_spaces_environment()
# Install flash-attn if needed
flash_attn_available = install_flash_attn()
# Apply transformers patches for performance optimization
def apply_transformers_patch():
"""
Apply YuEGP transformers patches for high-performance generation.
This function applies optimized transformers patches that provide:
- 2x speed improvement for low VRAM profiles
- 3x speed improvement for Stage 1 generation (16GB+ VRAM)
- 2x speed improvement for Stage 2 generation (all profiles)
The patches replace two key files in the transformers library:
- models/llama/modeling_llama.py (LLaMA model optimizations)
- generation/utils.py (generation utilities optimizations)
Includes smart detection to avoid re-applying patches on restart.
"""
try:
import shutil
import site
import hashlib
# Define source and target directories
source_dir = os.path.join(project_root, "YuEGP", "transformers")
# Get the site-packages directory where transformers is installed
site_packages = site.getsitepackages()
if not site_packages:
# Fallback for some environments
import transformers
transformers_path = os.path.dirname(transformers.__file__)
target_base = os.path.dirname(transformers_path)
else:
target_base = site_packages[0]
target_dir = os.path.join(target_base, "transformers")
# Check if source patches exist
if not os.path.exists(source_dir):
print("⚠️ YuEGP transformers patches not found, skipping optimization")
return False
if not os.path.exists(target_dir):
print("⚠️ Transformers library not found, skipping patches")
return False
# Check if patches are already applied by comparing file hashes
def get_file_hash(filepath):
"""Get MD5 hash of file content"""
if not os.path.exists(filepath):
return None
with open(filepath, 'rb') as f:
return hashlib.md5(f.read()).hexdigest()
# Key files to check for patch status
key_patches = [
"models/llama/modeling_llama.py",
"generation/utils.py"
]
patches_needed = False
for patch_file in key_patches:
source_file = os.path.join(source_dir, patch_file)
target_file = os.path.join(target_dir, patch_file)
if os.path.exists(source_file):
source_hash = get_file_hash(source_file)
target_hash = get_file_hash(target_file)
if source_hash != target_hash:
patches_needed = True
break
if not patches_needed:
print("✅ YuEGP transformers patches already applied, skipping re-installation")
print(" 📈 High-performance optimizations are active:")
print(" • Stage 1 generation: 3x faster (16GB+ VRAM)")
print(" • Stage 2 generation: 2x faster (all profiles)")
return True
# Apply patches by copying optimized files
print("🔧 Applying YuEGP transformers patches for high-performance generation...")
# Copy the patched files, preserving directory structure
for root, dirs, files in os.walk(source_dir):
# Calculate relative path from source_dir
rel_path = os.path.relpath(root, source_dir)
target_subdir = os.path.join(target_dir, rel_path) if rel_path != '.' else target_dir
# Ensure target subdirectory exists
os.makedirs(target_subdir, exist_ok=True)
# Copy all Python files in this directory
for file in files:
if file.endswith('.py'):
src_file = os.path.join(root, file)
dst_file = os.path.join(target_subdir, file)
shutil.copy2(src_file, dst_file)
print(f" ✅ Patched: {os.path.relpath(dst_file, target_base)}")
print("🚀 Transformers patches applied successfully!")
print(" 📈 Expected performance gains:")
print(" • Stage 1 generation: 3x faster (16GB+ VRAM)")
print(" • Stage 2 generation: 2x faster (all profiles)")
return True
except Exception as e:
print(f"❌ Error applying transformers patches: {e}")
print(" Continuing without patches - performance may be reduced")
return False
# Now import the rest of the dependencies
# Add project root to Python path for imports
project_root = os.path.dirname(os.path.abspath(__file__))
if project_root not in sys.path:
sys.path.insert(0, project_root)
from tools.groq_client import client as groq_client
from openai import OpenAI
from tools.generate_lyrics import generate_structured_lyrics, format_lyrics
# Apply patches after all imports are set up
patch_applied = apply_transformers_patch()
# Import CUDA info after flash-attn setup
import torch
if torch.cuda.is_available():
print(f"🎮 GPU: {torch.cuda.get_device_name(0)}")
print(f"💾 VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
else:
print("⚠️ No CUDA GPU detected")
@dataclass
class AppState:
"""
Maintains the application state throughout user interactions.
"""
conversation: list = field(default_factory=list)
stopped: bool = False
model_outs: Any = None
lyrics: str = ""
genre: str = "pop"
mood: str = "upbeat"
theme: str = "love"
def validate_api_keys():
"""Validate required API keys for Spaces deployment"""
required_keys = ["GROQ_API_KEY", "GEMINI_API_KEY"]
missing_keys = []
for key in required_keys:
if not os.getenv(key):
missing_keys.append(key)
if missing_keys:
print(f"⚠️ Missing API keys: {missing_keys}")
return False
print("✅ All API keys validated")
return True
def validate_file_structure():
"""Validate that required files and directories exist"""
required_paths = [
"YuEGP/inference/infer.py",
"YuEGP/inference/codecmanipulator.py",
"YuEGP/inference/mmtokenizer.py",
"tools/generate_lyrics.py",
"tools/groq_client.py",
"schemas/lyrics.py" # Required for lyrics structure models
]
missing_files = []
for path in required_paths:
if not os.path.exists(path):
missing_files.append(path)
if missing_files:
print(f"⚠️ Missing required files: {missing_files}")
return False
print("✅ All required files found")
return True
@spaces.GPU # Decorator for GPU access in Spaces
def generate_music_spaces(lyrics: str, genre: str, mood: str, progress=gr.Progress()) -> str:
"""
Generate music using YuE model with high-performance Spaces configuration
"""
if not lyrics.strip():
return "Please provide lyrics to generate music."
try:
progress(0.1, desc="Preparing lyrics...")
# Use lyrics directly (already formatted from chat interface)
formatted_lyrics = lyrics
# Create temporary files
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as genre_file:
genre_file.write(f"instrumental,{genre},{mood},male vocals")
genre_file_path = genre_file.name
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as lyrics_file:
lyrics_file.write(formatted_lyrics)
lyrics_file_path = lyrics_file.name
progress(0.2, desc="Setting up generation...")
# Generate music with high-performance Spaces configuration
output_dir = tempfile.mkdtemp()
# High-performance command based on Spaces GPU resources
# In Spaces, working directory is /app
infer_script_path = os.path.join(os.getcwd(), "YuEGP", "inference", "infer.py")
cmd = [
sys.executable,
infer_script_path,
"--cuda_idx", "0",
"--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
"--stage2_model", "m-a-p/YuE-s2-1B-general",
"--genre_txt", genre_file_path,
"--lyrics_txt", lyrics_file_path,
"--run_n_segments", "2", # Full segments for better quality
"--stage2_batch_size", "4", # Higher batch size for speed
"--output_dir", output_dir,
"--max_new_tokens", "3000", # Full token count
"--profile", "1", # Highest performance profile
"--verbose", "3",
"--prompt_start_time", "0",
"--prompt_end_time", "30", # Full 30-second clips
]
# Use flash attention if available, otherwise fallback
if not flash_attn_available:
cmd.append("--sdpa")
progress(0.3, desc="Starting music generation (Stage 1)...")
# Run with generous timeout for high-quality generation
print("🎵 Starting high-quality music generation...")
print(f"Working directory: {os.getcwd()}")
print(f"Infer script path: {infer_script_path}")
print(f"Command: {' '.join(cmd)}")
# Change to YuEGP/inference directory for execution
original_cwd = os.getcwd()
inference_dir = os.path.join(os.getcwd(), "YuEGP", "inference")
try:
os.chdir(inference_dir)
print(f"Changed to inference directory: {inference_dir}")
# Update command to use relative path since we're in the inference directory
cmd[1] = "infer.py"
# Run with real-time output for debugging
print(f"🚀 Executing command: {' '.join(cmd)}")
result = subprocess.run(cmd, capture_output=True, text=True, timeout=1200) # 20 minutes
# Print stdout and stderr for debugging
if result.stdout:
print(f"✅ Command output:\n{result.stdout}")
if result.stderr:
print(f"⚠️ Command stderr:\n{result.stderr}")
print(f"📊 Return code: {result.returncode}")
finally:
# Always restore original working directory
os.chdir(original_cwd)
progress(0.9, desc="Finalizing audio output...")
# Clean up input files
os.unlink(genre_file_path)
os.unlink(lyrics_file_path)
if result.returncode == 0:
# Find generated audio file
import glob
audio_files = glob.glob(os.path.join(output_dir, "**/*.mp3"), recursive=True)
if audio_files:
progress(1.0, desc="Music generation complete!")
return audio_files[0] # Return path to generated audio
else:
return "Music generation completed but no audio file found."
else:
error_msg = f"Return code: {result.returncode}\n"
if result.stderr:
error_msg += f"Error: {result.stderr[-1000:]}\n"
if result.stdout:
error_msg += f"Output: {result.stdout[-1000:]}"
return f"Music generation failed:\n{error_msg}"
except subprocess.TimeoutExpired:
return "Music generation timed out after 20 minutes. Please try again."
except Exception as e:
return f"Error during music generation: {str(e)}"
def respond(message, state):
"""Enhanced response function for lyrics generation"""
try:
# Add user message to conversation
state.conversation.append({"role": "user", "content": message})
# Generate response using your existing lyrics generation logic
song_structure = generate_structured_lyrics(
state.conversation,
state.genre,
state.mood,
state.theme
)
# Format the structured lyrics for display
response = format_lyrics(song_structure)
# Add assistant response
state.conversation.append({"role": "assistant", "content": response})
# Update lyrics if this looks like final lyrics
if any(marker in response.lower() for marker in ["[verse", "[chorus", "[bridge"]):
state.lyrics = response
# Return conversation for display
conversation_display = []
for msg in state.conversation:
role = "User" if msg["role"] == "user" else "Assistant"
conversation_display.append([msg["content"] if msg["role"] == "user" else None,
msg["content"] if msg["role"] == "assistant" else None])
return "", conversation_display, state
except Exception as e:
error_response = f"Sorry, I encountered an error: {str(e)}"
state.conversation.append({"role": "assistant", "content": error_response})
# Format conversation for display
conversation_display = []
for msg in state.conversation:
conversation_display.append([msg["content"] if msg["role"] == "user" else None,
msg["content"] if msg["role"] == "assistant" else None])
return "", conversation_display, state
def build_interface():
"""Build the Gradio interface optimized for Spaces with high performance"""
with gr.Blocks(
title="MiloMusic - AI Music Generation",
theme=gr.themes.Soft(),
css="""
.container { max-width: 1400px; margin: auto; }
.performance-notice { background-color: #d4edda; padding: 15px; border-radius: 5px; margin: 10px 0; }
.generation-status { background-color: #f8f9fa; padding: 10px; border-radius: 5px; }
"""
) as demo:
# Header
gr.Markdown("""
# 🎵 MiloMusic - AI Music Generation
### Professional AI-powered music creation from natural language
""")
# Performance notice for Spaces
gr.Markdown("""
<div class="performance-notice">
🚀 <strong>High-Performance Mode:</strong> Running on Spaces GPU with optimized settings for best quality.
Generation time: ~3-5 minutes for professional-grade music with vocals and instruments.
</div>
""")
state = gr.State(AppState())
with gr.Row():
with gr.Column(scale=2):
# Input controls
with gr.Group():
gr.Markdown("### 🎛️ Music Settings")
with gr.Row():
genre = gr.Dropdown(
choices=["pop", "rock", "jazz", "classical", "electronic", "folk", "r&b", "country", "hip-hop"],
value="pop", label="Genre"
)
mood = gr.Dropdown(
choices=["upbeat", "melancholic", "energetic", "calm", "romantic", "dark", "mysterious", "joyful"],
value="upbeat", label="Mood"
)
theme = gr.Dropdown(
choices=["love", "friendship", "adventure", "nostalgia", "freedom", "hope", "dreams", "nature"],
value="love", label="Theme"
)
# Voice Input
with gr.Group():
gr.Markdown("### 🎤 Voice Input")
input_audio = gr.Audio(
label="Speak Your Musical Ideas",
sources=["microphone"],
type="numpy",
streaming=False,
waveform_options=gr.WaveformOptions(waveform_color="#B83A4B"),
)
# Chat interface
with gr.Group():
gr.Markdown("### 💬 Lyrics Creation Chat")
chatbot = gr.Chatbot(height=400, label="AI Lyrics Assistant", show_copy_button=True)
with gr.Row():
text_input = gr.Textbox(
placeholder="Or type your song idea here...",
show_label=False,
scale=4,
lines=2
)
send_btn = gr.Button("Send", scale=1, variant="primary")
with gr.Column(scale=1):
# Output controls
with gr.Group():
gr.Markdown("### 🎵 Music Generation")
lyrics_display = gr.Textbox(
label="Current Lyrics",
lines=12,
interactive=True,
placeholder="Your generated lyrics will appear here..."
)
generate_btn = gr.Button("🎼 Generate High-Quality Music", variant="primary", size="lg")
with gr.Column():
music_output = gr.Audio(label="Generated Music", type="filepath", show_download_button=True)
gr.Markdown("""
<div class="generation-status">
<strong>Generation Features:</strong><br>
• Full 30-second clips<br>
• Professional vocals<br>
• Rich instrumentation<br>
• High-fidelity audio
</div>
""")
# Controls
with gr.Group():
gr.Markdown("### 🔧 Controls")
new_song_btn = gr.Button("🆕 Start New Song")
clear_btn = gr.Button("🧹 Clear Chat")
# Event handlers
def update_state_settings(genre_val, mood_val, theme_val, state):
state.genre = genre_val
state.mood = mood_val
state.theme = theme_val
return state
# Update state when settings change
for component in [genre, mood, theme]:
component.change(
fn=update_state_settings,
inputs=[genre, mood, theme, state],
outputs=[state]
)
# Voice recording functionality (from app.py)
stream = input_audio.start_recording(
process_audio,
[input_audio, state],
[input_audio, state],
)
respond_audio = input_audio.stop_recording(
response_audio, [state, input_audio, genre, mood, theme], [state, chatbot]
)
restart = respond_audio.then(start_recording_user, [state], [input_audio]).then(
lambda state: state, state, state, js=js_reset
)
# Text chat functionality
send_btn.click(
fn=respond,
inputs=[text_input, state],
outputs=[text_input, chatbot, state],
queue=True
)
text_input.submit(
fn=respond,
inputs=[text_input, state],
outputs=[text_input, chatbot, state],
queue=True
)
# Music generation with progress
generate_btn.click(
fn=generate_music_spaces,
inputs=[lyrics_display, genre, mood],
outputs=[music_output],
queue=True,
show_progress=True
)
# Control buttons
new_song_btn.click(
fn=lambda: (AppState(), [], "", None, gr.Audio(recording=False)),
outputs=[state, chatbot, lyrics_display, music_output, input_audio],
cancels=[respond_audio, restart]
)
clear_btn.click(
fn=lambda: [],
outputs=[chatbot]
)
# Auto-update lyrics display when state changes
state.change(
fn=lambda s: s.lyrics,
inputs=[state],
outputs=[lyrics_display]
)
# Instructions
gr.Markdown("""
### 📖 How to create your music:
1. **Set your preferences**: Choose genre, mood, and theme
2. **Voice or chat**: Either speak your ideas or type them in the chat
3. **Refine the lyrics**: Ask for changes, different verses, or style adjustments
4. **Generate music**: Click the generate button for professional-quality output
5. **Download & enjoy**: Your high-fidelity music with vocals and instruments
**Tips**: Be specific about your vision - mention instruments, vocal style, or song structure!
""")
# Footer
gr.Markdown("""
---
<center>
Made with ❤️ by the MiloMusic Team | Powered by YuE (乐) Model | 🤗 Hugging Face Spaces
</center>
""")
return demo
# Audio transcription functions (from app.py)
def process_whisper_response(completion):
"""
Process Whisper transcription response and filter out silence.
"""
if completion.segments and len(completion.segments) > 0:
no_speech_prob = completion.segments[0].get('no_speech_prob', 0)
print("No speech prob:", no_speech_prob)
if no_speech_prob > 0.7:
print("No speech detected")
return None
return completion.text.strip()
return None
def transcribe_audio(client, file_name):
"""
Transcribe an audio file using the Whisper model via the Groq API.
"""
if file_name is None:
return None
try:
with open(file_name, "rb") as audio_file:
with open("audio.wav", "wb") as f:
f.write(audio_file.read())
response = client.audio.transcriptions.create(
model="whisper-large-v3-turbo",
file=("audio.wav", audio_file),
response_format="text",
language="en",
)
# Process the response to filter out silence
# For text response format, we need to check if response is meaningful
if response and len(response.strip()) > 0:
return response.strip()
else:
return None
except Exception as e:
print(f"Transcription error: {e}")
return f"Error in audio transcription: {str(e)}"
def start_recording_user(state: AppState):
"""
Reset the audio recording component for a new user input.
"""
return None
def process_audio(audio: tuple, state: AppState):
"""
Process recorded audio in real-time during recording.
"""
return audio, state
@spaces.GPU(duration=40, progress=gr.Progress(track_tqdm=True))
def response_audio(state: AppState, audio: tuple, genre_value, mood_value, theme_value):
"""
Process recorded audio and generate a response based on transcription.
"""
if not audio:
return state, []
# Update state with current dropdown values
state.genre, state.mood, state.theme = genre_value, mood_value, theme_value
temp_dir = tempfile.gettempdir()
file_name = os.path.join(temp_dir, f"{xxhash.xxh32(bytes(audio[1])).hexdigest()}.wav")
sf.write(file_name, audio[1], audio[0], format="wav")
api_key = os.environ.get("GROQ_API_KEY")
if not api_key:
raise ValueError("Please set the GROQ_API_KEY environment variable.")
client = groq.Client(api_key=api_key)
# Transcribe the audio file
transcription = transcribe_audio(client, file_name)
if transcription:
if isinstance(transcription, str) and transcription.startswith("Error"):
transcription = "Error in audio transcription."
state.conversation.append({"role": "user", "content": transcription})
assistant_message = generate_chat_completion(client, state.conversation, state.genre, state.mood, state.theme)
state.conversation.append({"role": "assistant", "content": assistant_message})
# Update lyrics from conversation
state.lyrics = extract_lyrics_from_conversation(state.conversation)
os.remove(file_name)
# Format conversation for display
conversation_display = []
for msg in state.conversation:
conversation_display.append([msg["content"] if msg["role"] == "user" else None,
msg["content"] if msg["role"] == "assistant" else None])
return state, conversation_display
def extract_lyrics_from_conversation(conversation):
"""
Extract lyrics from conversation history.
"""
lyrics = ""
for message in reversed(conversation):
if message["role"] == "assistant" and "verse" in message["content"].lower() and "chorus" in message["content"].lower():
lyrics = message["content"]
break
return lyrics
def generate_chat_completion(client, history, genre, mood, theme):
"""
Generate an AI assistant response based on conversation history and song parameters.
"""
messages = []
system_prompt = f"""You are a creative AI music generator assistant. Help users create song lyrics in the {genre} genre with a {mood} mood about {theme}.
When generating lyrics, create a chorus and at least one verse. Format lyrics clearly with VERSE and CHORUS labels.
Ask if they like the lyrics or want changes. Be conversational, friendly, and creative.
Keep the lyrics appropriate for the selected genre, mood, and theme unless the user specifically requests changes."""
messages.append({
"role": "system",
"content": system_prompt,
})
for message in history:
messages.append(message)
try:
completion = client.chat.completions.create(
model="meta-llama/llama-4-scout-17b-16e-instruct",
messages=messages,
)
return completion.choices[0].message.content
except Exception as e:
return f"Error in generating chat completion: {str(e)}"
# JavaScript for frontend enhancements
js_reset = """
() => {
var record = document.querySelector('.record-button');
if (record) {
record.textContent = "Just Start Talking!"
record.style = "width: fit-content; padding-right: 0.5vw;"
}
}
"""
# Build the interface
demo = build_interface()
if __name__ == "__main__":
"""
Spaces entry point - optimized for high-performance deployment
"""
print("🚀 Starting MiloMusic High-Performance Mode on Hugging Face Spaces...")
print(f"📁 Working directory: {os.getcwd()}")
print(f"📂 Directory contents: {os.listdir('.')}")
# Validate file structure
if not validate_file_structure():
print("❌ Required files missing - please check your upload")
sys.exit(1)
# Validate environment
if not validate_api_keys():
print("⚠️ Some API keys missing - functionality may be limited")
# Launch with optimized settings for Spaces
demo.queue(
default_concurrency_limit=5, # Allow more concurrent users
max_size=20
).launch(
server_name="0.0.0.0",
server_port=7860,
share=False, # Spaces handles sharing
show_error=True,
quiet=False,
favicon_path=None,
ssl_verify=False
)