import gradio as gr
from random import randint
from all_models import models
# Função para carregar os modelos e armazená-los em um dicionário
def load_models(models):
models_dict = {}
for model in models:
try:
# Tente carregar o modelo e armazená-lo no dicionário
model_obj = gr.load(f'models/{model}')
models_dict[model] = model_obj
except Exception as e:
# Lidar com exceções ao carregar o modelo
print(f"Erro ao carregar o modelo {model}: {e}")
models_dict[model] = gr.Interface(lambda txt: None, ['text'], ['image'])
return models_dict
# Carregar os modelos e armazená-los na variável global
models_load = load_models(models)
num_models = 150
default_models = models[:num_models]
# Função para estender as escolhas com 'NA' se necessário
def extend_choices(choices):
return choices + ['NA'] * (num_models - len(choices))
# Função para atualizar a caixa de imagens com base nas escolhas
def update_imgbox(choices, output):
choices_plus = extend_choices(choices)
for m, o in zip(choices_plus, output):
o.label = m
o.visible = (m != 'NA')
# Função para gerar imagens com base no modelo e prompt selecionados
def generate_images(model_str, prompt, negative_prompt):
if model_str == 'NA':
return None
noise = str(randint(0, 99999999999))
prompt_with_noise = f'{prompt} {noise}'
if negative_prompt:
prompt_with_noise += ' Negative'
return models_load[model_str](prompt_with_noise)
# Interface Gradio
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)
negative_prompt = gr.Checkbox(label='Negative Prompt', value=False)
gen_button = gr.Button('Generate up to 150 images in up to 3 minutes total')
stop_button = gr.Button('Stop', variant='secondary', interactive=False)
# Evento de clique para iniciar a geração de imagens
gen_button.click(lambda s: gr.update(interactive=True), None, stop_button)
gr.HTML(
"""
Scroll down to see more images and select models.
"""
)
with gr.Row():
output = [gr.Image(label=m, min_width=512) 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):
# Evento de clique para gerar imagens para cada modelo selecionado
gen_event = gen_button.click(generate_images, [m, txt_input, negative_prompt], o)
# Evento de clique para parar a geração de imagens
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 available options!',
value=default_models,
multiselect=True,
max_choices=num_models,
interactive=True,
filterable=False
)
# Evento de mudança para atualizar a caixa de imagens com base nas escolhas
model_choice.change(update_imgbox, model_choice, output)
model_choice.change(extend_choices, model_choice, current_models)
with gr.Row():
gr.HTML(
"""