Spaces:
Running
Running
| import gradio as gr | |
| from gradio_client import Client | |
| 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): | |
| kwargs['timeout'] = self.timeout | |
| return super()._request(method, url, **kwargs) | |
| # Counter for image filenames to avoid overwriting | |
| count = 0 | |
| async def infer_single_gradio(client, prompt): | |
| global count | |
| # Prepare the inputs for the prediction | |
| inputs = { | |
| "prompt": prompt, | |
| "num_inference_steps": 10 # Number of inference steps for the model | |
| } | |
| try: | |
| # Send the request to the model and receive the image | |
| result = await asyncio.to_thread(client.predict, inputs, api_name="/infer") | |
| # Open the resulting image | |
| image = Image.open(result) | |
| # Create a unique filename to save the image | |
| filename = f"img_{count:08d}.jpg" | |
| while os.path.exists(filename): | |
| count += 1 | |
| filename = f"img_{count:08d}.jpg" | |
| # Save the image locally | |
| image.save(filename) | |
| print(f"Saved image as {filename}") | |
| # Return the image to be displayed in Gradio | |
| return image | |
| except Exception as e: | |
| # Handle any errors that occur | |
| print(f"An exception occurred: {str(e)}") | |
| print("Stack trace:") | |
| 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", | |
| 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 | |
| ) | |
| # Change the button text to "RUN:" and align it with the prompt input | |
| run_button = gr.Button("RUN") | |
| # 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_images) | |
| # Launch Gradio app | |
| demo.launch() | |