Update app.py

app.py

@@ -14,18 +14,12 @@ import torch
 import cv2
 from gradio_client import Client, file
 
-# The image for the bot and human
-bot_image = "path_to_bot_image.png"
-human_image = "path_to_human_image.png"
-
 def image_gen(prompt):
     client = Client("KingNish/Image-Gen-Pro")
-    return client.predict("Image Generation",None, prompt, api_name="/image_gen_pro")
+    return client.predict("Image Generation", None, prompt, api_name="/image_gen_pro")
 
 model_id = "llava-hf/llava-interleave-qwen-0.5b-hf"
-
 processor = LlavaProcessor.from_pretrained(model_id)
-
 model = LlavaForConditionalGeneration.from_pretrained(model_id)
 model.to("cpu")
 
@@ -34,12 +28,11 @@ def llava(message, history):
         image = message["files"][0]    
     else:
         for hist in history:
-            if type(hist[0])==tuple:
+            if type(hist[0]) == tuple:
                 image = hist[0][0]
         
     txt = message["text"]
         
-    gr.Info("Analyzing image")
     image = Image.open(image).convert("RGB")
     prompt = f"user <image>\n{txt}assistant"
     
@@ -93,7 +86,6 @@ def respond(message, history):
     func_caller = []
 
     user_prompt = message
-    # Handle image processing
     if message["files"]:
         inputs = llava(message, history)
         streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
@@ -161,52 +153,59 @@ def respond(message, history):
                     yield gr.Image(image[1])
                 except:
                     client_sd3 = InferenceClient("stabilityai/stable-diffusion-3-medium-diffusers")
-                    seed = random.randint(0,999999)
+                    seed = random.randint(0, 999999)
                     image = client_sd3.text_to_image(query, negative_prompt=f"{seed}")
                     yield gr.Image(image)
             elif json_data["name"] == "image_qna":
-                inputs = llava(message, history)
+                if "files" in message:
+                    image = message["files"][0]    
+                else:
+                    for hist in history:
+                        if type(hist[0]) == tuple:
+                            image = hist[0][0]
+                
+                txt = json_data["arguments"]["query"]
+
+                image = Image.open(image).convert("RGB")
+                prompt = f"user <image>\n{txt}assistant"
+                
+                inputs = processor(prompt, image, return_tensors="pt")
                 streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
                 generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
 
                 thread = Thread(target=model.generate, kwargs=generation_kwargs)
                 thread.start()
-
+            
                 buffer = ""
                 for new_text in streamer:
                     buffer += new_text
                     yield buffer
-        except Exception as ex:
-            gr.Error(str(ex))
-    
-# Gradio Blocks interface with a custom layout
-with gr.Blocks(css="#chatbot .chat-message.user { border-bottom: none; margin-bottom: 2px; }") as demo:
-    gr.Markdown("# OpenChat Mini 🚀")    
-    chatbot = gr.Chatbot(label="Nithish OpenChat").style(container=False).style(height=700)
-    with gr.Row():
-        with gr.Column(scale=0.95):
-            with gr.Row():
-                txt = gr.Textbox(show_label=False, placeholder="Type your message here...").style(container=False)
-            with gr.Row():
-                file_btn = gr.UploadButton("📁", file_types=["image", "video", "document"])
-            with gr.Row():
-                submit_btn = gr.Button("Send", size="small", variant="primary")
-        with gr.Column(scale=0.05, align="center"):
-            gr.Image.update(human_image, human_image)
-            
+        except:
+            messages = ""
+            for msg in history:
+                messages += f"\nuser\n{str(msg[0])}"
+                messages += f"\nassistant\n{str(msg[1])}"
+            messages += f"\nuser\n{str(message)}"
+            stream = client_llama.text_generation(messages, max_new_tokens=2000, do_sample=True, stream=True, details=True, return_full_text=False)
+            output = ""
+            for response in stream:
+                if not response.token.text == "":
+                    output += response.token.text
+                    yield output
+
+with gr.Blocks() as demo:
+    history = gr.Chatbot(label="ChatGPT Style Chatbot", height=500).style(container=False)
+    input_message = gr.Textbox(label="Message", placeholder="Type your message here...", show_label=False, container=False)
+    input_files = gr.File(label="Attach an image", file_types=["image"], container=False)
+
     with gr.Row():
-        with gr.Column(scale=0.95):
-            with gr.Row():
-                file_btn2 = gr.UploadButton("📁", file_types=["image", "video", "document"]).style(container=False)
-            with gr.Row():
-                bot_img = gr.Image(bot_image, bot_image).style(container=False)
-            with gr.Row():
-                submit_btn2 = gr.Button("Send", size="small", variant="primary")
-
-    txt.submit(respond, [txt, chatbot], [txt, chatbot], scroll_to_output=True)
-    file_btn.upload(respond, [txt, chatbot], [txt, chatbot], scroll_to_output=True)
-    submit_btn.click(respond, [txt, chatbot], [txt, chatbot], scroll_to_output=True)
-    file_btn2.upload(respond, [txt, chatbot], [txt, chatbot], scroll_to_output=True)
-    submit_btn2.click(respond, [txt, chatbot], [txt, chatbot], scroll_to_output=True)
-
-demo.launch(debug=True)
+        submit_button = gr.Button("Submit")
+        clear_button = gr.Button("Clear")
+
+    def send_message(message, files):
+        return respond(message, history)
+    
+    submit_button.click(send_message, inputs=[input_message, input_files], outputs=history)
+    clear_button.click(lambda: [], None, history)
+
+demo.launch()