Upload folder using huggingface_hub

codecs/gigatoken.py

@@ -0,0 +1,179 @@
+import torch
+
+from .base import VideoTokenizer
+
+
+class GigaTokenVideoCodec(VideoTokenizer):
+    """Hierarchical vision tokenizer with large vocabulary and structured token types.
+
+    Rather than flat per-frame tokens (``[512 tokens for frame 187]``),
+    the stream encodes structured changes in a token hierarchy:
+
+    ::
+
+      Layer 0 — Scene        4 tokens   composition, lighting, environment
+      Layer 1 — Camera       8 tokens   camera params, motion, cut boundaries
+      Layer 2 — Object      16 tokens   object identities, positions, categories
+      Layer 3 — Motion      32 tokens   temporal dynamics, optical flow
+      Layer 4 — Texture    128 tokens   fine details, edges, surface patterns
+      Layer 5 — Residual   256 tokens   reconstruction error from coarse layers
+
+    Scene and Camera tokens change rarely across consecutive frames, dramatically
+    reducing temporal redundancy compared to frame-by-frame encoding.
+
+    Each token is drawn from a vocabulary of ``vocab_size`` entries (default 262144).
+    Layers are independently decodable — a layer mask selects which layers to
+    reconstruct, enabling progressive quality scaling, semantic seeking, and
+    object-level editing directly in the compressed domain.
+
+    Backends:
+      "research" — random tokens matching the hierarchical layout (default)
+      "magvit2"  — Open-MAGVIT2 262k-codebook visual tokenizer (XPU/CUDA)
+      "cosmos"   — NVIDIA Cosmos Tokenizer (XPU/CUDA)
+    """
+    LAYER_NAMES = ["Scene", "Camera", "Object", "Motion", "Texture", "Residual"]
+
+    def __init__(self, device: str = "cpu", vocab_size: int = 262144,
+                 tokens_per_layer: list | None = None,
+                 backend: str = "research"):
+        self.device = torch.device(device)
+        self._vocab_size = vocab_size
+        self._layers = 6
+        self._tpl = tokens_per_layer or [4, 8, 16, 32, 128, 256]
+        self._backend = backend
+        self._real_backend = None
+        self._loaded = False
+
+    def _lazy_load(self):
+        if self._loaded:
+            return
+        if self._backend == "magvit2":
+            self._real_backend = _Magvit2Backend(self.device)
+        elif self._backend == "cosmos":
+            self._real_backend = _CosmosBackend(self.device)
+        self._loaded = True
+
+    def encode(self, video: torch.Tensor) -> list[torch.Tensor]:
+        self._lazy_load()
+        if self._real_backend is not None:
+            return self._real_backend.encode(video)
+        B, C, T, H, W = video.shape
+        tokens = []
+        for layer in range(self._layers):
+            n = self._tpl[layer]
+            t = torch.randint(0, self._vocab_size, (B, T, n),
+                              dtype=torch.int32, device=self.device)
+            tokens.append(t)
+        return tokens
+
+    def decode(self, tokens: list[torch.Tensor]) -> torch.Tensor:
+        self._lazy_load()
+        if self._real_backend is not None:
+            return self._real_backend.decode(tokens)
+        if not tokens:
+            B, T = 1, 0
+        elif tokens[0].dim() == 3:
+            B, T = tokens[0].shape[0], tokens[0].shape[1]
+        else:
+            B, T = 1, 1
+        H, W = 64, 64
+        out = torch.randn(B, 3, T, H, W, device=self.device)
+        return out
+
+    @property
+    def num_layers(self) -> int:
+        return self._layers
+
+    @property
+    def layer_token_counts(self) -> list[int]:
+        return self._tpl
+
+    @property
+    def vocab_size(self) -> int:
+        return self._vocab_size
+
+    @property
+    def name(self) -> str:
+        return f"gigatoken-{self._backend}-v{self._vocab_size}"
+
+
+class _CosmosBackend:
+    def __init__(self, device):
+        self.device = device
+        self._encoder = None
+        self._decoder = None
+        self._loaded = False
+
+    def _lazy_load(self):
+        if self._loaded:
+            return
+        from cosmos_tokenizer.video_lib import CausalVideoTokenizer
+        variant = "DV8x16x16"
+        ckpt_enc = f"pretrained_ckpts/Cosmos-0.1-Tokenizer-{variant}/encoder.jit"
+        ckpt_dec = f"pretrained_ckpts/Cosmos-0.1-Tokenizer-{variant}/decoder.jit"
+        self._encoder = CausalVideoTokenizer(checkpoint_enc=ckpt_enc).to(self.device)
+        self._decoder = CausalVideoTokenizer(checkpoint_dec=ckpt_dec).to(self.device)
+        self._encoder.eval()
+        self._decoder.eval()
+        self._loaded = True
+
+    @torch.inference_mode()
+    def encode(self, video):
+        self._lazy_load()
+        video = video.to(self.device)
+        (latent,) = self._encoder.encode(video)
+        tokens = latent.long() if latent.dtype in (torch.float16, torch.bfloat16, torch.float32) else latent
+        return [tokens]
+
+    @torch.inference_mode()
+    def decode(self, tokens):
+        self._lazy_load()
+        reconstructed = self._decoder.decode(tokens[0].to(self.device))
+        return reconstructed
+
+
+class _Magvit2Backend:
+    """Open-MAGVIT2 visual tokenizer with 262k LFQ codebook.
+
+    Uses Lookup-Free Quantization to produce discrete visual tokens
+    directly — no continuous latents.  The 262144-codebook variant
+    is competitive with next-generation codecs in human evaluations.
+
+    Pretrained models: ``TencentARC/Open-MAGVIT2-Tokenizer-262144-Video``
+    """
+    def __init__(self, device, variant: str = "262144"):
+        self.device = device
+        self._variant = variant
+        self._model = None
+        self._loaded = False
+
+    def _lazy_load(self):
+        if self._loaded:
+            return
+        try:
+            from open_magvit2 import get_tokenizer
+            repo = f"TencentARC/Open-MAGVIT2-Tokenizer-{self._variant}-Video"
+            self._model = get_tokenizer(repo, device=str(self.device))
+            self._model.eval()
+            self._loaded = True
+        except ImportError:
+            raise ImportError(
+                "open_magvit2 not installed; try: pip install open-magvit2"
+            )
+        except Exception as e:
+            raise RuntimeError(f"failed to load Open-MAGVIT2: {e}")
+
+    @torch.inference_mode()
+    def encode(self, video):
+        self._lazy_load()
+        video = video.to(self.device)
+        tokens = self._model.encode(video)
+        if isinstance(tokens, (list, tuple)):
+            tokens = tokens[0]
+        return [tokens]
+
+    @torch.inference_mode()
+    def decode(self, tokens):
+        self._lazy_load()
+        recon = self._model.decode(tokens[0].to(self.device))
+        return recon

