import gradio as gr from random import randint from all_models import models def load_models(model_names): loaded_models = {} for model_name in model_names: try: model = gr.load(f'models/{model_name}') except Exception: model = gr.Interface(lambda txt: "Model load error", inputs=["text"], outputs=["image"]) loaded_models[model_name] = model return loaded_models models_load = load_models(models) num_models = 6 default_models = models[:num_models] def extend_choices(choices, num_models=6): return choices + ['NA'] * (num_models - len(choices)) def update_image_boxes(choices): extended_choices = extend_choices(choices) return [gr.Image(label=model, visible=(model != 'NA')) for model in extended_choices] def generate_image(model_name, prompt): if model_name == 'NA': return None # If you decide to add randomness to the prompt, uncomment the next line # prompt += f" {randint(0, 999999999)}" return models_load[model_name](prompt) with gr.Blocks() as demo: with gr.Tab("The Dream"): txt_input = gr.Textbox(label="Your prompt:", lines=4).style(container=False, min_width=1200) gen_button = gr.Button("Generate up to 6 images") stop_button = gr.Button("Stop", variant="secondary") gen_button.click(lambda _: gr.update(interactive=True), None, stop_button) with gr.Row(): output_images = [gr.Image(label=model, min_width=480) for model in default_models] model_textboxes = [gr.Textbox(model, visible=False) for model in default_models] for model_textbox, output_image in zip(model_textboxes, output_images): gen_event = gen_button.click(generate_image, [model_textbox, txt_input], output_image) stop_button.click(lambda _: gr.update(interactive=False), None, stop_button, cancels=[gen_event]) with gr.Accordion("Model selection"): model_choice = gr.CheckboxGroup(choices=models, label=f"Choose up to {num_models} models from the 686 available!", value=default_models, multiselect=True, max_choices=num_models, interactive=True) model_choice.change(update_image_boxes, model_choice, output_images) model_choice.change(lambda choices: extend_choices(choices, num_models), model_choice, model_textboxes) with gr.Row(): gr.HTML(""" """) demo.queue(concurrency_count=200) demo.launch()