Video Upscale (Real-ESRGAN TRT)

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        variant = inputs.get("variant", "realesrgan-x4-plus")20        tile_size = int(inputs.get("tile_size", 1024))2122        if not video:23            raise ValueError("Video input is required")2425        video_path = os.path.join(INPUT_DIR, video)26        if not os.path.exists(video_path):27            raise FileNotFoundError(f"Input video not found: {video_path}")2829        os.makedirs(OUTPUT_DIR, exist_ok=True)3031        print(32            f"Processing: variant={variant}, tile_size={tile_size}",33            file=sys.stderr,34        )3536        result = Gais.video.upscale_realesrgan_trt(37            video=video_path,38            variant=variant,39            tile_size=tile_size,40        )4142        base_name = os.path.splitext(os.path.basename(video))[0]43        out_filename = f"upscaled_{base_name}.mp4"44        out_path = os.path.join(OUTPUT_DIR, out_filename)45        with open(out_path, "wb") as f:46            f.write(result.content)4748        inference_ms = result.metadata.get("inference_time_ms", "unknown")49        out_resolution = result.metadata.get("resolution", "unknown")50        frame_count = result.metadata.get("frame_count", "unknown")51        print(52            f"Done: variant={variant}, resolution={out_resolution}, "53            f"frames={frame_count}, inference_time={inference_ms}ms",54            file=sys.stderr,55        )5657        output = {58            "video": out_filename,59        }60        print(json.dumps(output, indent=2))6162    except Exception as e:63        error_output = {64            "error": str(e),65            "errorType": type(e).__name__,66            "traceback": traceback.format_exc(),67        }68        print(json.dumps(error_output), file=sys.stderr)69        sys.exit(1)707172if __name__ == "__main__":73    main()