Video Upscale

1import os2import sys3import json4import traceback56from gais import Gais78INPUT_DIR = "/data/input"9OUTPUT_DIR = "/data/output"101112def main():13    try:14        input_json = sys.stdin.read()15        execution_input = json.loads(input_json)16        inputs = execution_input.get("inputs", {})1718        video = inputs.get("video", "")19        mode = inputs.get("mode", "upscale")20        scale_factor = inputs.get("scale_factor", 2.0)21        sharpness = inputs.get("sharpness", "balanced")2223        if not video:24            raise ValueError("Video input is required")2526        video_path = os.path.join(INPUT_DIR, video)27        if not os.path.exists(video_path):28            raise FileNotFoundError(f"Input video not found: {video_path}")2930        # For denoise/deblur, scale_factor is irrelevant (set to 1.0)31        if mode in ("denoise", "deblur"):32            scale_factor = 1.03334        os.makedirs(OUTPUT_DIR, exist_ok=True)3536        print(37            f"Processing: mode={mode}, scale_factor={scale_factor}, sharpness={sharpness}",38            file=sys.stderr,39        )4041        result = Gais.video.upscale(42            video=video_path,43            mode=mode,44            scale_factor=scale_factor,45            sharpness=sharpness,46        )4748        # Save output video49        base_name = os.path.splitext(os.path.basename(video))[0]50        out_filename = f"upscaled_{base_name}.mp4"51        out_path = os.path.join(OUTPUT_DIR, out_filename)52        with open(out_path, "wb") as f:53            f.write(result.content)5455        # Log metadata56        inference_ms = result.metadata.get("inference_time_ms", "unknown")57        out_resolution = result.metadata.get("resolution", "unknown")58        frame_count = result.metadata.get("frame_count", "unknown")59        print(60            f"Done: resolution={out_resolution}, frames={frame_count}, "61            f"inference_time={inference_ms}ms",62            file=sys.stderr,63        )6465        # Probe output dimensions and emit aspect_ratio + resolution so66        # downstream nodes can chain.67        try:68            import av69            with av.open(out_path) as _c:70                _s = next(s for s in _c.streams if s.type == "video")71                _w = int(_s.codec_context.width)72                _h = int(_s.codec_context.height)73        except Exception:74            _w, _h = 0, 07576        # Flat output — keys match OUTPUT_SCHEMA77        output = {78            "video": out_filename,79            "aspect_ratio": round(_w / _h, 4) if _h else 0.0,80            "resolution": f"{_w}x{_h}",81        }82        print(json.dumps(output, indent=2))8384    except Exception as e:85        error_output = {86            "error": str(e),87            "errorType": type(e).__name__,88            "traceback": traceback.format_exc(),89        }90        print(json.dumps(error_output), file=sys.stderr)91        sys.exit(1)929394if __name__ == "__main__":95    main()