Update app.py

app.py

@@ -12,6 +12,11 @@ import re
 import time 
 import torch
 import cv2
+from gradio_client import Client, file
+
+def image_gen(prompt):
+    client = Client("KingNish/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"
 
@@ -148,13 +153,13 @@ def respond(message, history):
                         yield output
             elif json_data["name"] == "image_generation":
                 query = json_data["arguments"]["query"]
+                try:
+                    number_of_image = json_data["arguments"]["number_of_image"]
+                except:
+                    number_of_image = 1
                 gr.Info("Generating Image, Please wait 10 sec...")
-                client = InferenceClient("stabilityai/stable-diffusion-3-medium-diffusers")
-                seed = random.randint(0,999999)
-                image = client.text_to_image(f"{str(query)} {message_text}", negative_prompt=f"{seed}")
-                generated_images = gr.Image(image)
-                yield generated_images
-                gr.Info("We are going to more upgrade image generator in next update")
+                image = image_gen(f"{str(query)}")
+                yield gr.Image(image[1])
             elif json_data["name"] == "image_qna":
                 inputs = llava(message, history)
                 streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})