import gradio as gr
import time


model1 = gr.Interface.load("models/pimpilikipilapi1/NSFW_master", preprocess=True, postprocess=False)
model2 = gr.Interface.load("models/prashanth970/flux-lora-uncensored", preprocess=True, postprocess=False)
model3 = gr.Interface.load("models/DiegoJR1973/NSFW-TrioHMH-Flux", preprocess=True, postprocess=False)

models = {
    1: model1,
    2: model2,
    3: model3
}

def send_it_idx(idx):
    def send_it_fn(prompt):
        model = models.get(idx)
        if model:
            return model(prompt)
        return None
    return send_it_fn

def clear_fn():
    return None, None, None

with gr.Blocks(title="SD Models") as my_interface:
    with gr.Column(scale=12):
        with gr.Row():
            primary_prompt = gr.Textbox(label="Prompt", value="")
            run = gr.Button("Run", variant="primary")
            clear_btn = gr.Button("Clear")

        sd_outputs = {}
        for idx, model in models.items():
            with gr.Column(scale=3, min_width=200):
                with gr.Box():
                    sd_outputs[idx] = gr.Image(label=f"Model {idx}")

        with gr.Row(visible=False):
            start_box = gr.Number(interactive=False)
            end_box = gr.Number(interactive=False)
            tog_box = gr.Textbox(value=0, interactive=False)

        def all_task_start():
            t_stamp = time.time()
            return gr.update(value=t_stamp), gr.update(value=t_stamp), gr.update(value=0)

        primary_prompt.submit(all_task_start, None, [start_box, end_box, tog_box])
        run.click(all_task_start, None, [start_box, end_box, tog_box])

        for idx, model in models.items():
            run.click(model, inputs=[primary_prompt], outputs=[sd_outputs[idx]])

        clear_btn.click(
            clear_fn,
            None,
            [primary_prompt, *list(sd_outputs.values())]
        )

my_interface.queue(concurrency_count=100, status_update_rate=1)
my_interface.launch(inline=True, show_api=False)