infer.py

#!/usr/bin/env python3
"""Rewrite one MiniMax-H3 T2VA/I2VA/L2VA/FL2VA prompt with the 8B LoRA."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from PIL import Image, ImageOps
import torch
import transformers
from peft import PeftModel
from transformers import AutoProcessor

from prompt_template import build_messages, expected_image_count, normalize_task


DEFAULT_BASE_MODEL = "Qwen/Qwen3-VL-8B-Instruct"
DEFAULT_ADAPTER_REPO = "lightx2v/MiniMax-H3-Prompt-Rewriter-LoRA-8B"
REPO_ROOT = Path(__file__).resolve().parent


def get_model_class():
    """Prefer Qwen3-VL's concrete class, with portable AutoModel fallbacks."""

    candidates = (
        "Qwen3VLForConditionalGeneration",
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModelForMultimodalLM",
    )
    for name in candidates:
        model_class = getattr(transformers, name, None)
        if model_class is not None:
            return model_class
    raise RuntimeError(
        "This Transformers installation does not expose a Qwen3-VL-compatible "
        "conditional-generation class. Upgrade Transformers and retry."
    )


def default_adapter_path() -> str:
    """Use weights in this checkout when present, otherwise use the Hub repo."""

    if (REPO_ROOT / "adapter_config.json").is_file():
        return str(REPO_ROOT)
    return DEFAULT_ADAPTER_REPO


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-model", default=DEFAULT_BASE_MODEL)
    parser.add_argument(
        "--adapter-path",
        default=None,
        help=(
            "Local LoRA directory or Hugging Face repo ID. By default, use this "
            "checkout when adapter_config.json exists, otherwise use the 8B Hub repo."
        ),
    )
    parser.add_argument("--base-only", action="store_true", help="Run the base model without loading the LoRA")
    parser.add_argument("--task", default="t2va", help="T2VA, I2VA, L2VA, or FL2VA (common *2V/*2AV aliases work)")

    prompt_source = parser.add_mutually_exclusive_group(required=True)
    prompt_source.add_argument("--prompt", help="Original user prompt")
    prompt_source.add_argument("--prompt-file", type=Path, help="UTF-8 file containing the original prompt")

    parser.add_argument("--duration", type=int, choices=range(4, 16), default=10)
    parser.add_argument(
        "--resolution",
        choices=("adaptive", "21:9", "16:9", "4:3", "1:1", "3:4", "9:16"),
        help="Defaults to 16:9 for T2VA and adaptive for image-conditioned tasks",
    )
    parser.add_argument("--first-frame", type=Path, help="First-frame image for I2VA or FL2VA")
    parser.add_argument("--last-frame", type=Path, help="Last-frame image for L2VA or FL2VA")
    parser.add_argument("--output", type=Path, help="Optional .txt or .json output path")

    parser.add_argument("--max-new-tokens", type=int, default=4096)
    parser.add_argument("--min-pixels", type=int, default=256 * 256)
    parser.add_argument("--max-pixels", type=int, default=1024 * 1024)
    parser.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
    parser.add_argument("--attn-implementation", default="sdpa")
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--greedy", action="store_true", help="Disable sampling")
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top-p", type=float, default=0.8)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    try:
        args.task = normalize_task(args.task)
    except ValueError as exc:
        parser.error(str(exc))

    if args.prompt_file is not None:
        if not args.prompt_file.is_file():
            parser.error(f"prompt file not found: {args.prompt_file}")
        args.prompt = args.prompt_file.read_text(encoding="utf-8").strip()
    else:
        args.prompt = str(args.prompt or "").strip()
    if not args.prompt:
        parser.error("prompt must not be empty")

    if args.max_new_tokens <= 0:
        parser.error("--max-new-tokens must be positive")
    if args.min_pixels <= 0 or args.max_pixels < args.min_pixels:
        parser.error("pixel limits must satisfy 0 < --min-pixels <= --max-pixels")
    if not args.greedy and (args.temperature <= 0 or not 0 < args.top_p <= 1):
        parser.error("sampling requires --temperature > 0 and 0 < --top-p <= 1")

    required = expected_image_count(args.task)
    supplied = int(args.first_frame is not None) + int(args.last_frame is not None)
    if supplied != required:
        parser.error(f"{args.task} requires {required} reference image(s), but {supplied} were supplied")
    if args.task == "i2av" and args.last_frame is not None:
        parser.error("I2VA accepts --first-frame only")
    if args.task == "l2av" and args.first_frame is not None:
        parser.error("L2VA accepts --last-frame only")
    for path in (args.first_frame, args.last_frame):
        if path is not None and not path.is_file():
            parser.error(f"image not found: {path}")

    args.resolution = args.resolution or ("16:9" if args.task == "t2av" else "adaptive")
    args.adapter_path = args.adapter_path or default_adapter_path()
    return args


