changing model

main.py

@@ -2,22 +2,25 @@ import base64
 from io import BytesIO
 
 import torch
-from fastapi import FastAPI, Query
+from fastapi import Body, FastAPI, File, HTTPException, Query, UploadFile
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile, HTTPException
-from qwen_vl_utils import process_vision_info
-from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration, Qwen2VLForConditionalGeneration
-
-from fastapi import FastAPI, Body
 from pydantic import BaseModel
+from qwen_vl_utils import process_vision_info
+from transformers import (
+    AutoProcessor,
+    Qwen2_5_VLForConditionalGeneration,
+    Qwen2VLForConditionalGeneration,
+)
 
 app = FastAPI()
 
+
 # Define request model
 class PredictRequest(BaseModel):
     image_base64: str
     prompt: str
 
+
 # checkpoint = "Qwen/Qwen2-VL-2B-Instruct"
 # min_pixels = 256 * 28 * 28
 # max_pixels = 1280 * 28 * 28
@@ -31,13 +34,11 @@ class PredictRequest(BaseModel):
 #     # attn_implementation="flash_attention_2",
 # )
 
-checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
-min_pixels = 256*28*28
-max_pixels = 1280*28*28
+checkpoint = "Qwen/Qwen2.5-VL-7B-Instruct"
+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,
@@ -82,6 +83,7 @@ def read_root():
 #         print(f"❌ Error encoding image {image_path}: {e}")
 #         return None
 
+
 def encode_image(image_data: BytesIO, max_size=(800, 800), quality=85):
     """
     Converts an image from file data to a Base64-encoded string with optimized size.
@@ -96,6 +98,7 @@ def encode_image(image_data: BytesIO, max_size=(800, 800), quality=85):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error encoding image: {e}")
 
+
 @app.post("/encode-image/")
 async def upload_and_encode_image(file: UploadFile = File(...)):
     """
@@ -108,6 +111,7 @@ async def upload_and_encode_image(file: UploadFile = File(...)):
     except Exception as e:
         raise HTTPException(status_code=400, detail=f"Invalid file: {e}")
 
+
 @app.post("/predict")
 def predict(data: PredictRequest):
     """
@@ -121,7 +125,6 @@ def predict(data: PredictRequest):
         str: The generated description of the image.
     """
 
-
     # Create the input message structure
     messages = [
         {
@@ -160,6 +163,7 @@ def predict(data: PredictRequest):
 
     return {"response": output_text[0] if output_text else "No description generated."}
 
+
 # @app.get("/predict")
 # def predict(image_url: str = Query(...), prompt: str = Query(...)):