Update app.py

app.py

@@ -6,21 +6,23 @@ from transformers import AutoTokenizer, AutoModel
 from PIL import Image
 from typing import Optional
 import io
-import torch 
 
 app = FastAPI()
 
-device = "cpu"
+# "device" variable is no longer needed for the model, 
+# accelerate handles it via device_map="auto".
 model_id = "OpenGVLab/InternVL2_5-2B"
 
-# 2. Update the loading line to force bfloat16 and use low_cpu_mem_usage
+# 1. Load the Model with smart memory management
 model = AutoModel.from_pretrained(
     model_id, 
     trust_remote_code=True, 
-    torch_dtype=torch.bfloat16,  
+    torch_dtype=torch.bfloat16, 
     low_cpu_mem_usage=True,
-    # revision="65b9340"
-).to(device).eval()
+    device_map="auto",  # <--- THIS IS THE FIX: Prevents OOM by handling memory placement automatically
+    offload_folder="offload" # <--- Optional: Explicit folder for offloading if RAM is full
+).eval() # Removed .to(device) as device_map handles this
+
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 
 class GenerateRequest(BaseModel):
@@ -30,7 +32,6 @@ class GenerateRequest(BaseModel):
 
 @app.post("/generate")
 async def generate(image: UploadFile = File(...), request: GenerateRequest = Depends()):
-    # Read image
     image_bytes = await image.read()
     pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
@@ -39,7 +40,9 @@ async def generate(image: UploadFile = File(...), request: GenerateRequest = Dep
     else:
         prompt = f"<s><image>\nDescribe the image.</s>"
 
-    inputs = tokenizer(prompt, pil_image, return_tensors="pt").to(device)
+    # 2. Update inputs to use the model's device dynamically
+    # model.device will tell us where the model (or its first layer) lives
+    inputs = tokenizer(prompt, pil_image, return_tensors="pt").to(model.device)
 
     generation_args = {
         "max_new_tokens": request.max_new_tokens,
@@ -54,4 +57,4 @@ async def generate(image: UploadFile = File(...), request: GenerateRequest = Dep
 
 @app.get("/")
 async def read_root():
-    return {"message": "InternVL2_5-2B API. Go to /docs for API documentation."}
+    return {"message": "InternVL2_5-2B API. Go to /docs for API documentation."}