placed custom_fn for chatbot to handle image outputs

app.py

@@ -4,7 +4,7 @@ from transformers import TextIteratorStreamer, AutoProcessor, LlavaForConditiona
 from diffusers import DiffusionPipeline
 import gradio as gr
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageDraw
 import threading
 import openai
 import os
@@ -109,6 +109,7 @@ def multimodal_and_generation(message, history):
         thread.start()
 
         return streamer
+
     
 def diffusing(prompt):
     """
@@ -127,7 +128,7 @@ def diffusing(prompt):
         image=image
     ).images[0]
     return image
-
+    
 def check_cuda_availability():
     if torch.cuda.is_available():
         return f"GPU: {torch.cuda.get_device_name(0)}"
@@ -136,6 +137,30 @@ def check_cuda_availability():
 
 mode = ""
 
+# # Handling Gradio UI depending on mode
+# def mode_gradio(text: str):
+#     """
+#     Handles the interface gradio will output depending what mode
+#     it's in. It stores the mode in a dictionary that was placed by
+#     a condition inside bot_comms and switches the behaviour in gradio.
+#     """
+
+#     # Mode active
+#     mode_selected = text
+#     return
+    
+# # Image created from diffusing
+# image_created = {}
+
+# # Generate An Image from PIL Image grabbed from diffusing
+# def create_image(image):
+#     """
+#     Generates the image using the Image positions
+#     grabbed from PIL in the diffusing.
+#     """
+#     draw = ImageDraw.Draw(image)
+#     retu
+
 @spaces.GPU(duration=120)
 def bot_comms(message, history):
     """
@@ -166,7 +191,9 @@ def bot_comms(message, history):
         logger.debug("Processing imagery prompt.")
         message = message["text"]
         image = diffusing(message)
-        yield image
+        # mode_gradio("imagery")
+        # image_created["Image"] = image
+        yield ("Generated Image:", image)
         return
         
     buffer = ""
@@ -188,8 +215,16 @@ chat_input = gr.MultimodalTextbox(interactive=True, file_types=["images"], place
 
 with gr.Blocks(fill_height=True) as demo:
     gr.Markdown(DESCRIPTION)
+    # if mode_gradio == "imagery":
+    #     gr.Interface()
+    # Customize chatinterface to handle tuples
+    def custom_fn(*args, **kwargs):
+        result = bot_comms(*args, **kwargs):
+        if isinstance(result, tuple) and isinstance(result[1], Image.Image):
+            return [(result[0], None), (None, result[1])]
+        return result
     gr.ChatInterface(
-        fn=bot_comms,
+        fn=custom_fn,
         chatbot=chatbot,
         fill_height=True,
         multimodal=True,