Video Upscale

1import os2import sys3import json4import subprocess5import time6import traceback78try:9    import requests10except ImportError:11    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])12    import requests1314NATIVE_VIDEO_UPSCALE_SERVICE_URL = os.getenv(15    "NATIVE_VIDEO_UPSCALE_SERVICE_URL", "http://native-video-upscale-service:8113"16)17_EMBLEMA_VERSION = os.getenv("EMBLEMA_VERSION", "dev")18NATIVE_VIDEO_UPSCALE_SERVICE_IMAGE = os.getenv(19    "NATIVE_VIDEO_UPSCALE_SERVICE_IMAGE",20    f"emblema/native-video-upscale-service:{_EMBLEMA_VERSION}",21)22CONTAINER_NAME = "native-video-upscale-service"23INPUT_DIR = "/data/input"24OUTPUT_DIR = "/data/output"2526# Quality level mapping: (mode, quality) → nvvfx QualityLevel int27QUALITY_MAP = {28    ("upscale", "low"): 1,29    ("upscale", "medium"): 2,30    ("upscale", "high"): 3,31    ("upscale", "ultra"): 4,32    ("upscale_hbr", "low"): 16,33    ("upscale_hbr", "medium"): 17,34    ("upscale_hbr", "high"): 18,35    ("upscale_hbr", "ultra"): 19,36    ("denoise", "low"): 8,37    ("denoise", "medium"): 9,38    ("denoise", "high"): 10,39    ("denoise", "ultra"): 11,40    ("deblur", "low"): 12,41    ("deblur", "medium"): 13,42    ("deblur", "high"): 14,43    ("deblur", "ultra"): 15,44}454647def start_container():48    """Create and start native-video-upscale-service, removing any stale container first."""49    subprocess.run(50        ["docker", "rm", "-f", CONTAINER_NAME],51        capture_output=True, text=True52    )5354    print(f"Creating container {CONTAINER_NAME}...", file=sys.stderr)55    run_cmd = [56        "docker", "run", "-d",57        "--name", CONTAINER_NAME,58        "--network", "emblema",59        "--gpus", "all",60        "-e", "PORT=8113",61        "-e", "DEVICE=cuda",62        NATIVE_VIDEO_UPSCALE_SERVICE_IMAGE,63    ]64    result = subprocess.run(run_cmd, capture_output=True, text=True)65    if result.returncode != 0:66        print(f"docker run failed (exit {result.returncode}): {result.stderr}", file=sys.stderr)67        raise RuntimeError(f"Failed to start container: {result.stderr}")6869    # Poll health endpoint (120s — nvidia-vfx bundles models, no download needed)70    timeout = 12071    interval = 372    elapsed = 073    health_url = f"{NATIVE_VIDEO_UPSCALE_SERVICE_URL}/health"74    while elapsed < timeout:75        try:76            r = requests.get(health_url, timeout=5)77            if r.status_code == 200:78                print(f"Container healthy (waited {elapsed}s).", file=sys.stderr)79                return80        except requests.ConnectionError:81            pass82        time.sleep(interval)83        elapsed += interval8485    raise RuntimeError(f"Container did not become healthy within {timeout}s")868788def stop_container():89    """Remove the container."""90    try:91        subprocess.run(92            ["docker", "rm", "-f", CONTAINER_NAME],93            capture_output=True, text=True, timeout=3094        )95        print(f"Container {CONTAINER_NAME} removed.", file=sys.stderr)96    except Exception as e:97        print(f"Warning: failed to remove container: {e}", file=sys.stderr)9899100def main():101    try:102        input_json = sys.stdin.read()103        execution_input = json.loads(input_json)104        inputs = execution_input.get("inputs", {})105106        video = inputs.get("video", "")107        mode = inputs.get("mode", "upscale")108        quality = inputs.get("quality", "high")109        scale_factor = inputs.get("scale_factor", 2.0)110111        if not video:112            raise ValueError("Video input is required")113114        video_path = os.path.join(INPUT_DIR, video)115        if not os.path.exists(video_path):116            raise FileNotFoundError(f"Input video not found: {video_path}")117118        # Map mode + quality to nvvfx QualityLevel int119        quality_int = QUALITY_MAP.get((mode, quality))120        if quality_int is None:121            raise ValueError(f"Invalid mode/quality combination: {mode}/{quality}")122123        # For denoise/deblur, scale_factor is irrelevant (set to 1.0)124        if mode in ("denoise", "deblur"):125            scale_factor = 1.0126127        os.makedirs(OUTPUT_DIR, exist_ok=True)128129        # Start the container130        start_container()131132        try:133            print(134                f"Processing: mode={mode}, quality={quality} (level={quality_int}), "135                f"scale_factor={scale_factor}",136                file=sys.stderr,137            )138139            # Send video to native service140            with open(video_path, "rb") as f:141                resp = requests.post(142                    f"{NATIVE_VIDEO_UPSCALE_SERVICE_URL}/upscale",143                    files={"video": (os.path.basename(video_path), f, "video/mp4")},144                    data={145                        "quality": quality_int,146                        "scale_factor": scale_factor,147                    },148                    timeout=1800,149                )150151            if resp.status_code != 200:152                try:153                    error_detail = resp.json()154                except Exception:155                    error_detail = resp.text[:500]156                raise RuntimeError(157                    f"Native video upscale service returned {resp.status_code}: {error_detail}"158                )159160            # Save output video161            base_name = os.path.splitext(os.path.basename(video))[0]162            out_filename = f"upscaled_{base_name}.mp4"163            out_path = os.path.join(OUTPUT_DIR, out_filename)164            with open(out_path, "wb") as f:165                f.write(resp.content)166167            # Log metadata168            inference_ms = resp.headers.get("X-Inference-Time-Ms", "unknown")169            resolution = resp.headers.get("X-Resolution", "unknown")170            frame_count = resp.headers.get("X-Frame-Count", "unknown")171            print(172                f"Done: resolution={resolution}, frames={frame_count}, "173                f"inference_time={inference_ms}ms",174                file=sys.stderr,175            )176177            # Flat output — keys match OUTPUT_SCHEMA178            output = {179                "video": out_filename,180            }181            print(json.dumps(output, indent=2))182183        finally:184            stop_container()185186    except Exception as e:187        error_output = {188            "error": str(e),189            "errorType": type(e).__name__,190            "traceback": traceback.format_exc(),191        }192        print(json.dumps(error_output), file=sys.stderr)193        sys.exit(1)194195196if __name__ == "__main__":197    main()