import gradio as gr from random import randint from all_models import models 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 = 6 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')) for m in choices_plus] def gen_fn(model_str, prompt, negative_prompt=None, image_style="Default"): if model_str == 'NA': return None modified_prompt = prompt if image_style != "Default": modified_prompt += f", {image_style}" if negative_prompt: modified_prompt += f", not {negative_prompt}" return models_load[model_str](modified_prompt) with gr.Blocks() as demo: with gr.Tab('The Dream'): txt_input = gr.Textbox(label='Your prompt:', lines=4) with gr.Accordion("Advanced Settings", open=False) as advanced_settings: neg_prompt = gr.Textbox(label='Negative prompt (Optional):', placeholder='Enter undesirable attributes here', lines=2) image_style = gr.Dropdown(label='Select Style', choices=["Default", "Realistic", "Portrait", "Anime"], value="Default") gen_button = gr.Button('Generate up to 6 images in up to 2 minutes total') stop_button = gr.Button('Stop', variant='secondary', interactive=False) gen_button.click(lambda s: gr.update(interactive=True), None, stop_button, concurrency_limit=10) gr.HTML( """

Scroll down to see more images and select models.

""" ) with gr.Row(): output = [gr.Image(label=m, min_width=460) 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, neg_prompt, image_style], o, concurrency_limit=10) stop_button.click(lambda s: gr.update(interactive=False), None, stop_button, cancels=[gen_event]) with gr.Accordion('Model selection'): model_choice = gr.CheckboxGroup(models, label=f'Choose up to {num_models} different models from the best available!', value=default_models, interactive=True) model_choice.change(update_imgbox, model_choice, output) model_choice.change(extend_choices, model_choice, current_models) with gr.Row(): gr.HTML( """