Update app.py

app.py

@@ -1,72 +1,80 @@
 import gradio as gr
-from transformers import (
-    AutoModelForCausalLM,
-    AutoProcessor,
-    GenerationConfig,
-    BitsAndBytesConfig,
-)
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
 import torch
+import spaces
 
-# Configuration for 4-bit quantization and GPU offloading
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
+# Load the processor and model
+processor = AutoProcessor.from_pretrained(
+    'allenai/Molmo-7B-D-0924',
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto'
 )
 
-# Model repository
-repo_name = "cyan2k/molmo-7B-O-bnb-4bit"
-
-# Load the processor and model
-processor = AutoProcessor.from_pretrained(repo_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    repo_name,
-    torch_dtype=torch.float16,
-    device_map="auto",
+    'allenai/Molmo-7B-D-0924',
     trust_remote_code=True,
-    quantization_config=bnb_config,
+    torch_dtype='auto',
+    device_map='auto'
 )
 
-# Ensure model is on GPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
 
-def describe_images(images):
-    descriptions = []
-    for image in images:
-        if isinstance(image, str):
-            image = Image.open(image)
-        # Process the image
-        inputs = processor.process(
-            images=[image],
-            text="Describe this image in great detail.",
-        )
-        # Move inputs to the same device as the model
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-        # Generate output
-        with torch.no_grad():
-            output = model.generate_from_batch(
-                inputs,
-                GenerationConfig(max_new_tokens=200, stop_strings=["<|endoftext|>"]),
-                tokenizer=processor.tokenizer,
-            )
-        # Decode generated tokens to text
-        generated_tokens = output[0, inputs["input_ids"].size(1):]
-        generated_text = processor.tokenizer.decode(
-            generated_tokens, skip_special_tokens=True
-        )
-        descriptions.append(generated_text.strip())
-    return "\n\n".join(descriptions)
+@spaces.GPU(duration=120)
+def process_image_and_text(image, text):
+    # Process the image and text
+    inputs = processor.process(
+        images=[Image.fromarray(image)],
+        text=text
+    )
+
+    # Move inputs to the correct device and make a batch of size 1
+    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|>"),
+        tokenizer=processor.tokenizer
+    )
 
-# Gradio interface
+    # 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)
+
+    return generated_text
+
+def chatbot(image, text, history):
+    if image is None:
+        return history + [("Please upload an image first.", None)]
+
+    response = process_image_and_text(image, text)
+    history.append((text, response))
+    return history
+
+# Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("<h3><center>Image Description Generator</center></h3>")
+    gr.Markdown("# Image Chatbot with Molmo-7B-D-0924")
+    
     with gr.Row():
-        image_input = gr.File(
-            file_types=["image"], label="Upload Image(s)", multiple=True
-        )
-    generate_button = gr.Button("Generate Descriptions")
-    output_text = gr.Textbox(label="Descriptions", lines=15)
+        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")
+
+    state = gr.State([])
+
+    submit_button.click(
+        chatbot,
+        inputs=[image_input, text_input, state],
+        outputs=[chatbot_output]
+    )
 
-    generate_button.click(describe_images, inputs=image_input, outputs=output_text)
+    text_input.submit(
+        chatbot,
+        inputs=[image_input, text_input, state],
+        outputs=[chatbot_output]
+    )
 
-demo.launch()
+demo.launch()