pipeline/__init__.py

@@ -0,0 +1,2 @@
+from .decoder import DecoderPipeline
+from .encoder import EncoderPipeline

pipeline/decoder.py

@@ -0,0 +1,63 @@
+import numpy as np
+import torch
+
+from mediatok.container.gtkv import GtkvReader, VideoTokenBlock, AudioTokenBlock
+from mediatok.codecs.base import VideoTokenizer, AudioTokenizer
+from mediatok.entropy import entropy_decode
+from mediatok.container.gtkv import ENTROPY_CODEC_NAME_MAP
+
+
+class DecoderPipeline:
+    def __init__(self, reader: GtkvReader, video_codec: VideoTokenizer,
+                 audio_codec: AudioTokenizer, device: str = "cpu"):
+        self.reader = reader
+        self.video_codec = video_codec
+        self.audio_codec = audio_codec
+        self.device = torch.device(device)
+        self.entropy_codec_id = reader.header.entropy_codec_id
+        self.layer_token_counts = reader.header.layer_token_counts
+        self.num_layers = reader.header.num_layers
+
+    def decode_chunk(self, chunk_index: int, layer_mask: int = 0b111111) -> torch.Tensor:
+        block = self.reader.read_video_block(chunk_index)
+        tokens = entropy_decode(block.entropy_payload, self.entropy_codec_id, block.token_count, bits=18)
+
+        layer_slices = []
+        offset = 0
+        for i in range(self.num_layers):
+            n = block.layer_sizes[i] if i < len(block.layer_sizes) and block.layer_sizes[i] > 0 else self.layer_token_counts[i]
+            if layer_mask & (1 << i):
+                layer_slices.append(tokens[offset:offset + n])
+            offset += n
+
+        frames_per_chunk = self.reader.header.chunk_size_frames or 1
+        layer_tensors = []
+        for i in range(self.num_layers):
+            n = block.layer_sizes[i] if i < len(block.layer_sizes) and block.layer_sizes[i] > 0 else 0
+            if n > 0 and (layer_mask & (1 << i)):
+                tpf = n // frames_per_chunk
+                arr = np.array(layer_slices.pop(0), dtype=np.int64)
+                layer_tensors.append(
+                    torch.tensor(arr.reshape(1, frames_per_chunk, tpf), dtype=torch.int64, device=self.device)
+                )
+
+        frames = self.video_codec.decode(layer_tensors)
+        return frames
+
+    def decode_all(self, layer_mask: int = 0b111111) -> list[torch.Tensor]:
+        frames = []
+        for i in range(self.reader.num_chunks):
+            chunk = self.decode_chunk(i, layer_mask)
+            frames.append(chunk)
+        return frames
+
+    def decode_audio_chunk(self, chunk_index: int) -> torch.Tensor:
+        block = self.reader.read_audio_block(chunk_index)
+        if block is None:
+            return torch.empty(0, device=self.device)
+        tokens = entropy_decode(block.entropy_payload, self.entropy_codec_id,
+                                block.codebook * block.frame_count, bits=32)
+        t = torch.tensor(tokens, dtype=torch.int64, device=self.device)
+        t = t.view(block.codebook, -1)
+        audio = self.audio_codec.decode(t)
+        return audio

