PruneVid for aerial drone videos:
The previous article discussed using one of the three supported VLMs for PruneVid to improve token efficiency with the smallest model: PLLaVA. The other two are LLaVA-OneVision and ST-LLM (Liu et al., 2024b).
ST-LLM feeds all raw, unpooled spatial-temporal tokens directly into the LLM backbone. For an egomotion camera where changes between consecutive frames are incremental, preserving the raw spatial-temporal progression allows the LLM to model subtle frame-to-frame displacements, motion cues, and panning trajectories. To prevent the massive context length from slowing down the system, it uses a dynamic masking strategy during training. This means it can drop redundant tokens where frames overlap significantly while forcing the model to retain a highly granular understanding of the continuous stream
The following implementation uses LLaVA-OneVision instead of the PLLaVA and ST-LLM from the previous article:
"""Standalone LLaVA-OneVision aerial-video runner. No PruneVid/ST-LLM imports."""
import argparse
import faulthandler
import json
import logging
import math
from pathlib import Path
import platform
import time
LOG = logging.getLogger('onevision')
DEFAULT_MODEL = 'llava-hf/llava-onevision-qwen2-0.5b-ov-hf'
def sample_indices(total, fps, start, seconds, frames):
if total < 1 or not math.isfinite(fps) or fps <= 0:
raise ValueError('Invalid video frame count or FPS. Transcode to constant-frame-rate MP4.')
if not math.isfinite(start) or not math.isfinite(seconds) or start < 0 or seconds <= 0 or frames < 1:
raise ValueError('Require start >= 0, seconds > 0 and frames >= 1.')
first = int(start * fps)
last = min(total, int((start + seconds) * fps)) - 1
if last - first + 1 < frames:
raise ValueError('Clip is out of range or contains fewer source frames than --frames.')
return [first] if frames == 1 else [round(first + i * (last-first)/(frames-1)) for i in range(frames)]
def save_json(path, value):
temporary = path.with_suffix(path.suffix + '.tmp')
temporary.write_text(json.dumps(value, indent=2), encoding='utf-8')
temporary.replace(path)
def parse_args():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--video', required=True)
p.add_argument('--model-dir', '--model-id', dest='model_dir', default=DEFAULT_MODEL,
help='Complete Hugging Face OneVision checkpoint directory or repository ID.')
p.add_argument('--start', type=float, default=0)
p.add_argument('--seconds', type=float, default=20)
p.add_argument('--frames', type=int, default=16)
p.add_argument('--question', default='Describe the visible terrain, structures, roads, and moving objects. Distinguish camera motion from object motion. State uncertainty; do not invent GPS coordinates or metric distances.')
p.add_argument('--max-new-tokens', type=int, default=160)
p.add_argument('--output', default='runs/drone_onevision')
p.add_argument('--prepare-only', action='store_true')
p.add_argument('--device', choices=['cuda', 'cpu'], default='cuda')
p.add_argument('--dtype', choices=['auto', 'float16', 'bfloat16', 'float32'], default='auto')
p.add_argument('--load-in-4bit', action='store_true')
p.add_argument('--local-files-only', action='store_true')
p.add_argument('--revision', default='main', help='Hub revision; use a commit hash for reproducibility.')
a = p.parse_args()
if a.max_new_tokens < 1 or a.frames < 1 or not a.question.strip():
p.error('Require positive frames/max-new-tokens and a nonempty question.')
if a.device == 'cpu' and a.load_in_4bit:
p.error('This adapter supports bitsandbytes 4-bit loading on CUDA only.')
return a
def prepare(a, target):
import cv2
from PIL import Image, ImageDraw
source = Path(a.video).resolve()
if not source.is_file():
raise FileNotFoundError(source)
cap = cv2.VideoCapture(str(source))
images, samples = [], []
try:
if not cap.isOpened():
raise RuntimeError('OpenCV cannot open the video; check codec and path.')
fps = float(cap.get(cv2.CAP_PROP_FPS))
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
indices = sample_indices(total, fps, a.start, a.seconds, a.frames)
for i, index in enumerate(indices):
cap.set(cv2.CAP_PROP_POS_FRAMES, index)
ok, bgr = cap.read()
if not ok:
raise RuntimeError(f'Failed to decode frame {index}.')
image = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
images.append(image)
image.save(target / f'frame_{i:03d}.jpg')
samples.append({'frame_index': index, 'nominal_seconds': index / fps})
finally:
cap.release()
sheet = Image.new('RGB', (1280, math.ceil(a.frames / 4) * 205), 'white')
draw = ImageDraw.Draw(sheet)
for i, image in enumerate(images):
thumb = image.copy()
thumb.thumbnail((320, 180))
x, y = (i % 4)*320, (i // 4)*205
sheet.paste(thumb, (x, y))
draw.text((x+4, y+182), f'{samples[i]["nominal_seconds"]:.3f}s', fill='black')
sheet.save(target / 'contact_sheet.jpg')
meta = {'arguments': vars(a), 'video': str(source), 'fps': fps,
'source_frame_count': total, 'timestamps_are_nominal': True, 'samples': samples}
save_json(target / 'sampling.json', meta)
return images, meta
def infer(a, images, metadata):
import numpy as np
import torch
import transformers
from transformers import AutoConfig, AutoProcessor, LlavaOnevisionForConditionalGeneration
device = 'cuda:0' if a.device == 'cuda' else 'cpu'
if a.device == 'cuda':
if not torch.cuda.is_available():
raise RuntimeError('CUDA unavailable. Install a CUDA-enabled PyTorch wheel; no implicit CPU fallback.')
LOG.info('GPU: %s; capability: %s; torch: %s; runtime: %s',
torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0), torch.__version__, torch.version.cuda)
if torch.cuda.get_device_capability(0) == (12, 0) and 'sm_120' not in torch.cuda.get_arch_list():
LOG.warning('sm_120 not listed in this wheel. The following real GPU computation must pass.')
x = torch.randn(128, 128, device=device)
_ = x @ x
torch.cuda.synchronize()
del x, _
dtype_name = a.dtype if a.dtype != 'auto' else ('float32' if a.device == 'cpu' else 'float16')
if a.device == 'cpu' and dtype_name != 'float32':
raise ValueError('Use float32 on CPU in this adapter.')
if dtype_name == 'bfloat16' and a.device == 'cuda' and not torch.cuda.is_bf16_supported():
raise RuntimeError('This GPU does not report BF16 support.')
dtype = getattr(torch, dtype_name)
shared = {'revision': a.revision, 'local_files_only': a.local_files_only, 'trust_remote_code': False}
config = AutoConfig.from_pretrained(a.model_dir, **shared)
if config.model_type != 'llava_onevision':
raise ValueError(f'Expected llava_onevision, got {config.model_type!r}; PLLaVA/ST-LLM weights cannot be reused.')
processor = AutoProcessor.from_pretrained(a.model_dir, use_fast=False, **shared)
processor.tokenizer.padding_side = 'left'
options = dict(config=config, torch_dtype=dtype, low_cpu_mem_usage=True,
device_map={'': device}, attn_implementation='sdpa', **shared)
if a.load_in_4bit:
from transformers import BitsAndBytesConfig
options['quantization_config'] = BitsAndBytesConfig(
load_in_4bit=True, bnb_4bit_quant_type='nf4',
bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=dtype)
LOG.info('Loading %s onto %s; no automatic RAM/disk offload.', a.model_dir, device)
started = time.perf_counter()
model = LlavaOnevisionForConditionalGeneration.from_pretrained(a.model_dir, **options).eval()
load_seconds = time.perf_counter()-started
times = ', '.join(f'{row["nominal_seconds"]:.3f}' for row in metadata['samples'])
question = (f'The supplied video consists of {a.frames} sampled frames, in chronological order, '
f'at nominal source-video times [{times}] seconds. Events between samples may be unobserved.\n' + a.question)
messages = [{'role': 'user', 'content': [{'type': 'video'}, {'type': 'text', 'text': question}]}]
prompt = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
clip = np.stack([np.asarray(image) for image in images])
inputs = processor(text=prompt, videos=[clip], padding=True, return_tensors='pt')
inputs = inputs.to(device=device, dtype=dtype)
input_length = inputs['input_ids'].shape[1]
limit = getattr(config.text_config, 'max_position_embeddings', None)
if limit and input_length + a.max_new_tokens > limit:
raise ValueError(f'Context overflow: {input_length}+{a.max_new_tokens}>{limit}; reduce --frames.')
video_token_count = int((inputs['input_ids'] == config.video_token_index).sum().item())
if a.device == 'cuda':
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
started = time.perf_counter()
with torch.inference_mode():
generated = model.generate(**inputs, max_new_tokens=a.max_new_tokens,
do_sample=False, num_beams=1, use_cache=True)
if a.device == 'cuda':
torch.cuda.synchronize()
elapsed = time.perf_counter()-started
answer = processor.batch_decode(generated[:, input_length:], skip_special_tokens=True,
clean_up_tokenization_spaces=False)[0].strip()
return {**metadata, 'answer': answer, 'model_type': config.model_type,
'resolved_model_commit': getattr(config, '_commit_hash', None),
'model_load_seconds': load_seconds, 'generation_seconds': elapsed,
'input_tokens': input_length, 'video_placeholder_tokens': video_token_count,
'generated_tokens': generated.shape[1]-input_length,
'input_shapes': {k:list(v.shape) for k,v in inputs.items() if hasattr(v, 'shape')},
'peak_allocated_gpu_gib': torch.cuda.max_memory_allocated()/2**30 if a.device == 'cuda' else None,
'gpu': torch.cuda.get_device_name(0) if a.device == 'cuda' else None,
'torch_version': torch.__version__, 'transformers_version': transformers.__version__,
'python_version': platform.python_version(), 'resolved_dtype': dtype_name,
'note': 'Standard OneVision video inference, not PruneVid pruning or GeoVARNet evidence reasoning. Timestamps are frame-index/FPS approximations.'}
def main():
faulthandler.enable()
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
a = parse_args()
target = Path(a.output)
target.mkdir(parents=True, exist_ok=True)
try:
started = time.perf_counter()
images, metadata = prepare(a, target)
preparation_seconds = time.perf_counter()-started
if a.prepare_only:
LOG.info('Prepared %d frames, sampling.json and contact sheet; no model loaded.', len(images))
return
result = infer(a, images, metadata)
result['preparation_seconds'] = preparation_seconds
result['total_seconds'] = time.perf_counter()-started
save_json(target / 'result.json', result)
print(result['answer'])
except Exception as error:
save_json(target / 'error.json', {'error_type': type(error).__name__, 'message': str(error), 'arguments': vars(a)})
LOG.exception('Run failed. No fallback answer was generated.')
raise
if __name__ == '__main__':
main()
Some caveat remain as before:
1. Libraries I used:
# Approximate environment for run_onevision_drone.py
# Use in a clean environment. GPU inference has not been validated here.
# --extra-index-url https://download.pytorch.org/whl/cu128
torch==2.7.1+cu128
torchvision==0.22.1+cu128
transformers==4.51.3
accelerate>=1.6,<2
huggingface-hub>=0.30,<1
numpy>=1.26,<2
opencv-python==4.9.0.80
pillow>=10,<12
sentencepiece>=0.2,<0.3
protobuf>=4.25,<6
safetensors>=0.4.3,<1
# Optional: uncomment for --load-in-4bit (not needed for the default 0.5B run).
bitsandbytes>=0.47,<1
2.
There is a checkpoint/code compatibility issue.
prunevid is not a generic VLM wrapper. It is tightly coupled to the downloaded checkpoint.
and there are three checkpoints involved which is suspicious:
LlaVa-OneVision
llava-v1.6-vicuna-7b-hf
llava-1.5-7b-hf
The checkpoint architecture and the code architecture do not match.
1. Last known state:
Environment built successfully.
Model weights downloaded successfully.
Preprocessing works.
Model loading partially works.
Current blocker = intermittent native crash LlaVA-OneVision compatibility issue.
