import gradio as gr

# Load multiple models
model1 = gr.load("models/Lykon/dreamshaper-xl-turbo")
model2 = gr.load("models/hakurei/waifu-diffusion")
model3 = gr.load("models/runwayml/stable-diffusion-v1-5")
model4 = gr.load("models/stablediffusionapi/juggernaut-xl-v5")
model5 = gr.load("models/stabilityai/stable-diffusion-xl-base-1.0")

# Function to switch between models
def generate_image(model, text_input):
    return model(text_input)

# Define the models and their names
models = {
    "Model 1": model1,
    "Model 2": model2,
    "Model 3": model3,
    "Model 4": model4,
    "Model 5": model5
}

# Create a dropdown to select the model
model_dropdown = gr.Dropdown(choices=list(models.keys()), label="Select Model")

# Create the input text box
input_text = gr.Textbox(label="Input Text", placeholder="Enter text here")  # Set placeholder instead of default

# Create the output image
output_image = gr.Image()

# Create a function that will be called when the "Generate" button is clicked
def generate_button_click():
    selected_model_name = model_dropdown.value
    selected_model = models[selected_model_name]
    result_image = generate_image(selected_model, input_text.value)
    output_image.image(result_image)

# Create the interface
iface = gr.Interface(
    fn=generate_button_click,
    inputs=[model_dropdown, input_text],
    outputs=output_image,
)

# Launch the interface
iface.launch()
except Exception as e:
    print(f"Error: {e}")