revised it

app.py

@@ -3,15 +3,12 @@ from transformers import TextIteratorStreamer, AutoProcessor, LlavaForConditiona
 from diffusers import DiffusionPipeline
 import gradio as gr
 import numpy as np
-import accelerate
-import spaces
 from PIL import Image
 import threading
-from openai import OpenAI
+import openai
 import os
-import asyncio
-from typing import Any
 
+# Retrieve the OpenAI API key from the environment
 API_KEY = os.getenv('OPEN_AI_API_KEYS')
 
 DESCRIPTION = '''
@@ -21,7 +18,7 @@ DESCRIPTION = '''
 </div>
 '''
 
-# Llava Installed
+# Initialize the models
 llava_model = LlavaForConditionalGeneration.from_pretrained(
     "xtuner/llava-llama-3-8b-v1_1-transformers",
     torch_dtype=torch.float16,
@@ -29,10 +26,12 @@ llava_model = LlavaForConditionalGeneration.from_pretrained(
 )
 
 llava_model.to("cuda:0")
+
 processor = AutoProcessor.from_pretrained("xtuner/llava-llama-3-8b-v1_1-transformers")
+
 llava_model.generation_config.eos_token_id = 128009
 
-# Stable Diffusor Installed
+# Initialize Stable Diffusion pipelines
 base = DiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     torch_dtype=torch.float16,
@@ -52,27 +51,28 @@ refiner = DiffusionPipeline.from_pretrained(
 refiner.to('cuda')
 
 def multimodal_and_generation(message, history):
+    """
+    Generates a response based on the input message and optionally an image.
+    """
     print(f"Message:\n{message}\nType:\n{type(message)}")
     image_path = None
-    if message["files"]:
-        if isinstance(message["files"][-1], dict):
+    if "files" in message and message["files"]:
+        if type(message["files"][-1]) == dict:
             image_path = message["files"][-1]["path"]
         else:
             image_path = message["files"][-1]
     else:
         for hist in history:
-            if isinstance(hist[0], tuple):
+            if type(hist[0]) == tuple:
                 image_path = hist[0][0]
 
     if image_path is None:
         input_prompt = message["text"]
-        client = OpenAI(api_key=API_KEY)
-        stream = client.chat.completions.create(
+        client = openai.OpenAI(api_key=API_KEY)
+        stream = client.ChatCompletion.create(
             model="gpt-3.5-turbo",
-            messages=[
-                {"role": "system", "content": "You are a helpful assistant called 'chimera'."},
-                {"role": "user", "content": input_prompt}
-            ],
+            messages=[{"role": "system", "content": "You are a helpful assistant called 'chimera'."},
+                      {"role": "user", "content": input_prompt}],
             stream=True,
         )
         return stream
@@ -87,8 +87,11 @@ def multimodal_and_generation(message, history):
         thread.start()
 
         return streamer
-
+    
 def diffusing(prompt):
+    """
+    Generates an image using Stable Diffusion based on the input prompt.
+    """
     image = base(
         prompt=prompt,
         num_inference_steps=40,
@@ -105,124 +108,62 @@ def diffusing(prompt):
 
 def check_cuda_availability():
     if torch.cuda.is_available():
-        result = f"GPU: {torch.cuda.get_device_name(0)}"
-        return result
+        return f"GPU: {torch.cuda.get_device_name(0)}"
     else:
         return "No CUDA device found."
 
 mode = ""
 
-# @spaces.GPU(duration=120)
-# async def bot_comms(message, history):
-#     global mode
-
-#     if message == "check cuda":
-#         result = check_cuda_availability()
-#         yield result
-#         return
-
-#     if message == "imagery":
-#         mode = message
-#         yield "Imagery On! Type your prompt to make the image 🖼️"
-#         return
-
-#     if message == "chatting":
-#         mode = message
-#         yield "Imagery Off. Ask me any questions. ☄️"
-#         return
-
-#     if mode == "imagery":
-#         print("On imagery\n\n")
-#         image = diffusing(
-#             prompt=message,
-#         )
-#         yield image
-#         return
-
-#     if mode == "chatting" or mode == "":
-#         print("On chatting or no mode.\n\n")
-#         stream = multimodal_and_generation(
-#             message=message,
-#             history=history,
-#         )
-#         gpt_outputs = []
-#         async for chunk in stream:
-#             if chunk.choices[0].delta.content is not None:
-#                 text = chunk.choices[0].delta.content
-#                 gpt_outputs.append(text)
-#                 yield "".join(gpt_outputs)
-
-async def bot_comms_async(message, history):
+@spaces.GPU(duration=120)
+def bot_comms(message, history):
+    """
+    Handles communication between Gradio and the models.
+    """
     global mode
 
     if message == "check cuda":
-        result = check_cuda_availability()
-        return [result]
-
+        yield check_cuda_availability()
+        return
+    
     if message == "imagery":
         mode = message
-        return ["Imagery On! Type your prompt to make the image 🖼️"]
-
+        yield "Imagery On! Type your prompt to make the image 🖼️"
+        return
+    
     if message == "chatting":
         mode = message
-        return ["Imagery Off. Ask me any questions. ☄️"]
-
+        yield "Imagery Off. Ask me any questions. ☄️"
+        return
+    
     if mode == "imagery":
         print("On imagery\n\n")
-        image = diffusing(prompt=message)
-        return [image]
-
+        image = diffusing(message)
+        yield image
+        return
+        
+    buffer = ""
+    gpt_outputs = []
     if mode == "chatting" or mode == "":
         print("On chatting or no mode.\n\n")
-        stream = multimodal_and_generation(message=message, history=history)
-        gpt_outputs = []
-        async for chunk in stream:
-            if chunk.choices[0].delta.content is not None:
+        stream = multimodal_and_generation(message, history)
+        for chunk in stream:
+            if hasattr(chunk.choices[0].delta, "content"):
                 text = chunk.choices[0].delta.content
                 gpt_outputs.append(text)
-        return ["".join(gpt_outputs)]
-
-@spaces.GPU(duration=120)
-def bot_comms(message: str, history: Any):
-    return asyncio.run(bot_comms_async(message, history))
-
-# Define your Gradio UI as usual
-import gradio as gr
-
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    with gr.Row():
-        submit = gr.Button("Submit")
-
-    def user(message, history):
-        return "", history + [[message, None]]
-
-    def bot_response(message, history):
-        response = bot_comms(message, history)
-        return history + [[message, response]]
-
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, [msg, chatbot], [msg, chatbot]
-    )
-    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, [msg, chatbot], [msg, chatbot]
+                yield "".join(gpt_outputs)
+
+chatbot = gr.Chatbot(height=600, label="Chimera AI")
+chat_input = gr.MultimodalTextbox(interactive=True, file_types=["images"], placeholder="Enter your question or upload an image.", show_label=False)
+
+with gr.Blocks(fill_height=True) as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.ChatInterface(
+        fn=bot_comms,
+        chatbot=chatbot,
+        fill_height=True,
+        multimodal=True,
+        textbox=chat_input,
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
-
-# chatbot = gr.Chatbot(height=600, label="Chimera AI")
-# chat_input = gr.MultimodalTextbox(interactive=True, file_types=["images"], placeholder="Enter your question or upload an image.", show_label=False)
-# with gr.Blocks(fill_height=True) as demo:
-#     gr.Markdown(DESCRIPTION)
-#     gr.ChatInterface(
-#         fn=bot_comms,
-#         chatbot=chatbot,
-#         fill_height=True,
-#         multimodal=True,
-#         textbox=chat_input,
-#     )
-
-# if __name__ == "__main__":
-#     demo.launch()
+    demo.launch()