Update app.py

app.py

@@ -7,7 +7,7 @@ import torch
 @st.cache_resource
 def load_model():
     model_id = 'microsoft/Florence-2-large'
-    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype='auto').eval()
+    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).eval()
     processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
     return model, processor
 
@@ -19,7 +19,16 @@ def run_example(task_prompt, image, text_input=None):
         prompt = task_prompt
     else:
         prompt = task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to(torch.float32)  # Ensure CPU compatibility
+
+    # Prepare inputs
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+    inputs["input_ids"] = inputs["input_ids"].to(torch.float32)
+    inputs["pixel_values"] = inputs["pixel_values"].to(torch.float32)
+
+    # Ensure the model is in float32 mode
+    model = model.to(torch.float32)
+
+    # Generate predictions
     generated_ids = model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],
@@ -51,19 +60,22 @@ if uploaded_file is not None:
     st.subheader("Generated Captions")
     
     with st.spinner("Generating caption..."):
-        caption = run_example('<CAPTION>', image)
-        detailed_caption = run_example('<DETAILED_CAPTION>', image)
-        more_detailed_caption = run_example('<MORE_DETAILED_CAPTION>', image)
-    
-    st.write("**Caption:**", caption)
-    st.write("**Detailed Caption:**", detailed_caption)
-    st.write("**More Detailed Caption:**", more_detailed_caption)
+        try:
+            caption = run_example('<CAPTION>', image)
+            detailed_caption = run_example('<DETAILED_CAPTION>', image)
+            more_detailed_caption = run_example('<MORE_DETAILED_CAPTION>', image)
+
+            st.write("**Caption:**", caption)
+            st.write("**Detailed Caption:**", detailed_caption)
+            st.write("**More Detailed Caption:**", more_detailed_caption)
 
-    # Option to save the output
-    if st.button("Save Captions"):
-        output_path = "captions.txt"
-        with open(output_path, "w") as file:
-            file.write(f"Caption: {caption}\n")
-            file.write(f"Detailed Caption: {detailed_caption}\n")
-            file.write(f"More Detailed Caption: {more_detailed_caption}\n")
-        st.success(f"Captions saved to {output_path}!")
+            # Option to save the output
+            if st.button("Save Captions"):
+                output_path = "captions.txt"
+                with open(output_path, "w") as file:
+                    file.write(f"Caption: {caption}\n")
+                    file.write(f"Detailed Caption: {detailed_caption}\n")
+                    file.write(f"More Detailed Caption: {more_detailed_caption}\n")
+                st.success(f"Captions saved to {output_path}!")
+        except Exception as e:
+            st.error(f"Error: {e}")