Update app.py

app.py

@@ -4,14 +4,27 @@ from PIL import Image
 import torch
 import spaces
 
-# Model name and arguments
-repo_name = "cyan2k/molmo-7B-D-bnb-4bit"
-arguments = {"device_map": "auto", "torch_dtype": "auto", "trust_remote_code": True}
+# Flag to use GPU (set to False by default)
+USE_GPU = False
 
 # Load the processor and model
-processor = AutoProcessor.from_pretrained(repo_name, **arguments)
+device = torch.device("cuda" if USE_GPU and torch.cuda.is_available() else "cpu")
 
-model = AutoModelForCausalLM.from_pretrained(repo_name, **arguments)
+processor = AutoProcessor.from_pretrained(
+    'allenai/MolmoE-1B-0924',
+    trust_remote_code=True,
+    torch_dtype='auto',
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    'allenai/MolmoE-1B-0924',
+    trust_remote_code=True,
+    torch_dtype='auto',
+    device_map='auto' if USE_GPU else None
+)
+
+if not USE_GPU:
+    model.to(device)
 
 # Predefined prompts
 prompts = [
@@ -49,7 +62,6 @@ def process_image_and_text(image, text, max_new_tokens, temperature, top_p):
     generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
 
     return generated_text
- 
 
 def chatbot(image, text, history, max_new_tokens, temperature, top_p):
     if image is None:
@@ -64,7 +76,7 @@ def update_textbox(prompt):
 
 # Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Image Chatbot with Molmo-7B-D-0924")
+    gr.Markdown("# Image Chatbot with MolmoE-1B-0924")
     
     with gr.Row():
         image_input = gr.Image(type="numpy")