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

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}')
                print(f"Loaded model {model} with interface: {m}")
            except Exception as error:
                m = gr.Interface(lambda txt: None, ['text'], ['image'])
                print(f"Failed to load model {model}: {error}")
            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, guidance_scale, seed, clip_skip, scheduler, randomize_seed):
    if model_str == 'NA':
        return None
    
    if randomize_seed:
        seed = randint(0, 99999999)
    
    retries = 0
    while retries < 10:
        try:
            noise = str(seed)
            full_prompt = f"{prompt} {noise}"
            
            if negative_prompt:
                full_prompt += f" --negative_prompt {negative_prompt}"
            
            if guidance_scale:
                full_prompt += f" --guidance_scale {guidance_scale}"
            
            if clip_skip:
                full_prompt += f" --clip_skip {clip_skip}"
            
            if scheduler:
                full_prompt += f" --scheduler {scheduler}"
            
            result = models_load[model_str](full_prompt)
            return result
        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):
                print(f"CUDA out of memory or server error: {e}")
            else:
                print(f"Error generating image: {e}")
            
            retries += 1
            if retries >= 10:
                raise Exception(f"Failed to generate image after 10 retries.")
    
    return None

def insert_example_prompt(txt_input):
    example_prompt = "Masterpiece, Highest Quality, Best Quality, Eye Catching, Award Winning, General, 1Girl, Long Black Hair, Wavy And Curly Hair, Green Eyes, Square Rimmed Glasses, Slim Bodied, Tall, Lithe, Petite, Smiling At Viewer, Looking At Viewer, Business Casual Clothing"
    return example_prompt

def make_me():
    with gr.Blocks(css=custom_css) as demo:
        # Add text and HTML link at the top
        gr.Markdown("# Before you generate: Illustrious Diffusion uses a different prompting style than Pony Diffusion, relying less on booru tags and more on detail-oriented, natural language style prompts, please refer to the link provided for a guide on Illustrious prompting.")
        gr.HTML('<p>Useful guide for how to prompt for Illustrious: <a href="https://civitai.com/articles/8380/tips-for-illustrious-xl-prompting-updates" target="_blank">https://civitai.com/articles/8380/tips-for-illustrious-xl-prompting-updates</a>.</p>')
        
        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():
            guidance_scale = gr.Slider(minimum=0, maximum=20, step=0.1, label='Guidance Scale', value=7.5)
            seed = gr.Slider(minimum=0, maximum=99999999, step=1, label='Seed', value=randint(0, 99999999))
            clip_skip = gr.Slider(minimum=0, maximum=5, step=1, label='CLIP Skip', value=1)
            scheduler = gr.Dropdown(['DDIM', 'PNDM', 'LMSDiscrete', 'EulerAncestralDiscrete', 'DPMSolverMultistep', 'HeunDiscrete', 'EulerDiscrete', 'DPMSolverSinglestep', 'DEISMultistep', 'UniPCMultistep'], label='Scheduler', value='DDIM')
            randomize_seed = gr.Checkbox(label='Randomize Seed', value=False)
        
        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, negative_txt_input, guidance_scale, seed, clip_skip, scheduler, randomize_seed], 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("")
        
        # Add the clickable element at the bottom
        with gr.Row():
            insert_example_button = gr.Button('Insert Example Prompt', variant='primary')
            insert_example_button.click(insert_example_prompt, inputs=[txt_input], outputs=[txt_input])

    return demo

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;
}
@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;
    }
}
"""

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