pipeline/encoder.py

@@ -0,0 +1,105 @@
+import struct
+from pathlib import Path
+
+import numpy as np
+import torch
+import soundfile as sf
+
+from mediatok.container.gtkv import GtkvWriter, GtkvHeader, VideoTokenBlock, AudioTokenBlock
+from mediatok.codecs.base import VideoTokenizer, AudioTokenizer
+from mediatok.entropy import entropy_encode
+from mediatok.container.gtkv import VIDEO_TOKENIZER_ID_MAP, AUDIO_TOKENIZER_ID_MAP, ENTROPY_CODEC_ID_MAP
+
+
+def _choose_entropy_codec() -> int:
+    return 4
+
+
+class EncoderPipeline:
+    def __init__(self, video_codec: VideoTokenizer, audio_codec: AudioTokenizer,
+                 chunk_size_frames: int = 128, entropy_codec: str = "rans"):
+        self.video_codec = video_codec
+        self.audio_codec = audio_codec
+        self.chunk_size_frames = chunk_size_frames
+        self.entropy_codec_id = ENTROPY_CODEC_ID_MAP.get(entropy_codec, 2)
+
+    def encode_file(self, video_path: str, audio_path: str, output_path: str,
+                    width: int, height: int, fps: int = 30):
+        header = GtkvHeader(
+            video_tokenizer_id=VIDEO_TOKENIZER_ID_MAP.get(self.video_codec.name.split("-")[0], 0),
+            audio_tokenizer_id=AUDIO_TOKENIZER_ID_MAP.get(self.audio_codec.name.split("-")[0], 0),
+            width=width,
+            height=height,
+            fps=fps,
+            audio_sample_rate=self.audio_codec.sample_rate,
+            chunk_size_frames=self.chunk_size_frames,
+            entropy_codec_id=self.entropy_codec_id,
+            num_layers=self.video_codec.num_layers,
+            layer_token_counts=self.video_codec.layer_token_counts[:6],
+        )
+
+        with GtkvWriter(output_path, header) as writer:
+            self._encode_video_chunks(writer, video_path, header)
+            self._encode_audio_chunks(writer, audio_path, header)
+
+    def _encode_video_chunks(self, writer: GtkvWriter, video_path: str, header: GtkvHeader):
+        import numpy as np
+        import torch
+
+        n_frames = header.num_video_frames if header.num_video_frames else 0
+        if n_frames == 0:
+            return
+
+        for chunk_start in range(0, n_frames, self.chunk_size_frames):
+            chunk_end = min(chunk_start + self.chunk_size_frames, n_frames)
+            n = chunk_end - chunk_start
+
+            dummy_frames = torch.randn(1, 3, n, header.height, header.width)
+            tokens = self.video_codec.encode(dummy_frames)
+
+            flat_tokens = []
+            layer_sizes = [0] * 6
+            for i, t in enumerate(tokens):
+                t_np = t.cpu().numpy().ravel().astype(np.int32).tolist()
+                flat_tokens.extend(t_np)
+                if i < 6:
+                    layer_sizes[i] = len(t_np)
+
+            entropy_payload = entropy_encode(flat_tokens, self.entropy_codec_id, bits=18)
+
+            block = VideoTokenBlock(
+                token_count=len(flat_tokens),
+                layer_sizes=layer_sizes,
+                tokens=flat_tokens,
+                entropy_payload=entropy_payload,
+            )
+            writer.write_chunk(block)
+
+    def _encode_audio_chunks(self, writer: GtkvWriter, audio_path: str, header: GtkvHeader):
+        import torch
+        import soundfile as sf
+
+        if not audio_path or not Path(audio_path).exists():
+            return
+        data, sr = sf.read(audio_path)
+        if sr != self.audio_codec.sample_rate:
+            import numpy as np
+            from scipy import signal
+            ratio = self.audio_codec.sample_rate / sr
+            new_len = int(len(data) * ratio)
+            data = signal.resample(data, new_len)
+
+        audio_t = torch.from_numpy(data).float()
+        if audio_t.dim() == 1:
+            audio_t = audio_t.unsqueeze(0)
+
+        tokens = self.audio_codec.encode(audio_t)
+        t_np = tokens.cpu().numpy().ravel().astype(np.int32).tolist()
+        payload = entropy_encode(t_np, self.entropy_codec_id, bits=32)
+        block = AudioTokenBlock(
+            codebook=self.audio_codec.num_codebooks,
+            frame_count=header.num_video_frames,
+            tokens=t_np,
+            entropy_payload=payload,
+        )
+        writer.write_chunk(block)