def load_images(args: argparse.Namespace) -> list[Image.Image]:
    """Load images in the exact placeholder order used by prompt_template.py."""

    paths: list[Path] = []
    if args.task in {"i2av", "fl2av"}:
        paths.append(args.first_frame)
    if args.task in {"l2av", "fl2av"}:
        paths.append(args.last_frame)

    images: list[Image.Image] = []
    for path in paths:
        with Image.open(path) as image:
            images.append(ImageOps.exif_transpose(image).convert("RGB").copy())
    return images


def load_processor(args: argparse.Namespace):
    processor_source = args.base_model if args.base_only else args.adapter_path
    processor_kwargs = {
        "trust_remote_code": True,
        "min_pixels": args.min_pixels,
        "max_pixels": args.max_pixels,
    }
    try:
        return AutoProcessor.from_pretrained(processor_source, **processor_kwargs)
    except (OSError, ValueError, KeyError):
        if processor_source == args.base_model:
            raise
        return AutoProcessor.from_pretrained(args.base_model, **processor_kwargs)


def model_input_device(model: torch.nn.Module) -> torch.device:
    embedding_device = model.get_input_embeddings().weight.device
    if embedding_device.type != "meta":
        return embedding_device
    for parameter in model.parameters():
        if parameter.device.type != "meta":
            return parameter.device
    raise RuntimeError("Could not determine a real input device for the loaded model")


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    processor = load_processor(args)
    load_kwargs = {
        "torch_dtype": getattr(torch, args.dtype),
        "low_cpu_mem_usage": True,
        "trust_remote_code": True,
        "device_map": args.device_map,
    }
    if args.attn_implementation:
        load_kwargs["attn_implementation"] = args.attn_implementation

    model = get_model_class().from_pretrained(args.base_model, **load_kwargs)
    if not args.base_only:
        model = PeftModel.from_pretrained(model, args.adapter_path)
    model.eval()

    messages = build_messages(
        args.prompt,
        task=args.task,
        resolution=args.resolution,
        duration=args.duration,
    )
    rendered = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    processor_kwargs = {
        "text": [rendered],
        "return_tensors": "pt",
        "padding": False,
        "return_mm_token_type_ids": True,
    }
    images = load_images(args)
    if images:
        processor_kwargs["images"] = images
    inputs = processor(**processor_kwargs)
    input_device = model_input_device(model)
    inputs = {
        key: value.to(input_device) if isinstance(value, torch.Tensor) else value
        for key, value in inputs.items()
    }

    generation_kwargs = {"max_new_tokens": args.max_new_tokens}
    if args.greedy:
        generation_kwargs["do_sample"] = False
    else:
        generation_kwargs.update(
            do_sample=True,
            temperature=args.temperature,
            top_p=args.top_p,
        )

    with torch.inference_mode():
        output_ids = model.generate(**inputs, **generation_kwargs)
    generated_ids = output_ids[0, inputs["input_ids"].shape[1] :]
    rewritten_prompt = processor.decode(generated_ids, skip_special_tokens=True).strip()
    print(rewritten_prompt)

    if args.output is not None:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        if args.output.suffix.lower() == ".json":
            payload = {
                "task": args.task,
                "resolution": args.resolution,
                "duration": args.duration,
                "prompt": args.prompt,
                "first_frame": str(args.first_frame) if args.first_frame else None,
                "last_frame": str(args.last_frame) if args.last_frame else None,
                "base_model": args.base_model,
                "adapter_path": None if args.base_only else str(args.adapter_path),
                "enhanced_prompt": rewritten_prompt,
            }
            args.output.write_text(
                json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )
        else:
            args.output.write_text(rewritten_prompt + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()