app.py

app.py

@@ -1,115 +1,221 @@
 import gradio as gr
-import time
+from random import randint
+from all_models import models
 
-models = [
-    "adirik/flux-cinestill",
-    "kudzueye/boreal-flux-dev-v2",
-    "alvdansen/flux_film_foto",
-    "alvdansen/pola-photo-flux",
-    "XLabs-AI/flux-RealismLora",
-    "black-forest-labs/FLUX.1-dev",
-]
+def load_fn(models):
+    global models_load
+    models_load = {}
+    for model in models:
+        if model not in models_load.keys():
+            try:
+                m = gr.load(f'models/{model}')
+            except Exception as error:
+                m = gr.Interface(lambda txt: None, ['text'], ['image'])
+            models_load.update({model: m})
 
-model_functions = {}
-model_idx = 1
-for model_path in models:
-    try:
-        model_functions[model_idx] = gr.load(f"models/{model_path}", live=False)
-    except Exception as error:
-        def the_fn(txt):
-            return None
-        model_functions[model_idx] = gr.Interface(fn=the_fn, inputs=["text"], outputs=["image"])
-    model_idx += 1
+load_fn(models)
 
-def send_it_idx(idx):
-    def send_it_fn(prompt):
-        output = (model_functions.get(str(idx)) or model_functions.get(str(1)))(prompt)
-        return output
-    return send_it_fn
+num_models = len(models)
+default_models = models[:num_models]
 
-def get_prompts(prompt_text):
-    return prompt_text
+def extend_choices(choices):
+    return choices + (num_models - len(choices)) * ['NA']
 
-def clear_it(val):
-    if int(val) != 0:
-        val = 0
-    else:
-        val = 0
-    return val
+def update_imgbox(choices):
+    choices_plus = extend_choices(choices)
+    return [gr.Image(None, label=m, visible=(m != 'NA'), elem_id="custom_image") for m in choices_plus]
 
-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)
-    return d, tog
+def gen_fn(model_str, prompt):
+    if model_str == 'NA':
+        return None
+    noise = str(randint(0, 9999999))
+    return models_load[model_str](f'{prompt} {noise}')
 
-def all_task_start():
-    t = time.gmtime()
-    t_stamp = time.time()
-    current_time = time.strftime("%H:%M:%S", t)
-    return gr.update(value=t_stamp), gr.update(value=t_stamp), gr.update(value=0)
-
-def clear_fn():
-    nn = len(models)
-    return tuple([None, *[None for _ in range(nn)]])
-
-with gr.Blocks(title="SD Models") as my_interface:
-    with gr.Column(scale=12):
-        with gr.Row():
+def make_me():
+    with gr.Row():
+        with gr.Column(scale=1):
+            txt_input = gr.Textbox(label='Your prompt:', lines=3, container=False, elem_id="custom_textbox", placeholder="Prompt")
             with gr.Row():
-                primary_prompt = gr.Textbox(label="Prompt", value="")
-            with gr.Row():
-                run = gr.Button("Run", variant="primary")
-                clear_btn = gr.Button("Clear")
-        with gr.Row():
-            sd_outputs = {}
-            model_idx = 1
-            for model_path in models:
-                with gr.Column(scale=3, min_width=200):
-                    with gr.Box():
-                        sd_outputs[model_idx] = gr.Image(label=model_path)
-                model_idx += 1
-        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
-        )
+                gen_button = gr.Button('Generate images', elem_id="custom_gen_button")
+                stop_button = gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                def on_generate_click():
+                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=True, elem_id="custom_stop_button")
+                
+                def on_stop_click():
+                    return gr.Button('Generate images', elem_id="custom_gen_button"), gr.Button('Stop', variant='secondary', interactive=False, elem_id="custom_stop_button")
+                
+                gen_button.click(on_generate_click, inputs=None, outputs=[gen_button, stop_button])
+                stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
+    
+    with gr.Row():
+        output = [gr.Image(label=m, min_width=250, height=250, elem_id="custom_image") for m in default_models]
+        current_models = [gr.Textbox(m, visible=False) for m in default_models]
+        for m, o in zip(current_models, output):
+            gen_event = gen_button.click(gen_fn, [m, txt_input], o)
+            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button], cancels=[gen_event])
+    
+    with gr.Accordion('Model selection', elem_id="custom_accordion"):
+        model_choice = gr.CheckboxGroup(models, label=f'{num_models} different models selected', value=default_models, interactive=True, elem_id="custom_checkbox_group")
+        model_choice.change(update_imgbox, model_choice, output)
+        model_choice.change(extend_choices, model_choice, current_models)
+    
+    with gr.Row():
+        gr.HTML("")
 
