Update app.py

app.py

@@ -2,17 +2,15 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
 import torch
-import spaces
 
 # 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}
 
-# Load the processor and quantized model, ensure they are on the GPU
+# Load the processor and quantized model
 processor = AutoProcessor.from_pretrained(repo_name, **arguments)
 model = AutoModelForCausalLM.from_pretrained(repo_name, **arguments).to("cuda")
 
-@spaces.GPU(duration=120)
 def process_image_and_text(image, text):
     # Process the image and text
     inputs = processor.process(
@@ -24,16 +22,13 @@ def process_image_and_text(image, text):
     inputs = {k: v.to("cuda").unsqueeze(0) for k, v in inputs.items()}
 
     # Generate output
-    output = model.generate_from_batch(
-        inputs,
-        GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
-        tokenizer=processor.tokenizer
+    output = model.generate(
+        **inputs,
+        max_new_tokens=200
     )
 
     # Only get generated tokens; decode them to text
-    generated_tokens = output[0, inputs['input_ids'].size(1):]
-    generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-
+    generated_text = processor.decode(output, skip_special_tokens=True)
     return generated_text
 
 def chatbot(image, text, history):
@@ -46,7 +41,7 @@ def chatbot(image, text, history):
 
 # Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Image Chatbot with Molmo-7B-D-0924")
+    gr.Markdown("# Image Chatbot with Molmo-7B-4 Bit Quantized")
     
     with gr.Row():
         image_input = gr.Image(type="numpy")