cli.py

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()