import os
import gradio as gr

models = [
    "models/stabilityai/stable-diffusion-xl-base-0.9",
    "models/stabilityai/stable-diffusion-2-1",
    "models/stabilityai/stable-diffusion-xl-refiner-0.9",
    "models/stabilityai/stable-diffusion-2-1-base",
    "models/stabilityai/stable-diffusion-2",
    "models/stabilityai/stable-diffusion-2-inpainting",
    "models/stabilityai/stable-diffusion-x4-upscaler",
    "models/stabilityai/stable-diffusion-2-depth",
    "models/stabilityai/stable-diffusion-2-base",
    "models/stabilityai/stable-diffusion-2-1-unclip",
    "models/helenai/stabilityai-stable-diffusion-2-1-base-ov",
    "models/helenai/stabilityai-stable-diffusion-2-1-ov",
    "models/stabilityai/stable-diffusion-2-1-unclip-small"
]

def generate_image(model_name, input_text):
    model = gr.Interface.load(model_name, api_key=os.getenv("HUGGINGFACE_API_KEY"))
    output_image = model.process(input_text)
    return output_image

model_dropdown = gr.inputs.Dropdown(models, label="Select Model")
input_text = gr.inputs.Textbox(lines=5, label="Input Text")
output_image = gr.outputs.Image(label="Output Image")

def change_model(model_name):
    global interface
    interface.close()
    interface = gr.Interface(generate_image, inputs=[model_name, input_text], outputs=output_image)
    interface.launch()

interface = gr.Interface(generate_image, inputs=[model_dropdown, input_text], outputs=output_image, capture_session=True)

change_model_button = gr.Interface.load("models/stabilityai/stable-diffusion-2", capture_session=True).launch()
interface.launch()