Update app.py

app.py

@@ -12,9 +12,9 @@ florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base
 florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
 
 def generate_caption(image):
-    # Ensure that the image is a PIL image
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)  # Convert numpy array to PIL.Image if necessary
+    """Generate a caption for the given image using Florence 2"""
+    # Convert image to RGB format to avoid channel errors
+    image = image.convert("RGB")  
 
     # Prepare the input for the Florence model
     inputs = florence_processor(text="<MORE_DETAILED_CAPTION>", images=image, return_tensors="pt").to(device)
@@ -30,16 +30,9 @@ def generate_caption(image):
     )
     
     # Decode the generated text
-    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-    
-    # Post-process the generated text
-    parsed_answer = florence_processor.post_process_generation(
-        generated_text,
-        task="<MORE_DETAILED_CAPTION>",
-        image_size=(image.width, image.height)
-    )
+    generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    return parsed_answer["<MORE_DETAILED_CAPTION>"]
+    return generated_text
 
 # Streamlit UI
 st.title("Florence 2 Caption Generator")
@@ -60,20 +53,21 @@ if uploaded_image is not None:
         st.write(caption)
 
 # ✅ API Mode: Handle API Requests
-st.experimental_set_query_params()  # Ensure Streamlit can handle query params
-
 def handle_api_request():
     """Handle API request by checking URL query parameters."""
-    query_params = st.experimental_get_query_params()
-    
-    if "image" in query_params:
-        image_base64 = query_params["image"][0]  # Get Base64-encoded image
-        image_bytes = BytesIO(base64.b64decode(image_base64))
-        image = Image.open(image_bytes)
+    query_params = st.query_params
 
-        caption = generate_caption(image)
-        st.json({"caption": caption})  # Return JSON response
+    if "image" in query_params:
+        try:
+            image_base64 = query_params["image"]
+            image_bytes = BytesIO(base64.b64decode(image_base64))
+            image = Image.open(image_bytes).convert("RGB")  # Ensure it's RGB
+            
+            caption = generate_caption(image)
+            st.json({"caption": caption})  # Return JSON response
+        except Exception as e:
+            st.json({"error": str(e)})
 
 # Check if API mode is enabled
-if "image" in st.experimental_get_query_params():
+if "image" in st.query_params:
     handle_api_request()