Update app.py

app.py

@@ -1,15 +1,12 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
+import requests
+from io import BytesIO
 
 # Load the model and processor
 repo_name = "cyan2k/molmo-7B-O-bnb-4bit"
-arguments = {
-    "device_map": "auto",
-    "torch_dtype": "auto",
-    "trust_remote_code": True,
-    "load_in_8bit": True  # Use 8-bit for reduced memory footprint
-}
+arguments = {"device_map": "auto", "torch_dtype": "auto", "trust_remote_code": True}
 
 # Load the processor and model
 processor = AutoProcessor.from_pretrained(repo_name, **arguments)
@@ -19,12 +16,12 @@ def describe_image(image):
     # Process the uploaded image
     inputs = processor.process(
         images=[image],
-        text="Describe this image in great detail."
+        text="Describe this image."
     )
     
     # Move inputs to model device
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}  # Removed unsqueeze(0) to keep batch size
-
+    inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
+    
     # Generate output
     output = model.generate_from_batch(
         inputs,
@@ -57,4 +54,4 @@ def gradio_app():
     interface.launch()
 
 # Launch the Gradio app
-gradio_app()
+gradio_app()