-        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])
-        runs_dict = {}
-        model_idx = 1
-        for model_path in models:
-            runs_dict[model_idx] = run.click(model_functions[model_idx], inputs=[primary_prompt], outputs=[sd_outputs[model_idx]])
-            model_idx += 1
+custom_css = """
+:root {
+    --body-background-fill: #2d3d4f;
+}
+body {
+    background-color: var(--body-background-fill) !important;
+    color: #2d3d4f;
+    margin: 0;
+    padding: 0;
+    font-family: Arial, sans-serif;
+    height: 100vh;
+    overflow-y: auto;
+}
+.gradio-container {
+    background-color: #2d3d4f;
+    color: #c5c6c7;
+    padding: 20px;
+    border-radius: 8px;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+    width: 100%;
+    max-width: 1200px;
+    margin: 20px auto; /* Center the container horizontally */
+    display: block; /* Ensure it's a block element */
+    min-height: 100vh; /* Adjust the height of the container */
+}
+.app_title {
+    background-color: #2d3d4f;
+    color: #c5c6c7;
+    padding: 10px 20px;
+    border-bottom: 1px solid #3b4252;
+    text-align: center;
+    font-size: 24px;
+    font-weight: bold;
+    width: 100%;
+    box-sizing: border-box;
+    margin-bottom: 20px; /* Add margin to separate the header from content */
+}
+.custom_textbox {
+    background-color: #2d343f;
+    border: 1px solid #3b4252;
+    color: #7f8184;
+    padding: 10px;
+    border-radius: 4px;
+    margin-bottom: 10px;
+    width: 100%;
+    box-sizing: border-box;
+}
+.custom_gen_button {
+    background-color: #8b38ff;
+    border: 1px solid #ffffff;
+    color: blue;
+    padding: 15px 32px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin: 4px 2px;
+    cursor: pointer;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+    transition: transform 0.2s, box-shadow 0.2s;
+    border-radius: 4px;
+}
+.custom_gen_button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 10px rgba(0, 0, 0, 0.3);
+}
+.custom_stop_button {
+    background-color: #6200ea;
+    border: 1px solid #ffffff;
+    color: blue;
+    padding: 15px 32px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin: 4px 2px;
+    cursor: pointer;
+    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
+    transition: transform 0.2s, box-shadow 0.2s;
+    border-radius: 4px;
+}
+.custom_stop_button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 6px 10px rgba(0, 0, 0, 0.3);
+}
+.custom_image {
+    border: 1px solid #3b4252;
+    background-color: #2d343f;
+    border-radius: 4px;
+    margin: 10px;
+    max-width: 100%;
+    box-sizing: border-box;
+}
+.custom_accordion {
+    background-color: #2d3d4f;
+    color: #7f8184;
+    border: 1px solid #3b4252;
+    border-radius: 4px;
+    margin-top: 20px;
+    width: 100%;
+    box-sizing: border-box;
+    transition: margin 0.2s ease; /* Smooth transition for margin */
+}
+.custom_accordion .gr-accordion-header {
+    background-color: #2d3d4f;
+    color: #7f8184;
+    padding: 10px 20px;
+    border-bottom: 1px solid #5b6270;
+    cursor: pointer;
+    font-size: 18px;
+    font-weight: bold;
+    height: 40px; /* Fixed height for the header */
+    display: flex;
+    align-items: center;
+}
+.custom_accordion .gr-accordion-header:hover {
+    background-color: #2d3d4f;
+}
+.custom_accordion .gr-accordion-content {
+    padding: 10px 20px;
+    background-color: #2d3d4f;
+    border-top: 1px solid #5b6270;
+    max-height: 0; /* Start with no height */
+    overflow: hidden;
+    transition: max-height 0.2s ease;
+}
+.custom_accordion .gr-accordion-content.open {
+    max-height: 500px;
+}
+.custom_checkbox_group {
+    background-color: #2d343f;
+    border: 1px solid #3b4252;
+    color: #7f8184;
+    border-radius: 4px;
+    padding: 10px;
+    width: 100%;
+    box-sizing: border-box;
+}
+@media (max-width: 768px) {
+   .gradio-container {
+        width: 100%;
+        margin: 0;
+        padding: 10px;
+    }
+   .custom_textbox,.custom_image,.custom_checkbox_group {
+        width: 100%;
+        box-sizing: border-box;
+    }
+}
+"""
 
-        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())]
-        )
+with gr.Blocks(css=custom_css) as demo: 
+    make_me()
 
-my_interface.queue(concurrency_count=200, status_update_rate=1)
-my_interface.launch(inline=True, show_api=False)
+demo.queue(concurrency_count=200)
+demo.launch()