Update app.py

app.py

@@ -3,15 +3,6 @@ from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
 import torch
 
-import os
-
-# Example command to install a package (e.g., numpy)
-os.system("pip install bitsandbytes-cuda111")
-os.system("pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118")
-
-
-
-
 # Define the repository for the quantized model
 repo_name = "cyan2k/molmo-7B-D-bnb-4bit"
 arguments = {"device_map": "auto", "torch_dtype": "auto", "trust_remote_code": True}
@@ -28,7 +19,7 @@ def process_image_and_text(image, text):
     )
 
     # Move inputs to the same device as the model (GPU) and make a batch of size 1
-    inputs = {k: v.unsqueeze(0) for k, v in inputs.items()}
+    inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
 
     # Generate output
     output = model.generate(
@@ -51,11 +42,11 @@ def chatbot(image, text, history):
 # Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Image Chatbot with Molmo-7B-4 Bit Quantized")
-    
+
     with gr.Row():
         image_input = gr.Image(type="numpy")
         chatbot_output = gr.Chatbot()
-    
+
     text_input = gr.Textbox(placeholder="Ask a question about the image...")
     submit_button = gr.Button("Submit")
 
@@ -73,5 +64,4 @@ with gr.Blocks() as demo:
         outputs=[chatbot_output]
     )
 
-demo.launch()
- 
+demo.launch()