app.py

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

# Create a Client instance to communicate with the Hugging Face space
client = Client("hsuwill000/LCM_SoteMix_OpenVINO_CPU_Space_TAESD")

# Counter for image filenames to avoid overwriting
count = 0

# Gradio Interface Function to handle image generation
def infer_gradio(prompt: str):
    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 = 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

# 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(
            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
            elem_id="prompt-input"  # Optional: For CSS styling
        ).style(
            width="80%"  # Set the prompt input width to 4/5
        )
        
        # Change the button text to "RUN:" and align it with the prompt input
        run_button = gr.Button("RUN:").style(
            width="20%"  # Set the button width to 1/5
        )
    
    # 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()