File size: 2,441 Bytes
c2a8649 e03ee37 6031dc7 9dda882 e03ee37 9dda882 e03ee37 9dda882 c2a8649 9dda882 9af81fd e03ee37 9dda882 e03ee37 9dda882 e03ee37 9dda882 e03ee37 9dda882 f698509 9dda882 e03ee37 216062e 9dda882 e03ee37 9dda882 e03ee37 9dda882 e03ee37 9dda882 e03ee37 9dda882 c2a8649 44afc6e 9dda882 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import gradio as gr
from random import randint
from all_models import models
# Load models
def load_models(models):
models_load = {}
for model in models:
if model not in models_load:
try:
m = gr.load(f'models/{model}')
except Exception as error:
m = gr.Interface(lambda txt: None, ['text'], ['image'])
models_load[model] = m
return models_load
models_load = load_models(models)
num_models = 6
default_models = models[:num_models]
# Extend choices to a fixed number of models
def extend_choices(choices):
return choices + ['NA'] * (num_models - len(choices))
# Dynamically update image boxes based on number of choices
def update_imgbox(choices):
extended_choices = extend_choices(choices)
return [gr.Image(None, label=m, visible=(m != 'NA')) for m in extended_choices]
# Generate function with noise added to prompt
def generate_image(model_str, prompt):
if model_str == 'NA':
return None
noise = str(randint(0, 99999999999))
return models_load[model_str](f'{prompt} {noise}')
# Gradio interface setup
with gr.Blocks() as demo:
model_dropdown = gr.Dropdown(models, label='Choose model', value=models[0], filterable=False)
text_input = gr.Textbox(label='Prompt text')
max_images = 6
num_images_slider = gr.Slider(1, max_images, value=max_images, step=1, label='Number of images')
generate_button = gr.Button('Generate')
stop_button = gr.Button('Stop', variant='secondary', interactive=False)
# Enable the stop button when generation starts
generate_button.click(lambda: gr.update(interactive=True), None, stop_button)
with gr.Row():
output_images = [gr.Image(label='') for _ in range(max_images)]
for i, output in enumerate(output_images):
img_index = gr.Number(i, visible=False)
num_images_slider.change(
lambda idx, n: gr.update(visible=(idx < n)),
[img_index, num_images_slider], output
)
generate_event = generate_button.click(
lambda idx, n, model, prompt: generate_image(model, prompt) if idx < n else None,
[img_index, num_images_slider, model_dropdown, text_input], output
)
# Stop button functionality to cancel image generation
stop_button.click(lambda: gr.update(interactive=False), None, stop_button, cancels=[generate_event])
demo.launch()
|