Music Creation

1import os2import sys3import json4import subprocess5import time6import traceback78try:9    import requests10except ImportError:11    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])12    import requests1314NATIVE_MUSIC_CREATION_SERVICE_URL = os.getenv(15    "NATIVE_MUSIC_CREATION_SERVICE_URL", "http://native-music-creation-service:8110"16)17_EMBLEMA_VERSION = os.getenv("EMBLEMA_VERSION", "dev")18NATIVE_MUSIC_CREATION_SERVICE_IMAGE = os.getenv(19    "NATIVE_MUSIC_CREATION_SERVICE_IMAGE",20    f"emblema/native-music-creation-service:{_EMBLEMA_VERSION}",21)22HF_CACHE_HOST_PATH = os.getenv("HF_CACHE_HOST_PATH", "/root/.cache/huggingface")23CONTAINER_NAME = "native-music-creation-service"24OUTPUT_DIR = "/data/output"252627def start_container():28    """Create and start native-music-creation-service, removing any stale container first."""29    subprocess.run(30        ["docker", "rm", "-f", CONTAINER_NAME],31        capture_output=True, text=True32    )3334    hf_token = os.getenv("HUGGINGFACE_TOKEN", "")35    print(f"Creating container {CONTAINER_NAME}...", file=sys.stderr)36    run_cmd = [37        "docker", "run", "-d",38        "--name", CONTAINER_NAME,39        "--network", "emblema",40        "--gpus", "all",41        "-e", "PORT=8110",42        "-e", "DEVICE=cuda",43        "-e", f"HF_TOKEN={hf_token}",44        "-v", f"{HF_CACHE_HOST_PATH}:/root/.cache/huggingface",45        NATIVE_MUSIC_CREATION_SERVICE_IMAGE,46    ]47    result = subprocess.run(run_cmd, capture_output=True, text=True)48    if result.returncode != 0:49        print(f"docker run failed (exit {result.returncode}): {result.stderr}", file=sys.stderr)50        raise RuntimeError(f"Failed to start container: {result.stderr}")5152    # Poll health endpoint53    timeout = 18054    interval = 355    elapsed = 056    health_url = f"{NATIVE_MUSIC_CREATION_SERVICE_URL}/health"57    while elapsed < timeout:58        try:59            r = requests.get(health_url, timeout=5)60            if r.status_code == 200:61                print(f"Container healthy (waited {elapsed}s).", file=sys.stderr)62                return63        except requests.ConnectionError:64            pass65        time.sleep(interval)66        elapsed += interval6768    raise RuntimeError(f"Container did not become healthy within {timeout}s")697071def stop_container():72    """Remove the container."""73    try:74        subprocess.run(75            ["docker", "rm", "-f", CONTAINER_NAME],76            capture_output=True, text=True, timeout=3077        )78        print(f"Container {CONTAINER_NAME} removed.", file=sys.stderr)79    except Exception as e:80        print(f"Warning: failed to remove container: {e}", file=sys.stderr)818283def main():84    try:85        input_json = sys.stdin.read()86        execution_input = json.loads(input_json)87        inputs = execution_input.get("inputs", {})8889        prompt = inputs.get("prompt", "")90        lyrics = inputs.get("lyrics", "")91        duration = int(inputs.get("duration", 90))92        language = inputs.get("language", "en")93        time_signature = inputs.get("time_signature", "4")94        bpm = int(inputs.get("bpm", 190))95        keyscale = inputs.get("keyscale", "E minor")9697        if not prompt:98            raise ValueError("Prompt input is required")99        if not lyrics:100            raise ValueError("Lyrics input is required")101        if not (10 <= duration <= 300):102            raise ValueError(f"Duration must be between 10 and 300, got {duration}")103        if not (40 <= bpm <= 300):104            raise ValueError(f"BPM must be between 40 and 300, got {bpm}")105106        os.makedirs(OUTPUT_DIR, exist_ok=True)107108        # Start the container109        start_container()110111        try:112            # JSON POST (no file inputs)113            payload = {114                "prompt": prompt,115                "lyrics": lyrics,116                "duration": duration,117                "language": language,118                "time_signature": time_signature,119                "bpm": bpm,120                "keyscale": keyscale,121            }122123            resp = requests.post(124                f"{NATIVE_MUSIC_CREATION_SERVICE_URL}/generate",125                json=payload,126                timeout=600,127            )128129            if resp.status_code != 200:130                try:131                    error_detail = resp.json()132                except Exception:133                    error_detail = resp.text134                raise RuntimeError(135                    f"Music creation service returned {resp.status_code}: {error_detail}"136                )137138            # Save result139            out_filename = "generated_music.mp3"140            out_path = os.path.join(OUTPUT_DIR, out_filename)141            with open(out_path, "wb") as f:142                f.write(resp.content)143144            inference_time = resp.headers.get("X-Inference-Time-Ms", "unknown")145            print(146                f"Music generated: time={inference_time}ms, duration={duration}s, "147                f"bpm={bpm}, key={keyscale}, time_sig={time_signature}, language={language}",148                file=sys.stderr,149            )150151            # Flat output -- keys match OUTPUT_SCHEMA152            output = {153                "audio": out_filename,154            }155            print(json.dumps(output, indent=2))156157        finally:158            stop_container()159160    except Exception as e:161        error_output = {162            "error": str(e),163            "errorType": type(e).__name__,164            "traceback": traceback.format_exc(),165        }166        print(json.dumps(error_output), file=sys.stderr)167        sys.exit(1)168169170if __name__ == "__main__":171    main()