| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.append(str(Path(__file__).resolve().parents[1])) |
| from src.io_utils import apply_overrides, load_config |
| from src.qwen_video_utils import generate_one, load_model_for_training, load_processor |
| from src.train_utils import oom_help, quantization_config_from_config |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", required=True) |
| parser.add_argument("--video", required=True) |
| parser.add_argument("--model_name", default=None) |
| parser.add_argument("--prompt", default=None) |
| parser.add_argument("--max_frames", type=int, default=None) |
| parser.add_argument("--video_fps", type=float, default=None) |
| args = parser.parse_args() |
|
|
| cfg = apply_overrides(load_config(args.config), model_name=args.model_name) |
| processor = load_processor(cfg["model_name"]) |
| model = load_model_for_training( |
| cfg["model_name"], |
| quantization_config=quantization_config_from_config(cfg), |
| device_map="auto", |
| dtype="bfloat16" if cfg.get("bf16", True) else "float16", |
| ) |
| model.eval() |
| try: |
| prediction = generate_one( |
| model, |
| processor, |
| args.video, |
| args.prompt or cfg["prompt"], |
| cfg.get("generation", {}), |
| video_fps=args.video_fps if args.video_fps is not None else cfg.get("video_fps"), |
| max_frames=args.max_frames if args.max_frames is not None else cfg.get("max_frames"), |
| ) |
| except torch.cuda.OutOfMemoryError as exc: |
| raise RuntimeError(oom_help()) from exc |
| print(prediction) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
