import gradio as gr
import os

api_key = os.environ.get("HUGGINGFACE_API_KEY")

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

default_model = "stabilityai/stable-diffusion-2"
model_name = gr.inputs.Dropdown(choices=model_list, label="Select Model", default=default_model)

def generate_image(text, model_name):
    model = gr.Interface.load(model_name, source="huggingface", api_key=api_key)
    return model.predict(text)

input_text = gr.inputs.Textbox(label="Input Text")
output_image = gr.outputs.Image(label="Generated Image")

iface = gr.Interface(
    fn=generate_image,
    inputs=[input_text, model_name],
    outputs=output_image,
    title="Text to Image Generation",
    description="Generate an image from input text using a Hugging Face model."
)

iface.launch()