fixing image loading

main.py

@@ -1,17 +1,19 @@
+import base64
+from io import BytesIO
+
+import torch
 from fastapi import FastAPI, Query
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from PIL import Image
 from qwen_vl_utils import process_vision_info
-import torch
+from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
 
 app = FastAPI()
 
 checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
-min_pixels = 256*28*28
-max_pixels = 1280*28*28
+min_pixels = 256 * 28 * 28
+max_pixels = 1280 * 28 * 28
 processor = AutoProcessor.from_pretrained(
-    checkpoint,
-    min_pixels=min_pixels,
-    max_pixels=max_pixels
+    checkpoint, min_pixels=min_pixels, max_pixels=max_pixels
 )
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     checkpoint,
@@ -20,17 +22,64 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
     # attn_implementation="flash_attention_2",
 )
 
+
 @app.get("/")
 def read_root():
     return {"message": "API is live. Use the /predict endpoint."}
 
+
+def encode_image(image_path, max_size=(800, 800), quality=85):
+    """
+    Converts an image from a local file path to a Base64-encoded string with optimized size.
+
+    Args:
+        image_path (str): The path to the image file.
+        max_size (tuple): The maximum width and height of the resized image.
+        quality (int): The compression quality (1-100, higher means better quality but bigger size).
+
+    Returns:
+        str: Base64-encoded representation of the optimized image.
+    """
+    try:
+        with Image.open(image_path) as img:
+            # Convert to RGB (avoid issues with PNG transparency)
+            img = img.convert("RGB")
+
+            # Resize while maintaining aspect ratio
+            img.thumbnail(max_size, Image.LANCZOS)
+
+            # Save to buffer with compression
+            buffer = BytesIO()
+            img.save(
+                buffer, format="JPEG", quality=quality
+            )  # Save as JPEG to reduce size
+            return base64.b64encode(buffer.getvalue()).decode("utf-8")
+    except Exception as e:
+        print(f"❌ Error encoding image {image_path}: {e}")
+        return None
+
+
 @app.get("/predict")
 def predict(image_url: str = Query(...), prompt: str = Query(...)):
+
+    image = encode_image(image_url)
+
     messages = [
-        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
-        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
+        {
+            "role": "system",
+            "content": "You are a helpful assistant with vision abilities.",
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": f"data:image;base64,{image}"},
+                {"type": "text", "text": prompt},
+            ],
+        },
     ]
-    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor(
         text=[text],
@@ -41,8 +90,13 @@ def predict(image_url: str = Query(...), prompt: str = Query(...)):
     ).to(model.device)
     with torch.no_grad():
         generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :]
+        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
     output_texts = processor.batch_decode(
-        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        generated_ids_trimmed,
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=False,
     )
     return {"response": output_texts[0]}

model.py

@@ -1,10 +1,15 @@
 import requests
 
-url = "https://<uname>-<spacename>.hf.space/predict"
+# curl -G "https://<uname>-<spacename>.hf.space/predict" \
+#      --data-urlencode "image_url=https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" \
+#      --data-urlencode "prompt=Describe this image."
+
+
+url = "https://danilohssantana-qwen2-5-vl-api.hf.space/predict"
 
 # Define the parameters
 params = {
-    "image_url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
+    "image_url": "https://cdn.britannica.com/35/238335-050-2CB2EB8A/Lionel-Messi-Argentina-Netherlands-World-Cup-Qatar-2022.jpg",
     "prompt": "describe",
 }
 

test.py

@@ -0,0 +1,13 @@
+from io import BytesIO
+
+import requests
+from PIL import Image
+
+image_url = "https://cdn.britannica.com/35/238335-050-2CB2EB8A/Lionel-Messi-Argentina-Netherlands-World-Cup-Qatar-2022.jpg"
+response = requests.get(image_url, stream=True)
+
+if response.status_code == 200:
+    image = Image.open(BytesIO(response.content))
+    image.show()
+else:
+    print(f"Failed to download image. Status code: {response.status_code}")