Update app.py

app.py

@@ -3,11 +3,14 @@ from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
 import requests
 from io import BytesIO
-import subprocess
-import sys
+
 # Load the model and processor
 repo_name = "cyan2k/molmo-7B-O-bnb-4bit"
-arguments = {"device_map": "auto", "torch_dtype": "auto", "trust_remote_code": True}
+arguments = {
+    "device_map": "auto",   # Force CPU inference
+    "torch_dtype": "auto",  # Set model to use float32 precision
+    "trust_remote_code": True  # Allow the loading of remote code
+}
 
 # Load the processor and model
 processor = AutoProcessor.from_pretrained(repo_name, **arguments)
@@ -17,23 +20,23 @@ def describe_image(image):
     # Process the uploaded image
     inputs = processor.process(
         images=[image],
-        text="Describe this image."
+        text="Describe this image in great detail without missing any piece of information"
     )
-    
+
     # Move inputs to model device
     inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
-    
+
     # Generate output
     output = model.generate_from_batch(
         inputs,
-        GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
+        GenerationConfig(max_new_tokens=1024, stop_strings="<|endoftext|>"),
         tokenizer=processor.tokenizer,
     )
-    
+
     # Decode the generated tokens
     generated_tokens = output[0, inputs["input_ids"].size(1):]
     generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-    
+
     return generated_text
 
 
@@ -50,7 +53,7 @@ def gradio_app():
         title="Image Description App",
         description="Upload an image and get a detailed description using the Molmo 7B model"
     )
-    
+
     # Launch the interface
     interface.launch()