File size: 7,943 Bytes
0af19e2
 
 
 
 
 
0ed58ab
 
7bb94f0
0ed58ab
 
0af19e2
1310984
0af19e2
0ed58ab
1310984
0af19e2
 
 
 
 
1310984
0af19e2
 
 
 
 
1310984
0af19e2
 
 
 
 
0ed58ab
0af19e2
1310984
0af19e2
 
1310984
0af19e2
 
1310984
 
0af19e2
 
 
 
 
 
 
1310984
 
 
0398bb7
1310984
 
0398bb7
1310984
0ed58ab
1310984
0af19e2
 
 
 
1310984
0af19e2
 
1310984
 
0af19e2
1310984
 
 
0af19e2
1310984
0af19e2
 
 
 
1310984
0af19e2
1310984
0af19e2
1310984
0af19e2
 
 
 
 
 
 
1310984
 
0af19e2
 
1310984
 
 
 
 
0af19e2
 
0ed58ab
 
 
 
 
 
 
 
 
 
30405c2
1310984
 
0ed58ab
0af19e2
 
0ed58ab
 
1310984
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0af19e2
 
1310984
0af19e2
 
 
1310984
0af19e2
 
1310984
 
 
 
0af19e2
 
1310984
 
 
 
0af19e2
1310984
 
 
 
0af19e2
1310984
0af19e2
 
1310984
0af19e2
 
1310984
b4bc419
0af19e2
1310984
0af19e2
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import gradio as gr
from all_models import models
from externalmod import gr_Interface_load, save_image, randomize_seed
import asyncio
import os
from threading import RLock
from datetime import datetime

preSetPrompt = "High fashion studio foto shoot. tall slender 18+ caucasian woman. gorgeous face. photorealistic. f1.4"
negPreSetPrompt = "[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry, text, fuzziness"

lock = RLock()
HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None

def get_current_time():
    return datetime.now().strftime("%y-%m-%d %H:%M:%S")

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

load_fn(models)

num_models = 6
max_images = 6
inference_timeout = 400
default_models = models[:num_models]
MAX_SEED = 2**32 - 1

def extend_choices(choices):
    return choices[:num_models] + (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 random_choices():
    import random
    random.seed()
    return random.choices(models, k=num_models)

async def infer(model_str, prompt, nprompt="", height=0, width=0, steps=0, cfg=0, seed=-1, timeout=inference_timeout):
    kwargs = {"height": height if height > 0 else None, "width": width if width > 0 else None, 
              "num_inference_steps": steps if steps > 0 else None, "guidance_scale": cfg if cfg > 0 else None}
    
    if seed == -1:
        kwargs["seed"] = randomize_seed()
    else:
        kwargs["seed"] = seed

    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn, prompt=prompt, negative_prompt=nprompt, **kwargs, token=HF_TOKEN))
    
    try:
        result = await asyncio.wait_for(task, timeout=timeout)
    except asyncio.TimeoutError as e:
        print(f"Task timed out: {model_str}")
        task.cancel()
        result = None
    except Exception as e:
        print(f"Error generating image: {model_str} - {e}")
        task.cancel()
        result = None
    
    if result and not isinstance(result, tuple):
        png_path = f"{model_str.replace('/', '_')}_{get_current_time()}_{kwargs['seed']}.png"
        with lock:
            image = save_image(result, png_path, model_str, prompt, nprompt, height, width, steps, cfg, kwargs["seed"])
        return image
    return None

def gen_fn(model_str, prompt, nprompt="", height=0, width=0, steps=0, cfg=0, seed=-1):
    loop = asyncio.new_event_loop()
    try:
        result = loop.run_until_complete(infer(model_str, prompt, nprompt, height, width, steps, cfg, seed, inference_timeout))
    except (Exception, asyncio.CancelledError) as e:
        print(f"Error generating image: {e}")
        result = None
        raise gr.Error(f"Task aborted: {model_str}, Error: {e}")
    finally:
        loop.close()
    return result

def add_gallery(image, model_str, gallery):
    if image is not None:
        gallery = [(image, model_str)] + gallery[:5]  # Keep only the latest 6 images
    return gallery

