| """Fast-dVLM command-line chatbot with speculative block-causal decoding.""" |
|
|
| import argparse |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer, AutoProcessor |
| from qwen_vl_utils import process_vision_info |
|
|
|
|
| def build_messages(image, prompt): |
| content = [] |
| if image: |
| content.append({"type": "image", "image": image}) |
| content.append({"type": "text", "text": prompt}) |
| return [{"role": "user", "content": content}] |
|
|
|
|
| def run_inference(model, processor, image, prompt, args): |
| messages = build_messages(image, prompt) |
| text = processor.apply_chat_template( |
| messages, tokenize=False, add_generation_prompt=True, |
| ) |
| image_inputs, video_inputs = process_vision_info(messages) |
| inputs = processor( |
| text=[text], |
| images=image_inputs, |
| videos=video_inputs, |
| padding=True, |
| return_tensors="pt", |
| ).to(model.device) |
|
|
| mask_id = processor.tokenizer.encode(args.mask_token)[0] |
|
|
| gen_kwargs = { |
| "input_ids": inputs.input_ids, |
| "tokenizer": processor.tokenizer, |
| "block_size": args.block_size, |
| "max_tokens": args.max_tokens, |
| "mask_id": mask_id, |
| } |
| if hasattr(inputs, "pixel_values"): |
| gen_kwargs["pixel_values"] = inputs.pixel_values |
| if hasattr(inputs, "image_grid_thw"): |
| gen_kwargs["image_grid_thw"] = inputs.image_grid_thw |
|
|
| generated_ids = model.generate(**gen_kwargs) |
| generated_ids_trimmed = [ |
| out[len(inp):] for inp, out in zip(inputs.input_ids, generated_ids) |
| ] |
| return processor.batch_decode( |
| generated_ids_trimmed, skip_special_tokens=True, |
| )[0] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Fast-dVLM chatbot") |
| parser.add_argument( |
| "--model-name", default="Efficient-Large-Model/Fast_dVLM_3B", |
| help="HuggingFace model id or local path.", |
| ) |
| parser.add_argument( |
| "--image", default=None, |
| help="Image URL or local path. Leave empty for text-only.", |
| ) |
| parser.add_argument( |
| "--prompt", default=None, |
| help="Text prompt. If omitted, enters interactive mode.", |
| ) |
| parser.add_argument("--block-size", type=int, default=32) |
| parser.add_argument("--max-tokens", type=int, default=512) |
| parser.add_argument("--mask-token", default="|<MASK>|") |
| parser.add_argument("--device", default="cuda") |
| args = parser.parse_args() |
|
|
| print("Loading model...") |
| model = AutoModelForCausalLM.from_pretrained( |
| args.model_name, |
| torch_dtype=torch.bfloat16, |
| device_map="auto", |
| trust_remote_code=True, |
| ) |
| tokenizer = AutoTokenizer.from_pretrained(args.model_name) |
| processor = AutoProcessor.from_pretrained(args.model_name, use_fast=False) |
| processor.tokenizer = tokenizer |
|
|
| if args.prompt: |
| image = args.image or "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" |
| output = run_inference(model, processor, image, args.prompt, args) |
| print(f"\n{output}") |
| else: |
| print("Interactive mode. Type 'exit' to quit, 'clear' to reset.") |
| while True: |
| prompt = input("\nYou: ").strip() |
| if prompt.lower() == "exit": |
| break |
| if prompt.lower() == "clear": |
| print("History cleared.") |
| continue |
| if not prompt: |
| continue |
| image = args.image |
| output = run_inference(model, processor, image, prompt, args) |
| print(f"\nAssistant: {output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
