Spaces:
Running
Running
# app.py | |
import torch | |
import gradio as gr | |
from PIL import Image | |
from transformers import Blip2Processor, Blip2ForConditionalGeneration | |
class MiniGPT4Lite: | |
def __init__(self): | |
self.device = torch.device("cpu") # Force CPU for Hugging Face Space | |
self.processor = Blip2Processor.from_pretrained("Salesforce/blip2-flan-t5-xl") | |
self.model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-flan-t5-xl").to(self.device) | |
def chat(self, image, prompt="What is in this image?"): | |
inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device) | |
out = self.model.generate(**inputs, max_new_tokens=100) | |
return self.processor.tokenizer.decode(out[0], skip_special_tokens=True) | |
# Initialize model | |
minigpt = MiniGPT4Lite() | |
def respond(image, prompt): | |
return minigpt.chat(image, prompt) | |
demo = gr.Interface( | |
fn=respond, | |
inputs=[ | |
gr.Image(type="pil", label="Upload an Image"), | |
gr.Textbox(lines=2, placeholder="Ask a question about the image...", label="Prompt") | |
], | |
outputs=gr.Textbox(label="MiniGPT-4 Lite Response"), | |
title="MiniGPT-4 Lite (CPU)", | |
description="Upload an image and ask a question. Powered by BLIP-2 + Flan-T5. Works on CPU." | |
) | |
if __name__ == "__main__": | |
demo.launch() | |