Spaces:
Paused
Paused
File size: 4,092 Bytes
92076a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | 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)
|