import gradio as gr
import time

def generate_image(model,input_text):
    output_image = model(input_text)  # Générer une image à partir de l'entrée donnée
    return output_image

models = [
    "Yntec/elldrethSDreamMix",
    "Yntec/ResidentCNZCartoon3D",
    "Yntec/GoFish",
    "Yntec/IsThisArt",
    "Yntec/mixRealisticFantasy",
    "Yntec/iffyMix",
    "Yntec/GoldenEra",
    "Yntec/Hassanim",
]

model_functions = {}
for idx, model_path in enumerate(models, 1):
    try:
        print(idx)
        model_functions[idx] = gr.Interface.load(f"models/{model_path}", live=False, preprocess=True, postprocess=False)
    except Exception as error:
        def the_fn(txt):
            return None
        model_functions[idx] = gr.Interface(fn=the_fn, inputs=["text"], outputs=["image"])

def generate_image_api(prompt):
    return generate_image(model_functions[2], prompt)

def get_prompts(prompt_text):
    return prompt_text

def clear_it(val):
    if int(val) != 0:
        val = 0
    else:
        val = 0
        pass
    return val

def all_task_end(cnt, t_stamp):
    to = t_stamp + 360
    et = time.time()
    if et > to and t_stamp != 0:
        d = gr.update(value=0)
        tog = gr.update(value=1)
    else:
        if cnt != 0:
            d = gr.update(value=et)
        else:
            d = gr.update(value=0)
        tog = gr.update(value=0)
        pass
    return d, tog

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

def clear_fn():
    return tuple([None, *[None for _ in range(len(models))]])

with gr.Blocks(title="SD Models") as my_interface:
    with gr.Column(scale=12):
        with gr.Row():
            with gr.Row(scale=6):
                primary_prompt = gr.Textbox(label="Prompt", value="")
            with gr.Row(scale=6):
                with gr.Row():
                    run = gr.Button("Run", variant="primary")
                    clear_btn = gr.Button("Clear")
                    api = gr.Button("Api", variant="primary")
        with gr.Row():
            sd_outputs = {}
            api_image_output = gr.Image(label='Api OutPut')
            for idx, model_path in enumerate(models, 1):
                with gr.Column(scale=3, min_width=320):
                    with gr.Box():
                        sd_outputs[idx] = gr.Image(label=model_path)
        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)

        start_box.change(
            all_task_end,
            [start_box, end_box],
            [start_box, tog_box],
            every=1,
            show_progress=True)

        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])
        api.click(generate_image_api, inputs=[primary_prompt], outputs=[api_image_output], api_name='generate')

        runs_dict = {}
        for idx, model_path in enumerate(models, 1):
            runs_dict[idx] = run.click(model_functions[idx], inputs=[primary_prompt], outputs=[sd_outputs[idx]])

        clear_btn.click(
            clear_fn,
            None,
            [primary_prompt, *list(sd_outputs.values())],
            cancels=[*list(runs_dict.values())])
        tog_box.change(
            clear_it,
            tog_box,
            tog_box,
            cancels=[*list(runs_dict.values())])

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