Update app.py

app.py

@@ -1,71 +1,37 @@
 import torch
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 from PIL import Image
 import requests
-from io import BytesIO  # Importing BytesIO from the io module
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
-from qwen_vl_utils import process_vision_info
-
-# Check if CUDA is available and set the device accordingly
-device = "cuda" if torch.cuda.is_available() else "cpu"
+from io import BytesIO
 
-# Load the model on the available device
-model = Qwen2VLForConditionalGeneration.from_pretrained(
-    "Qwen/Qwen2-VL-2B-Instruct", torch_dtype="auto", device_map="auto"
-)
-model = model.to(device)
+# Initialize the model and processor
+model_name = "Qwen/Qwen2-VL-2B-Instruct"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Default processor
-processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
+model = Qwen2VLForConditionalGeneration.from_pretrained(model_name).to(device)
+processor = AutoProcessor.from_pretrained(model_name)
 
-# Resize the image to a smaller resolution (e.g., 512x512)
+# Load the image from URL
 image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
 response = requests.get(image_url)
-img = Image.open(BytesIO(response.content))  # Using BytesIO to handle image from the byte stream
+img = Image.open(BytesIO(response.content))
 
-# Resize the image
-img_resized = img.resize((512, 512))  # Resize the image to 512x512
-image_inputs = processor(images=img_resized, return_tensors="pt").to(device)
+# Ensure the image is resized and processed correctly
+img_resized = img.resize((224, 224))  # Resize as needed (adjust based on model requirements)
 
-# Prepare the text input
-messages = [
-    {
-        "role": "user",
-        "content": [
-            {
-                "type": "image",
-                "image": img_resized,
-            },
-            {"type": "text", "text": "Describe this image."},
-        ],
-    }
-]
+# Create a prompt or text input
+text_input = "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)
+# Process the image and the text input
 inputs = processor(
-    text=[text],
-    images=image_inputs,
-    videos=video_inputs,
-    padding=True,
+    images=img_resized,
+    text=text_input,
     return_tensors="pt",
-)
-inputs = inputs.to(device)  # Move inputs to the same device as the model
+).to(device)
 
-# Inference: Generation of the output
+# Inference
 generated_ids = model.generate(**inputs, max_new_tokens=128)
 
-# Trim the output tokens
-generated_ids_trimmed = [
-    out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
-]
-
-# Decode the generated text
-output_text = processor.batch_decode(
-    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
-)
-
-# Print the output
+# Decode the output
+output_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
 print(output_text)