app.py

import gradio as gr
from gradio_client import Client
from PIL import Image
import os
import traceback
import random
import time

# Create Client instances for the repositories
clients = [
    Client("hsuwill000/LCM_SoteMix_OpenVINO_CPU_Space_TAESD"),
    Client("HelloSun/LCM_Dreamshaper_v7-int8-ov")
]

# Counter for image filenames to avoid overwriting
count = 0

# Gradio Interface Function to handle image generation
def infer_gradio(prompt: str):
    global count
    random.seed(time.time())
    # Randomly choose a client
    client = random.choice(clients)
  
    # 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 result (image URL or file path)
        result = client.predict(inputs, api_name="/infer")
        
        # Assuming the result is a URL or path, use it to open the image
        image = Image.open(result)  # If the result is a URL, ensure you download it first
        
        # 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

# 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
    output_image = gr.Image(label="Generated Image")
    
    # Connecting the button click to the image generation function
    run_button.click(infer_gradio, inputs=prompt_input, outputs=output_image)

demo.launch()