Speech 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_SPEECH_CREATION_SERVICE_URL = os.getenv(15    "NATIVE_SPEECH_CREATION_SERVICE_URL", "http://native-speech-creation-service:8108"16)17_EMBLEMA_VERSION = os.getenv("EMBLEMA_VERSION", "dev")18NATIVE_SPEECH_CREATION_SERVICE_IMAGE = os.getenv(19    "NATIVE_SPEECH_CREATION_SERVICE_IMAGE",20    f"emblema/native-speech-creation-service:{_EMBLEMA_VERSION}",21)22HF_CACHE_HOST_PATH = os.getenv("HF_CACHE_HOST_PATH", "/root/.cache/huggingface")23CONTAINER_NAME = "native-speech-creation-service"24INPUT_DIR = "/data/input"25OUTPUT_DIR = "/data/output"262728def start_container():29    """Create and start native-speech-creation-service, removing any stale container first."""30    subprocess.run(31        ["docker", "rm", "-f", CONTAINER_NAME],32        capture_output=True, text=True33    )3435    hf_token = os.getenv("HUGGINGFACE_TOKEN", "")36    print(f"Creating container {CONTAINER_NAME}...", file=sys.stderr)37    run_cmd = [38        "docker", "run", "-d",39        "--name", CONTAINER_NAME,40        "--network", "emblema",41        "--gpus", "all",42        "-e", "PORT=8108",43        "-e", "DEVICE=cuda",44        "-e", f"HF_TOKEN={hf_token}",45        "-v", f"{HF_CACHE_HOST_PATH}:/root/.cache/huggingface",46        NATIVE_SPEECH_CREATION_SERVICE_IMAGE,47    ]48    result = subprocess.run(run_cmd, capture_output=True, text=True)49    if result.returncode != 0:50        print(f"docker run failed (exit {result.returncode}): {result.stderr}", file=sys.stderr)51        raise RuntimeError(f"Failed to start container: {result.stderr}")5253    # Poll health endpoint54    timeout = 18055    interval = 356    elapsed = 057    health_url = f"{NATIVE_SPEECH_CREATION_SERVICE_URL}/health"58    while elapsed < timeout:59        try:60            r = requests.get(health_url, timeout=5)61            if r.status_code == 200:62                print(f"Container healthy (waited {elapsed}s).", file=sys.stderr)63                return64        except requests.ConnectionError:65            pass66        time.sleep(interval)67        elapsed += interval6869    raise RuntimeError(f"Container did not become healthy within {timeout}s")707172def stop_container():73    """Remove the container."""74    try:75        subprocess.run(76            ["docker", "rm", "-f", CONTAINER_NAME],77            capture_output=True, text=True, timeout=3078        )79        print(f"Container {CONTAINER_NAME} removed.", file=sys.stderr)80    except Exception as e:81        print(f"Warning: failed to remove container: {e}", file=sys.stderr)828384def main():85    try:86        input_json = sys.stdin.read()87        execution_input = json.loads(input_json)88        inputs = execution_input.get("inputs", {})8990        prompt = inputs.get("prompt", "")91        voice_reference = inputs.get("voice_reference", "")92        language = inputs.get("language", "Italian")93        exaggeration = float(inputs.get("exaggeration", 0.8))94        temperature = float(inputs.get("temperature", 0.8))9596        if not prompt:97            raise ValueError("Prompt input is required")98        if not (0.0 <= exaggeration <= 1.0):99            raise ValueError(f"Exaggeration must be between 0.0 and 1.0, got {exaggeration}")100        if not (0.0 <= temperature <= 1.5):101            raise ValueError(f"Temperature must be between 0.0 and 1.5, got {temperature}")102103        os.makedirs(OUTPUT_DIR, exist_ok=True)104105        # Start the container106        start_container()107108        try:109            # Build multipart form data110            form_data = {111                "text": (None, prompt),112                "language": (None, language),113                "exaggeration": (None, str(exaggeration)),114                "temperature": (None, str(temperature)),115            }116117            # Optionally attach voice reference file118            if voice_reference:119                voice_path = os.path.join(INPUT_DIR, voice_reference)120                if not os.path.exists(voice_path):121                    raise FileNotFoundError(f"Voice reference not found: {voice_path}")122                form_data["voice_reference"] = (123                    os.path.basename(voice_path),124                    open(voice_path, "rb"),125                    "audio/mpeg",126                )127128            resp = requests.post(129                f"{NATIVE_SPEECH_CREATION_SERVICE_URL}/generate",130                files=form_data,131                timeout=600,132            )133134            if resp.status_code != 200:135                try:136                    error_detail = resp.json()137                except Exception:138                    error_detail = resp.text139                raise RuntimeError(140                    f"Speech creation service returned {resp.status_code}: {error_detail}"141                )142143            # Save result144            out_filename = "generated_speech.mp3"145            out_path = os.path.join(OUTPUT_DIR, out_filename)146            with open(out_path, "wb") as f:147                f.write(resp.content)148149            inference_time = resp.headers.get("X-Inference-Time-Ms", "unknown")150            print(151                f"Speech generated: time={inference_time}ms, language={language}, "152                f"exaggeration={exaggeration}, temperature={temperature}, "153                f"voice_cloning={bool(voice_reference)}",154                file=sys.stderr,155            )156157            # Flat output -- keys match OUTPUT_SCHEMA158            output = {159                "audio": out_filename,160            }161            print(json.dumps(output, indent=2))162163        finally:164            stop_container()165166    except Exception as e:167        error_output = {168            "error": str(e),169            "errorType": type(e).__name__,170            "traceback": traceback.format_exc(),171        }172        print(json.dumps(error_output), file=sys.stderr)173        sys.exit(1)174175176if __name__ == "__main__":177    main()