# Interface Layout
CSS = """
.gradio-container { max-width: 1200px; margin: 0 auto; }
.output { width: 112px; height: 112px; }
.gallery { min-width: 512px; min-height: 512px; }
"""

js_func = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""

with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', fill_width=True, css=CSS) as demo:
    gr.HTML("")
    
    with gr.Tab('6 Models'):
        with gr.Column(scale=2):
            with gr.Group():
                txt_input = gr.Textbox(label='Your prompt:', value=preSetPrompt, lines=3, autofocus=1)
                neg_input = gr.Textbox(label='Negative prompt:', value=negPreSetPrompt, lines=1)
                with gr.Accordion("Advanced", open=False):
                    width = gr.Slider(label="Width", maximum=1216, step=32, value=0)
                    height = gr.Slider(label="Height", maximum=1216, step=32, value=0)
                    steps = gr.Slider(label="Number of inference steps", maximum=100, step=1, value=0)
                    cfg = gr.Slider(label="Guidance scale", maximum=30.0, step=0.1, value=0)
                    seed = gr.Slider(label="Seed", minimum=-1, maximum=MAX_SEED, step=1, value=-1)
                    seed_rand = gr.Button("Randomize Seed 🎲", size="sm", variant="secondary")
                    seed_rand.click(randomize_seed, None, [seed], queue=False)
            gen_button = gr.Button(f'Generate up to {num_models} images', variant='primary', scale=3)
            random_button = gr.Button('Randomize Models', variant='secondary', scale=1)
            gr.Markdown("")

        with gr.Row():
            output = [gr.Image(label=m, show_download_button=True, elem_classes="output", interactive=False, width=112, height=112, format="png", visible=True) for m in default_models]
            current_models = [gr.Textbox(m, visible=False) for m in default_models]

        with gr.Column(scale=2):
            gallery = gr.Gallery(label="Output", show_download_button=True, elem_classes="gallery", interactive=False, format="png", preview=True, object_fit="cover", columns=2, rows=2)

        for m, o in zip(current_models, output):
            gen_event = gr.on(triggers=[gen_button.click, txt_input.submit], fn=gen_fn,
                              inputs=[m, txt_input, neg_input, height, width, steps, cfg, seed], outputs=[o])
            o.change(add_gallery, [o, m, gallery], [gallery])

        model_choice = gr.CheckboxGroup(models, label=f'Choose up to {num_models} models', value=default_models)
        model_choice.change(update_imgbox, model_choice, output)
        model_choice.change(extend_choices, model_choice, current_models)
        random_button.click(random_choices, None, model_choice)

    with gr.Tab('Single model'):
        model_choice2 = gr.Dropdown(models, label='Choose model', value=models[0])
        txt_input2 = gr.Textbox(label='Your prompt:', value=preSetPrompt, lines=3, autofocus=1)
        neg_input2 = gr.Textbox(label='Negative prompt:', value=negPreSetPrompt, lines=1)
        num_images = gr.Slider(1, max_images, value=max_images, step=1, label='Number of images')

        gen_button2 = gr.Button('Let the machine hallucinate', variant='primary', scale=2)
        
        output2 = [gr.Image(label='', show_download_button=True, elem_classes="output", interactive=False, width=112, height=112, format="png", show_share_button=False) for _ in range(max_images)]
        
        with gr.Column(scale=2):
            gallery2 = gr.Gallery(label="Output", show_download_button=True, elem_classes="gallery", interactive=False, show_share_button=True, format="png", preview=True, object_fit="cover", columns=2, rows=2)

        for i, o in enumerate(output2):
            num_images.change(lambda i, n: gr.update(visible=(i < n)), [i, num_images], o, queue=False)
            gen_event2 = gr.on(triggers=[gen_button2.click, txt_input2.submit],
                               fn=lambda i, n, m, t1, t2, n1, n2, n3, n4, n5: gen_fn(m, t1, t2, n1, n2, n3, n4, n5) if (i < n) else None,
                               inputs=[i, num_images, model_choice2, txt_input2, neg_input2, height, width, steps, cfg, seed], outputs=[o])
            o.change(add_gallery, [o, model_choice2, gallery2], [gallery2])

    gr.Markdown("")
demo.launch(show_api=False, max_threads=400)