#!/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()