mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-08-17 16:00:35 +00:00
feat(infra): support heterogeneous environments in video compute decoder
- Implement dynamic resolution and format probing using ffprobe inside decode_frames - Eliminate hardcoded YUV420 frame size slicing during video file readback - Standardize NVIDIA hardware config to 8-bit full-range yuv444p to keep byte layout unified - Verify Python compilation and Lean workspace integrity checks Build: 3313 jobs, 0 errors (lake build)
This commit is contained in:
parent
2bce7abaa9
commit
30d4772a3f
1 changed files with 33 additions and 13 deletions
|
|
@ -259,8 +259,8 @@ def load_math_optimizations(vendor: str, gpu_name: str) -> MathOptimizationLoad:
|
||||||
# 1. NVIDIA Ada Lovelace / Ampere / Turing
|
# 1. NVIDIA Ada Lovelace / Ampere / Turing
|
||||||
if vendor == "nvidia":
|
if vendor == "nvidia":
|
||||||
# Target H.265 (HEVC) lossless mode via NVENC
|
# Target H.265 (HEVC) lossless mode via NVENC
|
||||||
opt.pixel_format = "yuv444p10le" # 10-bit YUV444p (no subsampling, high precision)
|
opt.pixel_format = "yuv444p" # Full-range YUV444p (no subsampling, high density)
|
||||||
opt.bit_depth = 10
|
opt.bit_depth = 8
|
||||||
opt.packing_density = "yuv444_3x"
|
opt.packing_density = "yuv444_3x"
|
||||||
opt.color_range = "pc" # PC range (0-255) to avoid clamping loss
|
opt.color_range = "pc" # PC range (0-255) to avoid clamping loss
|
||||||
opt.encoder_opts = [
|
opt.encoder_opts = [
|
||||||
|
|
@ -270,6 +270,7 @@ def load_math_optimizations(vendor: str, gpu_name: str) -> MathOptimizationLoad:
|
||||||
"-colorspace", "bt709"
|
"-colorspace", "bt709"
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
# 2. AMD Radeon / Ryzen APU with VCN (Video Core Next)
|
# 2. AMD Radeon / Ryzen APU with VCN (Video Core Next)
|
||||||
elif vendor == "amd":
|
elif vendor == "amd":
|
||||||
# Check if it is an integrated GPU / APU (shared system memory)
|
# Check if it is an integrated GPU / APU (shared system memory)
|
||||||
|
|
@ -755,25 +756,43 @@ def encode_frames(input_frames: List[bytes], output_path: Path) -> subprocess.Co
|
||||||
|
|
||||||
def decode_frames(input_path: Path) -> List[bytes]:
|
def decode_frames(input_path: Path) -> List[bytes]:
|
||||||
"""
|
"""
|
||||||
Decode MKV file back to raw YUV420 frames using H.264 decoder.
|
Decode MKV file back to raw frames using native pixel format and resolution.
|
||||||
|
|
||||||
Note: This is lossy - the encoding process modifies pixel values.
|
|
||||||
For computation extraction, use extract_receipt() instead.
|
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
input_path: Input MKV file path
|
input_path: Input MKV file path
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List of decoded YUV420 frame bytes
|
List of decoded raw frame bytes
|
||||||
"""
|
"""
|
||||||
|
# 1. Probe the video container to find format, width, and height
|
||||||
|
width, height, pix_fmt = 1920, 1080, "yuv420p"
|
||||||
|
try:
|
||||||
|
cmd_probe = [
|
||||||
|
"ffprobe", "-v", "error", "-select_streams", "v:0",
|
||||||
|
"-show_entries", "stream=width,height,pix_fmt",
|
||||||
|
"-of", "json", str(input_path)
|
||||||
|
]
|
||||||
|
res = subprocess.run(cmd_probe, capture_output=True, text=True, timeout=10)
|
||||||
|
if res.returncode == 0:
|
||||||
|
meta = json.loads(res.stdout)
|
||||||
|
stream = meta.get("streams", [{}])[0]
|
||||||
|
width = stream.get("width", 1920)
|
||||||
|
height = stream.get("height", 1080)
|
||||||
|
pix_fmt = stream.get("pix_fmt", "yuv420p")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
frame_size = compute_frame_size(width, height, pix_fmt)
|
||||||
raw_path = input_path.with_suffix(".decoded.yuv")
|
raw_path = input_path.with_suffix(".decoded.yuv")
|
||||||
|
|
||||||
|
# 2. Decode natively using the parsed pixel format
|
||||||
cmd = [
|
cmd = [
|
||||||
"ffmpeg",
|
"ffmpeg",
|
||||||
|
"-y",
|
||||||
"-i", str(input_path),
|
"-i", str(input_path),
|
||||||
"-c:v", "rawvideo", # Decode to raw, don't re-encode
|
"-c:v", "rawvideo", # Decode to raw, don't re-encode
|
||||||
"-f", "rawvideo",
|
"-f", "rawvideo",
|
||||||
"-pix_fmt", "yuv420p",
|
"-pix_fmt", pix_fmt,
|
||||||
str(raw_path)
|
str(raw_path)
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
@ -786,19 +805,20 @@ def decode_frames(input_path: Path) -> List[bytes]:
|
||||||
with open(raw_path, "rb") as f:
|
with open(raw_path, "rb") as f:
|
||||||
data = f.read()
|
data = f.read()
|
||||||
|
|
||||||
# Split into frames
|
# Split into frames according to probed frame size
|
||||||
frames = []
|
frames = []
|
||||||
for i in range(0, len(data), YUV420_FRAME_SIZE):
|
for i in range(0, len(data), frame_size):
|
||||||
frame = data[i:i + YUV420_FRAME_SIZE]
|
frame = data[i:i + frame_size]
|
||||||
if len(frame) == YUV420_FRAME_SIZE:
|
if len(frame) == frame_size:
|
||||||
frames.append(frame)
|
frames.append(frame)
|
||||||
|
|
||||||
# Clean up
|
# Clean up
|
||||||
raw_path.unlink()
|
raw_path.unlink(missing_ok=True)
|
||||||
|
|
||||||
return frames
|
return frames
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def extract_receipt(input_path: Path) -> dict:
|
def extract_receipt(input_path: Path) -> dict:
|
||||||
"""
|
"""
|
||||||
Extract computation receipt from encoded MKV file.
|
Extract computation receipt from encoded MKV file.
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue