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 ST-LLM instead of the PLLaVA from the previous article:

"""Single-clip adapter for Visual-AI/PruneVid. Run from its repository root."""

"""ST-LLM official-API adapter; run from the ST-LLM repository root."""

import argparse

import json

from pathlib import Path

import time

def main():

    parser = argparse.ArgumentParser(description=__doc__)

    parser.add_argument('--video', required=True, help='Already-trimmed video clip')

    parser.add_argument('--ckpt-path', required=True)

    parser.add_argument('--qformer-path', required=True)

    parser.add_argument('--cfg-path', default='config/instructblipbase_stllm_conversation.yaml')

    parser.add_argument('--frames', type=int, default=16)

    parser.add_argument('--video-input', choices=['all', 'residual'], default='all')

    parser.add_argument('--max-new-tokens', type=int, default=160)

    parser.add_argument('--question', default='Describe terrain, structures, roads and moving objects. Distinguish camera motion from object motion. Do not invent GPS coordinates or distances.')

    parser.add_argument('--output', default='runs/drone_stllm.json')

    a = parser.parse_args()

    if a.frames < 16 or a.max_new_tokens < 1:

        parser.error('Use frames >=16 and max-new-tokens >=1.')

    for path in [a.video, a.qformer_path, a.cfg_path]:

        if not Path(path).is_file():

            parser.error('Missing file: ' + path)

    if not (Path(a.ckpt_path) / 'config.json').is_file():

        parser.error('ckpt-path must be the conversation_weight directory containing config.json.')

    import torch

    if not torch.cuda.is_available():

        raise RuntimeError('CUDA unavailable; fix PyTorch before loading ST-LLM.')

    if torch.cuda.get_device_capability(0) == (12, 0) and 'sm_120' not in torch.cuda.get_arch_list():

        raise RuntimeError('Install a Blackwell-compatible PyTorch wheel with sm_120 support.')

    x = torch.randn(64, 64, device='cuda')

    _ = x @ x

    torch.cuda.synchronize()

    from stllm.common.config import Config

    from stllm.common.registry import registry

    import stllm.models

    from stllm.conversation.conversation import Chat, CONV_instructblip_Vicuna0

    config_args = argparse.Namespace(cfg_path=a.cfg_path, options=None)

    cfg = Config(config_args)

    mc = cfg.model_cfg

    mc.ckpt = str(Path(a.ckpt_path).resolve())

    mc.llama_model = mc.ckpt

    mc.q_former_model = str(Path(a.qformer_path).resolve())

    mc.device_8bit = 0

    mc.use_grad_checkpoint = False

    # Keep residual modules when loading this residual-trained checkpoint.

    # Switch only its inference aggregation after loading.

    model_cls = registry.get_model_class(mc.arch)

    print('Loading official ST-LLM checkpoint; full FP16, no quantization/offload.', flush=True)

    model = model_cls.from_config(mc).to(device='cuda:0', dtype=torch.float16).eval()

    chat = Chat(model, device='cuda:0')

    chat.model.video_input = a.video_input

    conv = CONV_instructblip_Vicuna0.copy()

    images = []

    torch.cuda.reset_peak_memory_stats()

    start = time.perf_counter()

    with torch.inference_mode():

        chat.upload_video(str(Path(a.video).resolve()), conv, images,

                          num_frame=a.frames, text=a.question)

        chat.ask('###Human: ' + a.question + ' ###Assistant: ', conv)

        probe = conv.copy()

        probe.append_message(probe.roles[1], None)

        embeddings, _ = chat.get_context_emb_sim(probe, images, system=True)

        input_tokens = int(embeddings.shape[1])

        visual_tokens = int(images[0].shape[1])

        del embeddings

        context_limit = int(model.config.max_position_embeddings)

        required = input_tokens + a.max_new_tokens

        if required > context_limit:

            raise RuntimeError(f'Context overflow: {input_tokens} input + {a.max_new_tokens} output > {context_limit}. Reduce --frames; no silent truncation allowed.')

        answer, _ = chat.answer(conv=conv, img_list=images, do_sample=False,

                                num_beams=1, max_new_tokens=a.max_new_tokens,

                                max_length=context_limit)

    torch.cuda.synchronize()

    result = {'arguments': vars(a), 'answer': answer, 'visual_tokens': visual_tokens,

              'input_tokens': input_tokens, 'context_limit': context_limit,

              'encode_and_generate_seconds': time.perf_counter()-start,

              'peak_allocated_gpu_gib': torch.cuda.max_memory_allocated()/2**30,

              'torch_version': torch.__version__, 'gpu': torch.cuda.get_device_name(0),

              'note': 'all retains all per-frame encoder/Q-Former output tokens, not raw pixels or all ViT patches; residual uses checkpoint global-local path.'}

    dest = Path(a.output)

    dest.parent.mkdir(parents=True, exist_ok=True)

    dest.write_text(json.dumps(result, indent=2), encoding='utf-8')

    print(answer)

if __name__ == '__main__':

    main()

Some caveat remain as before:

1. Libraries I used:

Python 3.10.22

torch 2.2.1+cu118

transformers 4.40.2

peft 0.10.0

accelerate 0.26.1

safetensors 0.4.2

mmcv 1.7.2

decord 0.6.0

moviepy 1.0.3

imageio 2.27.0

matplotlib 3.8.3

protobuf 4.23.4

sentencepiece 0.2.0

einops 0.6.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:

stllm-blackwell

llava-v1.6-vicuna-7b-hf

llava-1.5-7b-hf

The checkpoint architecture and the code architecture do not match.

3. Last known state:

Environment built successfully.

Model weights downloaded successfully.

Preprocessing works.

Model loading partially works.

Current blocker = intermittent native crash / ST-LLM compatibility issue.