Views
No views yet

1from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
2from qwen_vl_utils import process_vision_info
3
4# Load the model on the available device(s)
5model = Qwen2VLForConditionalGeneration.from_pretrained(
6 "prithivMLmods/JSONify-Flux-Large", torch_dtype="auto", device_map="auto"
7)
8
9# Enable flash_attention_2 for better acceleration and memory efficiency
10# model = Qwen2VLForConditionalGeneration.from_pretrained(
11# "prithivMLmods/JSONify-Flux-Large",
12# torch_dtype=torch.bfloat16,
13# attn_implementation="flash_attention_2",
14# device_map="auto",
15# )
16
17# Default processor
18processor = AutoProcessor.from_pretrained("prithivMLmods/JSONify-Flux-Large")
19
20messages = [
21 {
22 "role": "user",
23 "content": [
24 {
25 "type": "image",
26 "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
27 },
28 {"type": "text", "text": "Describe this image in JSON format."},
29 ],
30 }
31]
32
33# Prepare inputs for inference
34text = processor.apply_chat_template(
35 messages, tokenize=False, add_generation_prompt=True
36)
37image_inputs, video_inputs = process_vision_info(messages)
38inputs = processor(
39 text=[text],
40 images=image_inputs,
41 videos=video_inputs,
42 padding=True,
43 return_tensors="pt",
44)
45inputs = inputs.to("cuda")
46
47# Inference: Generate JSON-formatted output
48generated_ids = model.generate(**inputs, max_new_tokens=128)
49generated_ids_trimmed = [
50 out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
51]
52output_text = processor.batch_decode(
53 generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
54)
55
56print(output_text) # JSON-formatted response1buffer = ""
2for new_text in streamer:
3 buffer += new_text
4 buffer = buffer.replace("<|im_end|>", "")
5 yield buffer