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()