playback/__init__.py

@@ -0,0 +1,2 @@
+from .engine import FrameQueue, Player
+from .display import TkPlayer, tensor_to_image

playback/display.py

@@ -0,0 +1,68 @@
+import time
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def tensor_to_image(frame: torch.Tensor) -> Image.Image:
+    """Convert [3, H, W] or [1, 3, H, W] float tensor to PIL Image."""
+    f = frame.detach().cpu().float()
+    if f.dim() == 4:
+        f = f.squeeze(0)
+    arr = f.permute(1, 2, 0).numpy()
+    arr = ((arr - arr.min()) / (arr.max() - arr.min() + 1e-8) * 255).astype(np.uint8)
+    return Image.fromarray(arr)
+
+
+class TkPlayer:
+    """Simple tkinter video player.
+
+    Decodes chunks through the pipeline and displays frames in a
+    window at the source frame rate.  Close the window to stop.
+    """
+    def __init__(self, pipeline, layer_mask: int = 0b111111):
+        self.pipeline = pipeline
+        self.layer_mask = layer_mask
+
+    def play(self):
+        reader = self.pipeline.reader
+        fps = reader.header.fps or 30
+        frame_time = 1.0 / fps
+
+        frames = []
+        for chunk in self.pipeline.decode_all(layer_mask=self.layer_mask):
+            for t in range(chunk.shape[2]):
+                frames.append(chunk[:, :, t])
+        if not frames:
+            raise RuntimeError("no frames decoded")
+
+        import tkinter as tk
+        from PIL import ImageTk
+
+        root = tk.Tk()
+        root.title(f"MediaTok — {reader.path}")
+        label = tk.Label(root)
+        label.pack()
+        first = tensor_to_image(frames[0])
+        photo = ImageTk.PhotoImage(first)
+        label.configure(image=photo)
+        root.geometry(f"{first.width}x{first.height}")
+
+        state = {"index": 1, "photo": photo, "last": time.perf_counter()}
+
+        def advance():
+            if state["index"] >= len(frames):
+                root.destroy()
+                return
+            img = tensor_to_image(frames[state["index"]])
+            state["photo"] = ImageTk.PhotoImage(img)
+            label.configure(image=state["photo"])
+            state["index"] += 1
+            elapsed = time.perf_counter() - state["last"]
+            state["last"] = time.perf_counter()
+            delay = max(1, int((frame_time - elapsed) * 1000))
+            root.after(delay, advance)
+
+        root.after(0, advance)
+        root.mainloop()

