File size: 3,389 Bytes
374ecac
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7169302
374ecac
 
 
 
 
 
 
 
 
 
527535b
374ecac
 
 
 
 
 
 
 
 
 
 
 
 
c9696a1
374ecac
 
 
a1c23c9
 
374ecac
 
19a47bd
 
b314858
9a9e681
d1e0334
374ecac
b314858
a1c23c9
 
5593057
374ecac
 
 
 
4735bd9
 
 
 
374ecac
4735bd9
 
 
374ecac
4735bd9
374ecac
4735bd9
 
9a9e681
4735bd9
 
374ecac
4735bd9
1c6e17e
374ecac
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import gradio as gr
from random import randint
from all_models import models



def load_fn(models):
    global models_load
    models_load = {}
    
    for model in models:
        if model not in models_load.keys():
            try:
                m = gr.load(f'models/{model}')
            except Exception as error:
                m = gr.Interface(lambda txt: None, ['text'], ['image'])
            models_load.update({model: m})


load_fn(models)


num_models = 6
default_models = models[:num_models]



def extend_choices(choices):
    return choices + (num_models - len(choices)) * ['NA']


def update_imgbox(choices):
    choices_plus = extend_choices(choices)
    return [gr.Image(None, label = m, visible = (m != 'NA')) for m in choices_plus]

    
def gen_fn(model_str, prompt):
    if model_str == 'NA':
        return None
    noise = str(randint(0, 99999999999))
    return models_load[model_str](f'{prompt} {noise}')



with gr.Blocks() as demo:
    with gr.Tab('Multiple models'): 
        with gr.Accordion('Model selection'):
            model_choice = gr.Dropdown(models, label = f'Choose up to {num_models} different models', value = default_models, multiselect = True, max_choices = num_models, interactive = True, filterable = False)

        txt_input = gr.Textbox(label = 'Prompt text')
        gen_button = gr.Button('Generate')
        stop_button = gr.Button('Stop', variant = 'secondary', interactive = False)
        gen_button.click(lambda s: gr.update(interactive = True), None, stop_button)
           
        with gr.Row():
            output = [gr.Image(label = m) for m in default_models]
            current_models = [gr.Textbox(m, visible = False) for m in default_models]
            
            model_choice.change(update_imgbox, model_choice, output, show_progress = False)
            model_choice.change(extend_choices, model_choice, current_models)

            for m, o in zip(current_models, output):
                gen_event = gen_button.click(gen_fn, [m, txt_input], o)
                stop_button.click(lambda s: gr.update(interactive = False), None, stop_button, cancels = [gen_event])
                    
            
    with gr.Tab('Single model'):
        model_choice2 = gr.Dropdown(models, label = 'Choose model', value = models[0], filterable = False)
        txt_input2 = gr.Textbox(label = 'Prompt text')
        
        max_images = 6
        num_images = gr.Slider(1, max_images, value = max_images, step = 1, label = 'Number of images')
        
        gen_button2 = gr.Button('Generate')
        stop_button2 = gr.Button('Stop', variant = 'secondary', interactive = False)
        gen_button2.click(lambda s: gr.update(interactive = True), None, stop_button2)
        
        with gr.Row():
            output2 = [gr.Image(label = '') for _ in range(max_images)]

        for i, o in enumerate(output2):
            img_i = gr.Number(i, visible = False)
            num_images.change(lambda i, n: gr.update(visible = (i < n)), [img_i, num_images], o, show_progress = False)
            gen_event2 = gen_button2.click(lambda i, n, m, t: gen_fn(m, t) if (i < n) else None, [img_i, num_images, model_choice2, txt_input2], o)
            stop_button2.click(lambda s: gr.update(interactive = False), None, stop_button2, cancels = [gen_event2])

    
demo.queue(concurrency_count = 36)                        
demo.launch()