Update app.py

app.py

@@ -1,72 +1,60 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
-import torch
-import os
-import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-os.system('pip install -U bitsandbytes')
-
-
-
-# Define the repository for the quantized model
-repo_name = "cyan2k/molmo-7B-D-bnb-4bit"
-
-# Load processor and model with GPU optimization
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-processor = AutoProcessor.from_pretrained(repo_name, trust_remote_code=True)
-
-# Load model with 4-bit quantization
-model = AutoModelForCausalLM.from_pretrained(repo_name, 
- device_map="auto", 
- torch_dtype=torch.float16, 
- load_in_4bit=True, 
- trust_remote_code=True)
-model.to(device)
-
-def process_image_and_text(image, text):
- # Convert numpy image to PIL format
- pil_image = Image.fromarray(image)
-
- # Process image and text with processor
- inputs = processor(images=[pil_image], text=text, return_tensors="pt").to(device)
-
- # Generate output using the model
- output = model.generate(**inputs, max_new_tokens=200)
-
- # Decode the generated output
- generated_text = processor.decode(output[0], skip_special_tokens=True)
+# Load the model and processor
+repo_name = "cyan2k/molmo-7B-O-bnb-4bit"
+arguments = {
+ "device_map": "auto",
+ "torch_dtype": "auto",
+ "trust_remote_code": True,
+ "load_in_8bit": True # Use 8-bit for reduced memory footprint
+}
+
+# Load the processor and model
+processor = AutoProcessor.from_pretrained(repo_name, **arguments)
+model = AutoModelForCausalLM.from_pretrained(repo_name, **arguments)
+
+def describe_image(image):
+ # Process the uploaded image
+ inputs = processor.process(
+ images=[image],
+ text="Describe this image in great detail."
+ )
+ 
+ # Move inputs to model device
+ inputs = {k: v.to(model.device) for k, v in inputs.items()} # Removed unsqueeze(0) to keep batch size
+
+ # Generate output
+ output = model.generate_from_batch(
+ inputs,
+ GenerationConfig(max_new_tokens=200, 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
 
-def chatbot(image, text, history):
- # Check if the image is uploaded
- if image is None:
- return history + [("Please upload an image first.", None)]
-
- # Get response by processing the image and text
- response = process_image_and_text(image, text)
-
- # Append question and response to the chat history
- history.append((text, response))
- return 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")
-
- state = gr.State([])
-
- # Connect the submit button and textbox to the chatbot function
- submit_button.click(fn=chatbot, inputs=[image_input, text_input, state], outputs=chatbot_output)
- text_input.submit(fn=chatbot, inputs=[image_input, text_input, state], outputs=chatbot_output)
 
-# Launch the Gradio app with GPU
-demo.launch(share=True)
+def gradio_app():
+ # Define Gradio interface
+ image_input = gr.Image(type="pil", label="Upload Image")
+ output_text = gr.Textbox(label="Image Description", interactive=False)
+
+ # Create Gradio interface
+ interface = gr.Interface(
+ fn=describe_image,
+ inputs=image_input,
+ outputs=output_text,
+ title="Image Description App",
+ description="Upload an image and get a detailed description using the Molmo 7B model"
+ )
+ 
+ # Launch the interface
+ interface.launch()
+
+# Launch the Gradio app
+gradio_app()