Update app.py

app.py

@@ -4,33 +4,31 @@ from PIL import Image
 import os
 import time
 import traceback
+import asyncio
 
+# Your Hugging Face API key (ensure this is set in your environment or replace directly)
 api_key = os.getenv('MY_API_KEY')
 
+# List of repos (private spaces)
 repos = [
     "hsuwill000/LCM_SoteMix_OpenVINO_CPU_Space_TAESD",
     "hsuwill000/LCM_SoteMix_OpenVINO_CPU_Space_TAESD_0"
 ]
+
 class CustomClient(Client):
     def __init__(self, *args, timeout=30, **kwargs):
         super().__init__(*args, **kwargs)
         self.timeout = timeout
 
     def _request(self, method, url, **kwargs):
-        # 设置 timeout 参数
         kwargs['timeout'] = self.timeout
         return super()._request(method, url, **kwargs)
-   
+
 # Counter for image filenames to avoid overwriting
 count = 0
-repo_index = 0  # This will keep track of the current repository
 
-# Gradio Interface Function to handle image generation
-def infer_gradio(prompt: str):
-    global count, repo_index
-    # Create a Client instance to communicate with the Hugging Face space
-    client = CustomClient(repos[repo_index], hf_token=api_key,timeout=300)
-    
+async def infer_single_gradio(client, prompt):
+    global count
     # Prepare the inputs for the prediction
     inputs = {
         "prompt": prompt,
@@ -39,7 +37,7 @@ def infer_gradio(prompt: str):
 
     try:
         # Send the request to the model and receive the image
-        result = client.predict(inputs, api_name="/infer")
+        result = await asyncio.to_thread(client.predict, inputs, api_name="/infer")
         
         # Open the resulting image
         image = Image.open(result)
@@ -54,9 +52,6 @@ def infer_gradio(prompt: str):
         image.save(filename)
         print(f"Saved image as {filename}")
         
-        # Increment the repo_index to choose the next repository in the list
-        repo_index = (repo_index + 1) % len(repos)  # Cycle through repos list
-        
         # Return the image to be displayed in Gradio
         return image
     
@@ -67,12 +62,25 @@ def infer_gradio(prompt: str):
         traceback.print_exc()  # Print stack trace for debugging
         return None  # Return nothing if an error occurs
 
+async def infer_gradio(prompt: str):
+    # Create a list of tasks (one for each repo)
+    tasks = []
+    for repo in repos:
+        # Create a CustomClient instance for each repo
+        client = CustomClient(repo, hf_token=api_key, timeout=300)
+        task = infer_single_gradio(client, prompt)
+        tasks.append(task)
+    
+    # Run all tasks concurrently (i.e., generate images from all repos)
+    results = await asyncio.gather(*tasks)
+    return results  # Return all the images as a list
+
 # Define Gradio Interface
 with gr.Blocks() as demo:
     with gr.Row():  # Use a Row to place the prompt input and the button side by side
         prompt_input = gr.Textbox(
             label="Enter Your Prompt", 
-            show_label = "False",
+            show_label="False",
             placeholder="Type your prompt for image generation here",
             lines=1,  # Set the input to be only one line tall
             interactive=True  # Allow user to interact with the textbox
@@ -81,10 +89,11 @@ with gr.Blocks() as demo:
         # Change the button text to "RUN:" and align it with the prompt input
         run_button = gr.Button("RUN")
     
-    # Output image display area
-    output_image = gr.Image(label="Generated Image")
+    # Output image display area (will show multiple images)
+    output_images = gr.Gallery(label="Generated Images", elem_id="gallery", show_label=False)
     
     # Connecting the button click to the image generation function
-    run_button.click(infer_gradio, inputs=prompt_input, outputs=output_image)
+    run_button.click(infer_gradio, inputs=prompt_input, outputs=output_images)
 
+# Launch Gradio app
 demo.launch()