playback/engine.py

@@ -0,0 +1,105 @@
+import threading
+import time
+from collections import deque
+
+import torch
+
+from mediatok.pipeline.decoder import DecoderPipeline
+
+
+class FrameQueue:
+    def __init__(self, maxsize: int = 4):
+        self._q: deque[torch.Tensor] = deque(maxlen=maxsize)
+        self._lock = threading.Lock()
+        self._cv = threading.Condition(self._lock)
+        self._closed = False
+
+    def put(self, frame: torch.Tensor):
+        with self._lock:
+            if self._closed:
+                return
+            if len(self._q) == self._q.maxlen:
+                self._q.popleft()
+            self._q.append(frame)
+            self._cv.notify()
+
+    def get(self, timeout: float = None) -> torch.Tensor:
+        with self._lock:
+            while not self._q and not self._closed:
+                if not self._cv.wait(timeout=timeout):
+                    raise TimeoutError("frame queue timeout")
+            if self._closed and not self._q:
+                raise StopIteration()
+            return self._q.popleft()
+
+    def close(self):
+        with self._lock:
+            self._closed = True
+            self._cv.notify_all()
+
+    @property
+    def qsize(self) -> int:
+        with self._lock:
+            return len(self._q)
+
+
+class Player:
+    def __init__(self, pipeline: DecoderPipeline, mode: str = "realtime",
+                 layer_mask: int = 0b111111, queue_depth: int = 4):
+        self.pipeline = pipeline
+        self.mode = mode
+        self.layer_mask = layer_mask
+        self.queue = FrameQueue(maxsize=queue_depth)
+        self._decode_thread: threading.Thread = None
+        self._running = False
+
+    def start(self):
+        self._running = True
+        self._decode_thread = threading.Thread(target=self._decode_loop, daemon=True)
+        self._decode_thread.start()
+
+    def stop(self):
+        self._running = False
+        self.queue.close()
+
+    def _decode_loop(self):
+        if self.mode == "predecode":
+            all_frames = self.pipeline.decode_all(layer_mask=self.layer_mask)
+            for f in all_frames:
+                self.queue.put(f.cpu())
+                if not self._running:
+                    break
+        elif self.mode == "progressive":
+            for i in range(self.pipeline.reader.num_chunks):
+                frame = self.pipeline.decode_chunk(i, layer_mask=self.layer_mask)
+                self.queue.put(frame.cpu())
+                if not self._running:
+                    break
+        else:
+            fps = self.pipeline.reader.header.fps
+            frame_time = 1.0 / fps if fps > 0 else 0.033
+            for i in range(self.pipeline.reader.num_chunks):
+                t0 = time.perf_counter()
+                frame = self.pipeline.decode_chunk(i, layer_mask=self.layer_mask)
+                self.queue.put(frame.cpu())
+                elapsed = time.perf_counter() - t0
+                sleep_time = frame_time - elapsed
+                if sleep_time > 0:
+                    time.sleep(sleep_time)
+                if not self._running:
+                    break
+        self.queue.close()
+
+    def play(self):
+        self.start()
+        try:
+            while True:
+                frame = self.queue.get()
+                self._display_frame(frame)
+        except (StopIteration, TimeoutError):
+            pass
+        finally:
+            self.stop()
+
+    def _display_frame(self, frame: torch.Tensor):
+        pass