import time
from transformers import AutoModelForImageTextToText, AutoProcessor, AutoModel, TextIteratorStreamer
from qwen_vl_utils import process_vision_info
from threading import Thread

# default: Load the model on the available device(s)
model = AutoModelForImageTextToText.from_pretrained(
    # "src/models/Qwen2.5-VL-3B-Instruct", torch_dtype="auto", device_map="auto", trust_remote_code=True
    "Qwen/Qwen2.5-VL-3B-Instruct", torch_dtype="auto", device_map="auto", trust_remote_code=True
)

# We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
# model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
#     "Qwen/Qwen2.5-VL-3B-Instruct",
#     torch_dtype=torch.bfloat16,
#     attn_implementation="flash_attention_2",
#     device_map="auto",
# )

# default processer
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", use_fast=False)

# The default range for the number of visual tokens per image in the model is 4-16384.
# You can set min_pixels and max_pixels according to your needs, such as a token range of 256-1280, to balance performance and cost.
# min_pixels = 256*28*28
# max_pixels = 1280*28*28
# processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)

messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
            },
            {"type": "text", "text": "Describe this image."},
        ],
    }
]

# Preparation for inference
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",
)
inputs = inputs.to("cuda")

# Inference: Generation with streaming output
print("-" * 80)
print("Streaming output:")
print("-" * 80, flush=True)

# Create streamer for real-time output
streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True)

# Run generation in a separate thread
generation_kwargs = dict(inputs, max_new_tokens=128, streamer=streamer)
thread = Thread(target=model.generate, kwargs=generation_kwargs)

start_time = time.time()
thread.start()

# Stream the output token by token
generated_text = ""
token_count = 0
for new_text in streamer:
    generated_text += new_text
    token_count += 1
    print(new_text, end="", flush=True)

thread.join()

# Print stats
elapsed = time.time() - start_time
print()
print("-" * 80)
print(f"\033[90m[{token_count} tokens in {elapsed:.2f}s = {token_count/elapsed:.1f} tok/s]\033[0m")
