import gradio as gr


models = {
    "Flux Lora": "models/prashanth970/flux-lora-uncensored",
    "TrioHMH Flux": "models/DiegoJR1973/NSFW-TrioHMH-Flux",
    "Master": "models/pimpilikipilapi1/NSFW_master"
}


def generate_image(text, model_name):
    model_path = models[model_name]
    model = gr.load(model_path)
    result_image = model(text)
    return result_image


interface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Type here your imagination:", placeholder="Type your description here..."),
        gr.Dropdown(label="Select Model", choices=list(models.keys()), type="value", default="Flux Lora")
    ],
    outputs=gr.Image(label="Generated Image"),
    theme="NoCrypt/miku",
    description="Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.",
)

interface.launch()