added chat functionality with the model from UI

app.py

@@ -2,7 +2,6 @@ import torch
 import gradio as gr
 from transformers import TextIteratorStreamer, AutoProcessor, LlavaForConditionalGeneration
 from PIL import Image
-import requests
 import threading
 import spaces
 import accelerate
@@ -23,12 +22,39 @@ model = LlavaForConditionalGeneration.from_pretrained(
 
 processor = AutoProcessor.from_pretrained(model_id)
 
+model.generation_config.eos_token_id = 128009
+
 @spaces.GPU(duration=120)
-def krypton(input_image):
+def krypton(input,
+            history):
+    """
+    Recieves inputs (prompts with images if they were added),
+    the image is formated for pil and prompt is formated for the model,
+    to place it's output to the user, these prompts and images are passed in
+    the processor and generation of the model, than the output is decoded from the processor,
+    onto the UI.
+    """
+    if input["files"]:
+        if type(input["files"][-1]) == dict:
+            image = input["files"][-1]["path"]
+        else:
+            image = input["files"][-1]
+    else:
+        # If no images were passed now, look at the past images to keep up as reference still to the prompts
+        # kept inside in tuples, the last one
+        for hist in history:
+            if type(hist[0]) == tuple:
+                image = hist[0][0]
+    try:
+        if image is None:
+            gr.Error("You need to upload an image please for krypton to work.")
+    except NameError:
+        # Image is not defined at all
+        gr.Error("Uplaod an image for Krypton to work")
 
-    pil_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
     # image = Image.open(requests.get(url, stream=True).raw)
-    prompt = ("<|start_header_id|>user<|end_header_id|>\n\n<image>\nWhat are these?<|eot_id|>"
+    prompt = ("<|start_header_id|>user<|end_header_id|>\n\n<image>\n{input['text']}<|eot_id|>"
             "<|start_header_id|>assistant<|end_header_id|>\n\n")
     inputs = processor(prompt, pil_image, return_tensors='pt').to('cuda', torch.float16)
     outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False)
@@ -36,12 +62,15 @@ def krypton(input_image):
     print(output_text)
     return output_text
 
+chatbot=gr.Chatbot(fill_height=600, label="Krypt AI")
+chat_input = gr.MultimodalTextbox(Interactive=True, file_types=["image"], placeholder="Enter your question or upload an image.", show_label=False)
 with gr.Blocks(fill_height=True) as demo:
     gr.Markdown(DESCRIPTION)
-    gr.Interface(
+    gr.ChatInterface(
         fn=krypton,
-        inputs="image",
-        outputs="text",
+        chatbot=chatbot,
+        multimodal=True,
+        textbox=chat_input,
         fill_height=True
     )