import argparse import sys from pathlib import Path import torch from mediatok.container.gtkv import GtkvReader from mediatok.codecs.video import DummyVideoCodec from mediatok.codecs.gigatoken import GigaTokenVideoCodec from mediatok.codecs.audio import EnCodecAudioCodec, DummyAudioCodec from mediatok.pipeline import EncoderPipeline, DecoderPipeline from mediatok.playback import Player def _xpu_usable() -> bool: """Level Zero may enumerate a stub device in WSL2 without GPU render nodes (/dev/dri). Verify the OS actually exposes a GPU before committing to the XPU device, otherwise tensor creation segfaults in the driver.""" import os return torch.xpu.is_available() and os.path.exists("/dev/dri") def _detect_device() -> str: if _xpu_usable(): return "xpu" return "cpu" def cmd_encode(args): device = _detect_device() if args.dummy: vc = DummyVideoCodec(device=device) ac = DummyAudioCodec(device=device) else: vc = GigaTokenVideoCodec(device=device) ac = EnCodecAudioCodec(device=device) pipeline = EncoderPipeline(vc, ac, chunk_size_frames=args.chunk_frames, entropy_codec=args.entropy) pipeline.encode_file( video_path=args.input, audio_path=args.audio or "", output_path=args.output, width=args.width, height=args.height, fps=args.fps, ) print(f"encoded {args.input} -> {args.output}") def cmd_decode(args): device = _detect_device() reader = GtkvReader(args.input) if args.dummy: vc = DummyVideoCodec(device=device) ac = DummyAudioCodec(device=device) else: vc = GigaTokenVideoCodec(device=device) ac = EnCodecAudioCodec(device=device) pipeline = DecoderPipeline(reader, vc, ac, device=device) frames = pipeline.decode_all(layer_mask=args.layers) if frames: import torch out = torch.cat(frames, dim=2) # save with torch.save for now (conversion to mp4 via ffmpeg) torch.save(out, args.output) print(f"decoded {args.input} -> {args.output}") def cmd_play(args): device = _detect_device() reader = GtkvReader(args.input) if args.dummy: vc = DummyVideoCodec(device=device) else: vc = GigaTokenVideoCodec(device=device) pipeline = DecoderPipeline(reader, vc, DummyAudioCodec(device=device), device=device) from mediatok.playback.display import TkPlayer player = TkPlayer(pipeline, layer_mask=args.layers) player.play() reader.close() def cmd_info(args): reader = GtkvReader(args.input) h = reader.header print(f"GTKV Container: {args.input}") print(f" Version: {h.version}") print(f" Frames: {h.num_video_frames}") print(f" Resolution: {h.width}x{h.height}") print(f" FPS: {h.fps}") print(f" Audio: {h.audio_sample_rate} Hz") print(f" Chunk size: {h.chunk_size_frames} frames") print(f" Num chunks: {reader.num_chunks}") print(f" Layers: {h.num_layers} ({h.layer_token_counts})") print(f" Entropy: {['none','rans','zstd','arithmetic'][h.entropy_codec_id]}") print(f" Video tokenizer: {['gigatoken','cosmos','open_magvit2','vidtok'][h.video_tokenizer_id]}") print(f" Audio tokenizer: {['encodec','dac'][h.audio_tokenizer_id]}") reader.close() def cmd_seek(args): device = _detect_device() reader = GtkvReader(args.input) chunk = reader.header.chunk_size_frames ci = args.frame // chunk fi = args.frame % chunk if args.dummy: vc = DummyVideoCodec(device=device) else: vc = GigaTokenVideoCodec(device=device) pipeline = DecoderPipeline(reader, vc, DummyAudioCodec(device=device), device=device) frames = pipeline.decode_chunk(ci, layer_mask=args.layers) if frames.shape[2] > fi: from PIL import Image import numpy as np frame = frames[:, :, fi].cpu() arr = frame.squeeze(0).permute(1, 2, 0).numpy() arr = ((arr - arr.min()) / (arr.max() - arr.min() + 1e-8) * 255).astype(np.uint8) img = Image.fromarray(arr) out_path = args.output or f"frame_{args.frame}.png" img.save(out_path) print(f"saved frame {args.frame} -> {out_path}") reader.close() def cmd_bench(args): import time import numpy as np device = _detect_device() reader = GtkvReader(args.input) if args.dummy: vc = DummyVideoCodec(device=device) else: vc = GigaTokenVideoCodec(device=device) ac = DummyAudioCodec(device=device) pipeline = DecoderPipeline(reader, vc, ac, device=device) stages = args.stages.split(",") times = {} n = min(reader.num_chunks, args.num_chunks) if "all" in stages or "read" in stages: t0 = time.perf_counter() for i in range(n): reader.read_video_block(i) times["read"] = (time.perf_counter() - t0) / n if "all" in stages or "decode" in stages: t0 = time.perf_counter() for i in range(n): pipeline.decode_chunk(i, layer_mask=args.layers) times["decode"] = (time.perf_counter() - t0) / n if "all" in stages or "transfer" in stages and device != "cpu": t0 = time.perf_counter() dummy = torch.randint(0, 1000, (1, 256), dtype=torch.int64) for _ in range(n): dummy.to(device) times["transfer"] = (time.perf_counter() - t0) / n print(f"Benchmark: {args.input} ({n} chunks, device={device})") for k, v in times.items(): print(f" {k}: {v*1000:.2f} ms") reader.close() def main(): parser = argparse.ArgumentParser(description="MediaTok — token-native media engine") parser.add_argument("--dummy", action="store_true", help="use dummy codecs (no GPU model required)") sub = parser.add_subparsers(dest="command") p_encode = sub.add_parser("encode") p_encode.add_argument("input") p_encode.add_argument("output") p_encode.add_argument("--audio", default=None) p_encode.add_argument("--width", type=int, default=1920) p_encode.add_argument("--height", type=int, default=1080) p_encode.add_argument("--fps", type=int, default=30) p_encode.add_argument("--chunk-frames", type=int, default=128) p_encode.add_argument("--entropy", default="rans") p_decode = sub.add_parser("decode") p_decode.add_argument("input") p_decode.add_argument("output") p_decode.add_argument("--layers", type=int, default=0b111111) p_decode.add_argument("--dummy", action="store_true") p_play = sub.add_parser("play") p_play.add_argument("input") p_play.add_argument("--layers", type=int, default=0b111111) p_play.add_argument("--dummy", action="store_true") p_info = sub.add_parser("info") p_info.add_argument("input") p_seek = sub.add_parser("seek") p_seek.add_argument("input") p_seek.add_argument("--frame", type=int, required=True) p_seek.add_argument("--output", default=None) p_seek.add_argument("--layers", type=int, default=0b111111) p_seek.add_argument("--dummy", action="store_true") p_bench = sub.add_parser("bench") p_bench.add_argument("input") p_bench.add_argument("--stages", default="all") p_bench.add_argument("--layers", type=int, default=0b111111) p_bench.add_argument("--num-chunks", type=int, default=10) p_bench.add_argument("--dummy", action="store_true") args = parser.parse_args() if not args.command: parser.print_help() return commands = { "encode": cmd_encode, "decode": cmd_decode, "play": cmd_play, "info": cmd_info, "seek": cmd_seek, "bench": cmd_bench, } commands[args.command](args) if __name__ == "__main__": main()