app.py

"""MediaTok Player — Upload media, encode to .gtkv, play back in browser or via WebGPU."""

import os
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

import gradio as gr
import torch
import numpy as np
import subprocess as sp

from mediatok.container.gtkv import GtkvWriter, GtkvHeader, VideoTokenBlock, GtkvReader
from mediatok.codecs.gigatoken import GigaTokenVideoCodec
from mediatok.codecs.audio import DummyAudioCodec
from mediatok.entropy import entropy_encode
from mediatok.entropy.delta import delta_encode
from mediatok.pipeline.decoder import DecoderPipeline

device = "cpu"
video_codec = GigaTokenVideoCodec(device=device, backend="research")


BBB_ZIP_URL = "https://download.blender.org/demo/movies/BBB/bbb_sunflower_1080p_60fps_normal.mp4.zip"
BBB_FILENAME = "bbb_sunflower_1080p_60fps_normal.mp4"


def download_demo(progress=gr.Progress()) -> str:
    """Download the Big Buck Bunny demo video (zip) and return the extracted mp4 path."""
    import urllib.request, zipfile
    out_dir = tempfile.mkdtemp()
    zip_path = os.path.join(out_dir, "bbb.zip")

    try:
        progress(0, desc="Connecting to Blender server...")
        req = urllib.request.Request(
            BBB_ZIP_URL,
            headers={"User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"},
        )
        with urllib.request.urlopen(req, timeout=600) as src, open(zip_path, "wb") as dst:
            total = int(src.headers.get("Content-Length", 0))
            downloaded = 0
            while True:
                chunk = src.read(65536)
                if not chunk:
                    break
                dst.write(chunk)
                downloaded += len(chunk)
                if total:
                    pct = downloaded / total
                    progress(pct, desc=f"Downloading {downloaded // 1048576}MB / {total // 1048576}MB")
        progress(0.95, desc="Extracting...")
        with zipfile.ZipFile(zip_path, "r") as zf:
            for name in zf.namelist():
                if name.endswith(".mp4"):
                    zf.extract(name, out_dir)
                    mp4_path = os.path.join(out_dir, name)
                    progress(1, desc="Ready")
                    return mp4_path
        raise RuntimeError("No mp4 found in zip")
    except Exception as e:
        raise RuntimeError(f"Download failed: {e}")


def _probe_video(path: str) -> tuple:
    """Probe a video file, return (W, H, fps, total_frames)."""
    probe = sp.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=width,height,r_frame_rate,nb_frames",
         "-of", "csv=p=0", path],
        capture_output=True, text=True, timeout=30,
    )
    parts = probe.stdout.strip().split(",")
    W, H = int(parts[0]), int(parts[1])
    num, den = map(int, parts[2].split("/"))
    fps = num // den if den else 30
    total = int(parts[3]) if len(parts) > 3 and parts[3] else 0
    return W, H, fps, total


def _encode_chunk(frames_tensor, chunk_size) -> bytes:
    """Encode a batch of frames to GigaToken tokens, return entropy payload."""
    chunk = frames_tensor.unsqueeze(0).float() / 255.0  # [1, C, T, H, W]
    with torch.no_grad():
        tokens = video_codec.encode(chunk)
    flat = torch.cat([t.cpu().view(-1).to(torch.int32) for t in tokens]).tolist()
    nf = chunk.shape[2]
    per_frame = [flat[i*444:(i+1)*444] for i in range(nf)]
    deltas = delta_encode(per_frame)
    return entropy_encode(deltas, 2, bits=18), flat


