app.py

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()