import gradio as gr
from random import randint
from all_models import models
from externalmod import gr_Interface_load
import asyncio
from threading import RLock
lock = RLock()
def load_fn(models):
global models_load
models_load = {}
for model in models:
if model not in models_load.keys():
try:
m = gr_Interface_load(f'models/{model}')
except Exception as error:
print(error)
m = gr.Interface(lambda: None, ['text'], ['image'])
models_load.update({model: m})
load_fn(models)
num_models = 6
default_models = models[:num_models]
timeout = 60
def extend_choices(choices):
return choices[:num_models] + (num_models - len(choices[:num_models])) * ['NA']
def update_imgbox(choices):
choices_plus = extend_choices(choices[:num_models])
return [gr.Image(None, label = m, visible = (m != 'NA')) for m in choices_plus]
def update_imgbox_gallery(choices):
choices_plus = extend_choices(choices[:num_models])
return [gr.Gallery(None, label = m, visible = (m != 'NA')) for m in choices_plus]
async def infer(model_str, prompt, timeout):
noise = ""
rand = randint(1, 500)
for i in range(rand):
noise += " "
task = asyncio.create_task(asyncio.to_thread(models_load[model_str], f'{prompt} {noise}'))
await asyncio.sleep(0)
try:
result = await asyncio.wait_for(task, timeout=timeout)
except (Exception, asyncio.TimeoutError) as e:
print(e)
print(f"Task timed-out: {model_str}")
if not task.done(): task.cancel()
result = None
return result
def gen_fn(model_str, prompt):
if model_str == 'NA':
return None
try:
loop = asyncio.new_event_loop()
result = loop.run_until_complete(infer(model_str, prompt, timeout))
except (Exception, asyncio.CancelledError) as e:
print(e)
print(f"Task aborted: {model_str}")
result = None
finally:
loop.close()
return result
def gen_fn_single_gallery(model_str, prompt):
if model_str == 'NA':
return None, gr.update(visible=True)
try:
loop = asyncio.new_event_loop()
result = loop.run_until_complete(infer(model_str, prompt, timeout))
except (Exception, asyncio.CancelledError) as e:
print(e)
print(f"Task aborted: {model_str}")
result = None
finally:
loop.close()
return result, gr.update(visible=True)
def add_gallery(image, model_str, gallery):
if gallery is None: gallery = []
with lock:
gallery.append((image, model_str))
return gallery
def gen_fn_gallery(model_str, prompt, gallery):
if gallery is None: gallery = []
if model_str == 'NA':
yield gallery
try:
loop = asyncio.new_event_loop()
result = loop.run_until_complete(infer(model_str, prompt, timeout))
with lock:
if result: gallery.append(result)
except (Exception, asyncio.CancelledError) as e:
print(e)
print(f"Task aborted: {model_str}")
finally:
loop.close()
yield gallery
CSS="""
.output { width=120px; height=120px !important; }
"""
with gr.Blocks(css=CSS) as demo:
with gr.Tab('The Dream'):
txt_input = gr.Textbox(label = 'Your prompt:', lines=4) #.style(container=False,min_width=1200)
with gr.Row():
gen_button = gr.Button('Generate up to 6 images in up to 3 minutes total', scale=2)
stop_button = gr.Button('Stop', variant = 'secondary', interactive = False, scale=1)
gen_button.click(lambda: 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, show_download_button=True, elem_classes="output", interactive=False, show_share_button=True, visible=True) for m in default_models]
#output = [gr.Image(label = m, show_download_button=True, elem_classes="output", interactive=False, show_share_button=True) for m in default_models]
#output = [gr.Gallery(label = m, show_download_button=True, elem_classes="output", interactive=False, show_share_button=True, container=True, format="png", object_fit="cover") for m in default_models]
current_models = [gr.Textbox(m, visible = False) for m in default_models]
gallery = gr.Gallery(label="Output", show_download_button=True, elem_classes="output", interactive=False, show_share_button=True, container=True, format="png", object_fit="cover")
for m, o in zip(current_models, output):
#gen_event = gen_button.click(gen_fn, [m, txt_input], o)
#gen_event = gen_button.click(gen_fn_gallery, [m, txt_input, o], o)
gen_event = gen_button.click(gen_fn, [m, txt_input], [o])
o.change(add_gallery, [o, m, gallery], [gallery])
stop_button.click(lambda: 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 866 available!', value = default_models, interactive = True)
#model_choice.change(update_imgbox, model_choice, output)
#model_choice.change(update_imgbox_gallery, model_choice, output)
model_choice.change(extend_choices, model_choice, current_models)
with gr.Row():
gr.HTML(
"""