Update app.py

app.py

@@ -48,19 +48,19 @@ async def predict(files: List[UploadFile] = File(...)):
             # Generate description
             prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, people, and environmental context."
             inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
-            outputs = model.generate(**inputs, max_length=512)
+            outputs = model.generate(**inputs, max_new_tokens=256)
             description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
             
             # Extract signs/number plates (OCR)
             prompt_ocr = "<image> Extract all visible text in the image, such as road signs, license plates, or billboards, with exact wording."
             inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
-            ocr_outputs = model.generate(**inputs_ocr, max_length=512)
+            ocr_outputs = model.generate(**inputs_ocr, max_new_tokens=256)
             signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
             
             # Detect harmful objects/blood with confidence simulation
             prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly and estimate confidence (0-100%) for each detection based on clarity."
             inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
-            detect_outputs = model.generate(**inputs_detect, max_length=512)
+            detect_outputs = model.generate(**inputs_detect, max_new_tokens=256)
             detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
             harmful_detected = []
             for obj in harmful_objects:
@@ -123,19 +123,19 @@ def gradio_predict(*images):
             # Generate description
             prompt_desc = "<image> Provide a detailed description of the image, including objects, colors, people, and environmental context."
             inputs = processor(text=[prompt_desc], images=[image], return_tensors="pt", padding=True)
-            outputs = model.generate(**inputs, max_length=512)
+            outputs = model.generate(**inputs, max_new_tokens=256)
             description = processor.decode(outputs[0], skip_special_tokens=True).replace(prompt_desc, "").strip()
             
             # Extract signs/number plates (OCR)
             prompt_ocr = "<image> Extract all visible text in the image, such as road signs, license plates, or billboards, with exact wording."
             inputs_ocr = processor(text=[prompt_ocr], images=[image], return_tensors="pt", padding=True)
-            ocr_outputs = model.generate(**inputs_ocr, max_length=512)
+            ocr_outputs = model.generate(**inputs_ocr, max_new_tokens=256)
             signs_text = processor.decode(ocr_outputs[0], skip_special_tokens=True).replace(prompt_ocr, "").strip()
             
             # Detect harmful objects/blood with confidence simulation
             prompt_detect = "<image> Identify any harmful objects (e.g., knife, gun, blood, syringe, bomb, blade) in the image. List them explicitly and estimate confidence (0-100%) for each detection based on clarity."
             inputs_detect = processor(text=[prompt_detect], images=[image], return_tensors="pt", padding=True)
-            detect_outputs = model.generate(**inputs_detect, max_length=512)
+            detect_outputs = model.generate(**inputs_detect, max_new_tokens=256)
             detected_objects = processor.decode(detect_outputs[0], skip_special_tokens=True).lower()
             harmful_detected = []
             for obj in harmful_objects: