multimodal-bot

app.py

@@ -1,58 +1,38 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
 import gradio as gr
-import os
-import time
 
-def print_like_dislike(x: gr.LikeData):
-    print(x.index, x.value, x.liked)
+model = AutoModelForCausalLM.from_pretrained(
+    "MILVLG/imp-v1-3b", 
+    torch_dtype=torch.float16, 
+    device_map="auto",
+    trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("MILVLG/imp-v1-3b", trust_remote_code=True)
 
 
-def add_text(history, text):
-    history = history + [(text, None)]
-    return history, gr.Textbox(value="", interactive=False)
+def generate_answer(text, image):
 
+    input_ids = tokenizer(text, return_tensors='pt').input_ids
+    image_tensor = model.image_preprocess(image)
 
-def add_file(history, file):
-    history = history + [((file.name,), None)]
-    return history
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=100,
+        images=image_tensor,
+        use_cache=True)[0]
+    
+    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
 
+text_input = gr.Textbox(lines=5, label="Enter text")
+image_input = gr.Image(shape=(224, 224), label="Upload Image")
 
-def bot(history):
-    response = "**That's cool!**"
-    history[-1][1] = ""
-    for character in response:
-        history[-1][1] += character
-        time.sleep(0.05)
-        yield history
+iface = gr.Interface(
+    fn=generate_answer, 
+    inputs=[text_input, image_input], 
+    outputs="text",
+    title="DD360-Bot-Multimodal",
+    description="Enter text and upload an image to receive a response from the chatbot."
+)
 
-
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(
-        [],
-        elem_id="chatbot",
-        bubble_full_width=False,
-        avatar_images=(None, (os.path.join(os.path.abspath(''), "avatar.png"))),
-    )
-
-    with gr.Row():
-        txt = gr.Textbox(
-            scale=4,
-            show_label=False,
-            placeholder="Enter text and press enter, or upload an image",
-            container=False,
-        )
-        btn = gr.UploadButton("📁", file_types=["image", "video", "audio"])
-
-    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-        bot, chatbot, chatbot, api_name="bot_response"
-    )
-    txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
-    file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    )
-
-    chatbot.like(print_like_dislike, None, None)
-
-
-demo.queue()
-if __name__ == "__main__":
-    demo.launch()
+iface.launch()