def encode_to_gtkv(video_file: str, progress=gr.Progress()) -> str:
    """Stream-encode a video to .gtkv — never holds all frames in RAM."""
    fd_o, out_path = tempfile.mkstemp(suffix=".gtkv")
    os.close(fd_o)

    try:
        progress(0.1, desc="Probing video...")
        W, H, fps, total_frames = _probe_video(video_file)
        frame_bytes = W * H * 3
        chunk_sz = min(128, max(1, fps))

        header = GtkvHeader(
            num_video_frames=total_frames or 0,
            width=W, height=H, fps=fps,
            chunk_size_frames=chunk_sz,
            num_layers=video_codec.num_layers,
            layer_token_counts=video_codec.layer_token_counts[:6],
        )

        progress(0.2, desc=f"Streaming {W}x{H} {fps}fps...")
        fd_e, err_path = tempfile.mkstemp(suffix=".err")
        os.close(fd_e)

        proc = sp.Popen(
            ["ffmpeg", "-i", video_file, "-f", "rawvideo", "-pix_fmt", "rgb24",
             "-an", "-sn", "-dn", "-"],
            stdout=sp.PIPE, stderr=open(err_path, "wb"),
        )

        writer = GtkvWriter(out_path, header)
        chunk_frames = []
        chunk_idx = 0
        n_total = 0
        try:
            while True:
                try:
                    raw = proc.stdout.read(frame_bytes)
                except ValueError:
                    break
                if not raw or len(raw) < frame_bytes:
                    break
                arr = np.frombuffer(raw, dtype=np.uint8).reshape(H, W, 3).copy()
                chunk_frames.append(torch.tensor(arr, dtype=torch.uint8))
                n_total += 1

                if len(chunk_frames) >= chunk_sz:
                    progress(0.2 + 0.7 * (n_total / max(total_frames, 1)),
                             desc=f"Encoding chunk {chunk_idx+1} ({n_total} frames)...")
                    frames_tensor = torch.stack(chunk_frames, dim=0).permute(3, 0, 1, 2)
                    payload, flat = _encode_chunk(frames_tensor, chunk_sz)
                    block = VideoTokenBlock(
                        token_count=len(flat),
                        layer_sizes=[frames_tensor.shape[1] * frames_tensor.shape[0]],
                        tokens=flat,
                        entropy_payload=payload,
                    )
                    writer.write_chunk(block)
                    chunk_frames = []
                    chunk_idx += 1

            # Flush remaining frames
            if chunk_frames:
                progress(0.9, desc=f"Encoding final chunk ({n_total} frames)...")
                frames_tensor = torch.stack(chunk_frames, dim=0).permute(3, 0, 1, 2)
                payload, flat = _encode_chunk(frames_tensor, len(chunk_frames))
                block = VideoTokenBlock(
                    token_count=len(flat),
                    layer_sizes=[frames_tensor.shape[1] * frames_tensor.shape[0]],
                    tokens=flat,
                    entropy_payload=payload,
                )
                writer.write_chunk(block)

            proc.wait()
            if proc.returncode != 0:
                with open(err_path) as f:
                    raise RuntimeError(f"ffmpeg error: {f.read()[:300]}")

            # Update header
            header.num_video_frames = n_total
            writer.patch_header()
            writer.finalize()
        except:
            try:
                os.unlink(out_path)
            except OSError:
                pass
            raise
        finally:
            os.unlink(err_path)

        progress(1, desc=f"Done — {n_total} frames encoded")
        return out_path
    except Exception as e:
        raise RuntimeError(f"Encode failed: {e}")
    except Exception as e:
        raise RuntimeError(f"Encode failed: {e}")


def play_gtkv_webgpu():
    """Return HTML snippet for the WebGPU player tab."""
    html_path = os.path.join(os.path.dirname(__file__), "index.html")
    with open(html_path) as f:
        return f.read()


