Video Creation (Helios)

1import os2import sys3import json4import subprocess5import time6import traceback7import random89try:10    import requests11except ImportError:12    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])13    import requests1415NATIVE_HELIOS_CREATION_SERVICE_URL = os.getenv(16    "NATIVE_HELIOS_CREATION_SERVICE_URL", "http://native-helios-creation-service:8112"17)18_EMBLEMA_VERSION = os.getenv("EMBLEMA_VERSION", "dev")19NATIVE_HELIOS_CREATION_SERVICE_IMAGE = os.getenv(20    "NATIVE_HELIOS_CREATION_SERVICE_IMAGE",21    f"emblema/native-helios-creation-service:{_EMBLEMA_VERSION}",22)23HF_CACHE_HOST_PATH = os.getenv("HF_CACHE_HOST_PATH", "/root/.cache/huggingface")24CONTAINER_NAME = "native-helios-creation-service"25INPUT_DIR = "/data/input"26OUTPUT_DIR = "/data/output"272829def start_container():30    """Create and start native-helios-creation-service, removing any stale container first."""31    subprocess.run(32        ["docker", "rm", "-f", CONTAINER_NAME],33        capture_output=True, text=True34    )3536    hf_token = os.getenv("HUGGINGFACE_TOKEN", "")37    print(f"Creating container {CONTAINER_NAME}...", file=sys.stderr)38    run_cmd = [39        "docker", "run", "-d",40        "--name", CONTAINER_NAME,41        "--network", "emblema",42        "--gpus", "all",43        "-e", "PORT=8112",44        "-e", "DEVICE=cuda",45        "-e", f"HF_TOKEN={hf_token}",46        "-v", f"{HF_CACHE_HOST_PATH}:/root/.cache/huggingface",47        NATIVE_HELIOS_CREATION_SERVICE_IMAGE,48    ]49    result = subprocess.run(run_cmd, capture_output=True, text=True)50    if result.returncode != 0:51        print(f"docker run failed (exit {result.returncode}): {result.stderr}", file=sys.stderr)52        raise RuntimeError(f"Failed to start container: {result.stderr}")5354    # Poll health endpoint (300s timeout for 14B model download on first run)55    timeout = 30056    interval = 557    elapsed = 058    health_url = f"{NATIVE_HELIOS_CREATION_SERVICE_URL}/health"59    while elapsed < timeout:60        try:61            r = requests.get(health_url, timeout=5)62            if r.status_code == 200:63                print(f"Container healthy (waited {elapsed}s).", file=sys.stderr)64                return65        except requests.ConnectionError:66            pass67        time.sleep(interval)68        elapsed += interval6970    raise RuntimeError(f"Container did not become healthy within {timeout}s")717273def stop_container():74    """Remove the container."""75    try:76        subprocess.run(77            ["docker", "rm", "-f", CONTAINER_NAME],78            capture_output=True, text=True, timeout=3079        )80        print(f"Container {CONTAINER_NAME} removed.", file=sys.stderr)81    except Exception as e:82        print(f"Warning: failed to remove container: {e}", file=sys.stderr)838485def main():86    try:87        input_json = sys.stdin.read()88        execution_input = json.loads(input_json)89        inputs = execution_input.get("inputs", {})9091        mode = inputs.get("mode", "t2v").lower()92        prompt = inputs.get("prompt", "")93        negative_prompt = inputs.get("negative_prompt", "")94        resolution = inputs.get("resolution", "640x384")95        num_frames = int(inputs.get("num_frames", 240))96        guidance_scale = float(inputs.get("guidance_scale", 1.0))97        pyramid_steps = inputs.get("pyramid_steps", "2,2,2")9899        seed_mode = inputs.get("seed_mode", "random")100        seed_input = int(inputs.get("seed", -1))101        if seed_mode == "fixed" and seed_input >= 0:102            seed_value = seed_input103        else:104            seed_value = random.randint(0, 2**31 - 1)105106        if mode == "t2v" and not prompt:107            raise ValueError("Prompt is required for T2V mode")108109        os.makedirs(OUTPUT_DIR, exist_ok=True)110111        # Start the container112        start_container()113114        try:115            # Build multipart request116            files = {}117            data = {118                "prompt": prompt,119                "negative_prompt": negative_prompt,120                "resolution": resolution,121                "num_frames": num_frames,122                "guidance_scale": guidance_scale,123                "seed": seed_value,124                "pyramid_steps": pyramid_steps,125            }126127            # Include image file for I2V mode128            if mode == "i2v":129                image_path = inputs.get("image", "")130                if not image_path:131                    raise ValueError("Image is required for I2V mode")132                local_path = os.path.join(INPUT_DIR, image_path)133                if not os.path.exists(local_path):134                    raise FileNotFoundError(f"Input image not found: {local_path}")135                files["image"] = (os.path.basename(local_path), open(local_path, "rb"), "image/png")136137            resp = requests.post(138                f"{NATIVE_HELIOS_CREATION_SERVICE_URL}/generate",139                files=files if files else None,140                data=data,141                timeout=1800,142            )143144            if resp.status_code != 200:145                try:146                    error_detail = resp.json()147                except Exception:148                    error_detail = resp.text149                raise RuntimeError(150                    f"Helios creation service returned {resp.status_code}: {error_detail}"151                )152153            # Save result154            out_filename = "generated_video.mp4"155            out_path = os.path.join(OUTPUT_DIR, out_filename)156            with open(out_path, "wb") as f:157                f.write(resp.content)158159            inference_time = resp.headers.get("X-Inference-Time-Ms", "unknown")160            result_mode = resp.headers.get("X-Mode", mode.upper())161            result_seed = resp.headers.get("X-Seed", str(seed_value))162            print(163                f"Video generated: mode={result_mode}, time={inference_time}ms, seed={result_seed}",164                file=sys.stderr,165            )166167            output = {168                "video": out_filename,169            }170            print(json.dumps(output, indent=2))171172        finally:173            stop_container()174175    except Exception as e:176        error_output = {177            "error": str(e),178            "errorType": type(e).__name__,179            "traceback": traceback.format_exc(),180        }181        print(json.dumps(error_output), file=sys.stderr)182        sys.exit(1)183184185if __name__ == "__main__":186    main()