battlewords/generate_sounds.py

#!/usr/bin/env python3
"""
Standalone script to generate sound effects using Hugging Face API.
Uses only built-in Python libraries (no external dependencies).
"""

import os
import json
import urllib.request
import time
from pathlib import Path

# Load environment variables from .env if present
env_path = Path(__file__).parent / ".env"
if env_path.exists():
    with open(env_path) as f:
        for line in f:
            if line.strip() and not line.startswith("#"):
                key, _, value = line.strip().partition("=")
                os.environ[key] = value

# Get Hugging Face API token from environment variable
HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
if not HF_API_TOKEN:
    print("Warning: HF_API_TOKEN not set in environment or .env file.")

# Using your UnlimitedMusicGen Gradio Space
SPACE_URL = "https://surn-unlimitedmusicgen.hf.space"
GRADIO_API_URL = f"{SPACE_URL}/api/predict"
GRADIO_STATUS_URL = f"{SPACE_URL}/call/predict/{{event_id}}"

# Sound effects to generate
EFFECT_PROMPTS = {
    "correct_guess": {"prompt": "A short, sharp ding sound for a correct guess", "duration": 2},
    "incorrect_guess": {"prompt": "A low buzz sound for an incorrect guess", "duration": 2},
    "miss": {"prompt": "A soft thud sound for a miss", "duration": 1},
    "hit": {"prompt": "A bright chime sound for a hit", "duration": 1},
    "congratulations": {"prompt": "A triumphant fanfare sound for congratulations", "duration": 3}
}

def generate_sound_effect_gradio(effect_name: str, prompt: str, duration: float, output_dir: Path) -> bool:
    """Generate a single sound effect using Gradio API (async call)."""

    print(f"\nGenerating: {effect_name}")
    print(f"  Prompt: {prompt}")
    print(f"  Duration: {duration}s")

    # Step 1: Submit generation request
    payload = json.dumps({
        "data": [prompt, duration]
    }).encode('utf-8')

    headers = {
        "Content-Type": "application/json"
    }

    try:
        print(f"  Submitting request to Gradio API...")

        # Submit the job
        req = urllib.request.Request(GRADIO_API_URL, data=payload, headers=headers, method='POST')

        with urllib.request.urlopen(req, timeout=30) as response:
            if response.status == 200:
                result = json.loads(response.read().decode())
                event_id = result.get("event_id")

                if not event_id:
                    print(f"  ✗ No event_id returned")
                    return False

                print(f"  Job submitted, event_id: {event_id}")

                # Step 2: Poll for results
                status_url = GRADIO_STATUS_URL.format(event_id=event_id)

                for poll_attempt in range(30):  # Poll for up to 5 minutes
                    time.sleep(10)
                    print(f"  Polling for results (attempt {poll_attempt + 1}/30)...")

                    status_req = urllib.request.Request(status_url, headers=headers)

                    try:
                        with urllib.request.urlopen(status_req, timeout=30) as status_response:
                            # Gradio returns streaming events, read until we get the result
                            for line in status_response:
                                line = line.decode('utf-8').strip()
                                if line.startswith('data: '):
                                    event_data = json.loads(line[6:])  # Remove 'data: ' prefix

                                    if event_data.get('msg') == 'process_completed':
                                        # Get the audio file URL
                                        output_data = event_data.get('output', {}).get('data', [])
                                        if output_data and len(output_data) > 0:
                                            audio_url = output_data[0].get('url')
                                            if audio_url:
                                                # Download the audio file
                                                full_audio_url = f"https://surn-unlimitedmusicgen.hf.space{audio_url}"
                                                print(f"  Downloading from: {full_audio_url}")

                                                audio_req = urllib.request.Request(full_audio_url)
                                                with urllib.request.urlopen(audio_req, timeout=30) as audio_response:
                                                    audio_data = audio_response.read()

                                                    # Save to file
                                                    output_path = output_dir / f"{effect_name}.wav"
                                                    with open(output_path, "wb") as f:
                                                        f.write(audio_data)

                                                    print(f"  ✓ Success! Saved to: {output_path}")
                                                    print(f"  File size: {len(audio_data)} bytes")
                                                    return True

                                    elif event_data.get('msg') == 'process_error':
                                        print(f"  ✗ Generation error: {event_data.get('output')}")
                                        return False

                    except Exception as poll_error:
                        print(f"  Polling error: {poll_error}")
                        continue

                print(f"  ✗ Timeout waiting for generation")
                return False

            else:
                print(f"  ✗ Error {response.status}: {response.read().decode()}")
                return False

    except Exception as e:
        print(f"  ✗ Error: {e}")
        return False

def main():
    """Generate all sound effects."""

    print("=" * 70)
    print("Sound Effects Generator for BattleWords")
    print("=" * 70)
    print(f"Using UnlimitedMusicGen Gradio API")
    print(f"API URL: {GRADIO_API_URL}")
    print(f"\nGenerating {len(EFFECT_PROMPTS)} sound effects...\n")

    # Create output directory
    output_dir = Path(__file__).parent / "assets" / "audio"
    output_dir.mkdir(parents=True, exist_ok=True)
    print(f"Output directory: {output_dir}\n")

    # Generate each effect
    success_count = 0
    for effect_name, config in EFFECT_PROMPTS.items():
        if generate_sound_effect_gradio(
            effect_name,
            config["prompt"],
            config["duration"],
            output_dir
        ):
            success_count += 1

        # Small delay between requests
        if effect_name != list(EFFECT_PROMPTS.keys())[-1]:
            print("  Waiting 5 seconds before next request...")
            time.sleep(5)

    print("\n" + "=" * 70)
    print(f"Generation complete! {success_count}/{len(EFFECT_PROMPTS)} successful")
    print("=" * 70)

    if success_count == len(EFFECT_PROMPTS):
        print("\n✓ All sound effects generated successfully!")
    else:
        print(f"\n⚠ {len(EFFECT_PROMPTS) - success_count} sound effects failed to generate")

if __name__ == "__main__":
    main()