app.py

import gradio as gr
from random import randint
from all_models import models
from datetime import datetime
import time

def get_current_time():
    now = datetime.now()
    now2 = now
    current_time = now2.strftime("%Y-%m-%d %H:%M:%S")
    ki = f'{current_time}'
    return ki

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})

load_fn(models)

num_models = len(models)
default_models = models[:num_models]

def extend_choices(choices):
    return choices + (num_models - len(choices)) * ['NA']

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 gen_fn(model_str, prompt, negative_prompt, max_retries=10, progress=gr.Progress()):
    if model_str == 'NA':
        return None, ""
    
    retries = 0
    feedback = ""
    while retries < max_retries:
        try:
            noise = str(randint(0, 9999999))
            with progress(tracker=True) as pbar:
                for step in range(1, 101):
                    time.sleep(0.05)  # Simulate work being done
                    pbar(step / 100)
            
            if hasattr(models_load[model_str], 'negative_prompt'):
                result = models_load[model_str](f'{prompt} {noise}', negative_prompt=negative_prompt)
            else:
                result = models_load[model_str](f'{prompt} {noise}')
            feedback += f"{get_current_time()} - Attempt {retries + 1}: Success for {model_str}\n"
            return result, feedback
        except Exception as e:
            # Check for specific error messages or status codes
            if "CUDA out of memory" in str(e) or "500" in str(e):
                feedback += f"{get_current_time()} - CUDA out of memory or server error for {model_str}: {e}\n"
            else:
                feedback += f"{get_current_time()} - Error generating image for {model_str}: {e}\n"
            
            retries += 1
            if retries >= max_retries:
                feedback += f"{get_current_time()} - Failed to generate image after {max_retries} retries for {model_str}\n"
                raise Exception(f"Failed to generate image after {max_retries} retries for {model_str}.")
    
    return None, feedback

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")
            negative_txt_input = gr.Textbox(label='Negative prompt:', lines=3, container=False, elem_id="custom_negative_textbox", placeholder="Negative Prompt")
            with gr.Row():
                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 = []
        progress_bars = []
        progress_labels = []
        feedback_box = gr.Textbox(label="Feedback", lines=10, interactive=False, elem_id="custom_feedback_box")
        current_models = [gr.Textbox(m, visible=False) for m in default_models]
        
        for m in default_models:
            if m != 'NA':
                progress_label = gr.Markdown(value=f"**Progress for {m}**")
                progress_bar = gr.Progress()
                output.append(gr.Image(label=m, min_width=250, height=250, elem_id="custom_image"))
                progress_labels.append(progress_label)
                progress_bars.append(progress_bar)
            else:
                output.append(gr.Image(label=m, visible=False, elem_id="custom_image"))
                progress_labels.append(gr.Markdown(value="", visible=False))
                progress_bars.append(gr.Progress(visible=False))
        
        for m, o, p, l, c in zip(current_models, output, progress_bars, progress_labels, default_models):
            gen_event = gen_button.click(gen_fn, [m, txt_input, negative_txt_input, gr.Slider(1, 10, value=10, label="Max Retries"), p], [o, feedback_box])
            stop_button.click(on_stop_click, inputs=None, outputs=[gen_button, stop_button])
    
    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("")

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;
    display: block;
    min-height: 100vh;
}
.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;
}
.custom_textbox, .custom_negative_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;
}
.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;
    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;
    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;
}
.custom_feedback_box {
    background-color: #2d343f;
    border: 1px solid #3b4252;
    color: #7f8184;
    padding: 10px;
    border-radius: 4px;
    margin-top: 10px;
    width: 100%;
    box-sizing: border-box;
    overflow-y: auto;
}
@media (max-width: 768px) {
   .gradio-container {
        width: 100%;
        margin: 0;
        padding: 10px;
    }
   .custom_textbox, .custom_negative_textbox, .custom_image, .custom_checkbox_group {
        width: 100%;
        box-sizing: border-box;
    }
}
"""

with gr.Blocks(css=custom_css) as demo: 
    make_me()

demo.queue(concurrency_count=500)
demo.launch()