|
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 |
|
|
|
|
|
return models_load[model_name](prompt) |
|
|
|
with gr.Blocks() as demo: |
|
with gr.Tab("Zexxxi"): |
|
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(""" |
|
<div class="footer"> |
|
<p>zexxxi nsfw model</p> |
|
</div> |
|
""") |
|
|
|
demo.queue(concurrency_count=200) |
|
demo.launch() |
|
|
|
|
