File size: 3,315 Bytes
7011765 8321447 7011765 8321447 7011765 8321447 7011765 | 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 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | from __future__ import annotations
import sys
from pathlib import Path
SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR / "vllm"))
from vllm_caption_runtime import build_vllm_request, load_llm, load_processor, sampling_params
# =============================================================================
# Paths
# =============================================================================
VIDEO_PATH = "/workspace/test7.mp4"
MODEL_DIR = str(SCRIPT_DIR / "model")
# =============================================================================
# Prompt Settings
# =============================================================================
PROMPT_OVERRIDE = ""
CAPTION_LENGTH = "very large"
INCLUDE_WATERMARK_INFO = False
HAS_THINKING = True
VULGARITY = "low"
UNCERTAINTY = "low"
CHARACTER_NAMES = "none"
FLUFF = "none"
HAS_REPETITION = False
SPECULATION = "low"
TEMPORAL_DETAIL = "medium"
VISUAL_SPECIFICITY = "moderate"
CAMERA_DETAIL = "medium"
CAPTION_STYLE = "plain"
# =============================================================================
# vLLM / Generation Hyperparameters
# =============================================================================
NUM_FRAMES = 12
SAMPLING_RATE = 16_000
MAX_MODEL_LEN = 4096
MAX_NUM_SEQS = 1
GPU_MEMORY_UTILIZATION = 0.88
DTYPE = "bfloat16"
# Online vLLM FP8 for Gemma linear weights. Parakeet is left unchanged.
USE_FP8 = True
FP8_QUANTIZATION = "fp8_per_tensor"
ENFORCE_EAGER = False
ENABLE_PREFIX_CACHING = False
TRUST_REMOTE_CODE = False
MAX_TOKENS = 1200
TEMPERATURE = 0.0
TOP_P = 0.9
REPETITION_PENALTY = 1.1
def prompt_settings() -> dict[str, object]:
return {
"caption_length": CAPTION_LENGTH,
"include_watermark_info": INCLUDE_WATERMARK_INFO,
"has_thinking": HAS_THINKING,
"vulgarity": VULGARITY,
"uncertainty": UNCERTAINTY,
"character_names": CHARACTER_NAMES,
"fluff": FLUFF,
"has_repetition": HAS_REPETITION,
"speculation": SPECULATION,
"temporal_detail": TEMPORAL_DETAIL,
"visual_specificity": VISUAL_SPECIFICITY,
"camera_detail": CAMERA_DETAIL,
"caption_style": CAPTION_STYLE,
}
def main() -> None:
quantization = FP8_QUANTIZATION if USE_FP8 else None
processor = load_processor(MODEL_DIR)
request = build_vllm_request(
processor=processor,
model_dir=MODEL_DIR,
video_path=VIDEO_PATH,
num_frames=NUM_FRAMES,
sampling_rate=SAMPLING_RATE,
prompt_override=PROMPT_OVERRIDE,
prompt_settings=prompt_settings(),
)
llm = load_llm(
model_dir=MODEL_DIR,
max_model_len=MAX_MODEL_LEN,
max_num_seqs=MAX_NUM_SEQS,
gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
dtype=DTYPE,
quantization=quantization,
enforce_eager=ENFORCE_EAGER,
enable_prefix_caching=ENABLE_PREFIX_CACHING,
trust_remote_code=TRUST_REMOTE_CODE,
)
outputs = llm.generate(
[request],
sampling_params=sampling_params(
temperature=TEMPERATURE,
max_tokens=MAX_TOKENS,
top_p=TOP_P,
repetition_penalty=REPETITION_PENALTY,
),
)
print(outputs[0].outputs[0].text)
if __name__ == "__main__":
main()
|