with gr.Blocks(title="MediaTok Player") as demo:
    gr.Markdown("# MediaTok Player")
    gr.Markdown("Encode video to `.gtkv` neural token format, or play existing `.gtkv` files.")

    with gr.Tab("Encode"):
        gr.Markdown("Upload a video file to encode it to .gtkv format.")

        with gr.Accordion("🎬 Demo: Big Buck Bunny", open=False):
            gr.Markdown(
                "[Big Buck Bunny](https://peach.blender.org/) is the classic open-source "
                "movie from the Blender Foundation. Download the 1080p60 clip below and "
                "encode it to .gtkv."
            )
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown(
                        f"Source: [`{BBB_FILENAME}`]({BBB_ZIP_URL.replace('.zip', '')})\n\n"
                        f"File size: ~277 MB (zipped) | 1080p @ 60fps | H.264\n\n"
                        f"All BBB variants: [download.blender.org/demo/movies/BBB/](https://download.blender.org/demo/movies/BBB/)"
                    )
                with gr.Column(scale=1):
                    demo_dl_btn = gr.Button("⬇ Download & Encode", variant="secondary")
                    demo_status = gr.Textbox(label="Status", interactive=False)

        with gr.Row():
            with gr.Column():
                video_input = gr.Video(label="Input video", sources=["upload"])
                encode_btn = gr.Button("Encode to .gtkv", variant="primary")
            with gr.Column():
                gtkv_output = gr.File(label="Download .gtkv")

        encode_btn.click(
            fn=encode_to_gtkv,
            inputs=[video_input],
            outputs=[gtkv_output],
        )

        demo_dl_btn.click(
            fn=download_demo,
            inputs=[],
            outputs=[demo_status],
        ).then(
            fn=encode_to_gtkv,
            inputs=[demo_status],
            outputs=[gtkv_output],
        )

    with gr.Tab("Play (WebGPU)"):
        gr.HTML(play_gtkv_webgpu())

    with gr.Tab("Play (Server decode)"):
        gr.Markdown("Upload a .gtkv file to decode and play back.")
        with gr.Row():
            with gr.Column():
                gtkv_input = gr.File(label="Upload .gtkv", file_types=[".gtkv"])
                play_btn = gr.Button("Play", variant="primary")
            with gr.Column():
                video_output = gr.Video(label="Playback")
                play_info = gr.Textbox(label="Info", interactive=False)

        def play_gtkv(file, progress=gr.Progress()):
            if file is None:
                return None, "Upload a .gtkv file first."
            from mediatok.pipeline.decoder import DecoderPipeline

            progress(0.1, desc="Parsing container...")
            reader = GtkvReader(file.name)
            h = reader.header
            progress(0.2, desc=f"Decoding {h.num_video_frames} frames across {reader.num_chunks} chunks...")
            pipeline = DecoderPipeline(reader, video_codec, DummyAudioCodec(device=device), device=device)
            frames = pipeline.decode_all(layer_mask=0b111111)

            progress(0.8, desc="Encoding to mp4 with ffmpeg...")
            video = torch.cat(frames, dim=2)  # [B, C, T, H, W]
            B, C, T, H, W = video.shape
            video = video.squeeze(0).clamp(0, 1)  # [C, T, H, W]
            arr = (video.permute(1, 2, 3, 0).cpu().numpy() * 255).astype(np.uint8)

            out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
            cmd = [
                "ffmpeg", "-y", "-f", "rawvideo",
                "-vcodec", "rawvideo", "-s", f"{W}x{H}",
                "-pix_fmt", "rgb24", "-r", str(h.fps or 30),
                "-i", "-", "-c:v", "libx264", "-preset", "fast",
                "-crf", "23", "-pix_fmt", "yuv420p", out_path,
            ]
            proc = sp.Popen(cmd, stdin=sp.PIPE, stdout=sp.PIPE, stderr=sp.PIPE)
            proc.stdin.write(arr.tobytes())
            proc.stdin.close()
            proc.wait()

            reader.close()
            progress(1, desc="Done")
            return out_path, f"{W}x{H} @ {h.fps}fps, {T} frames"

        play_btn.click(fn=play_gtkv, inputs=[gtkv_input], outputs=[video_output, play_info])


if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), show_error=True, max_file_size="500mb")