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

repos = [
    "hsuwill000/LCM_SoteMix_OpenVINO_CPU_Space_TAESD",
    "HelloSun/LCM_Dreamshaper_v7-int8-ov"
]

# 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 = Client(repos[repo_index])
    
    # 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}")
        
        # 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
    
    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()