Update app.py

app.py

@@ -1,144 +1,135 @@
 import gradio as gr
-# import os
-# import sys
-# from pathlib import Path
-import time
-
-models =[
-    "Yntec/elldrethSDreamMix",
-    "Yntec/ResidentCNZCartoon3D",
-    "Yntec/Based64",
-    "Yntec/GoFish",
-    "Yntec/IsThisArt",
-    "Yntec/mixRealisticFantasy",
-    "Yntec/iffyMix",
-    "Yntec/GoldenEra",
-    "Yntec/Hassanim",
-]
-
-
-model_functions = {}
-model_idx = 1
-for model_path in models:
+from random import randint
+from all_models import models
+
+from externalmod import gr_Interface_load, randomize_seed
+
+import asyncio
+import os
+from threading import RLock
+lock = RLock()
+HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.
+
+
+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}', hf_token=HF_TOKEN)
+            except Exception as error:
+                print(error)
+                m = gr.Interface(lambda: None, ['text'], ['image'])
+            models_load.update({model: m})
+
+
+load_fn(models)
+
+
+num_models = 9
+
+default_models = models[:num_models]
+inference_timeout = 600
+MAX_SEED=666666666
+starting_seed = randint(666666000, 666666666)
+
+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]
+
+async def infer(model_str, prompt, seed=1, timeout=inference_timeout):
+    from pathlib import Path
+    kwargs = {}
+    noise = ""
+    kwargs["seed"] = seed
+    task = asyncio.create_task(asyncio.to_thread(models_load[model_str].fn,
+                               prompt=f'{prompt} {noise}', **kwargs, token=HF_TOKEN))
+    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
+    if task.done() and result is not None:
+        with lock:
+            png_path = "image.png"
+            result.save(png_path)
+            image = str(Path(png_path).resolve())
+            return image
+    return None
+
+
+
+
+def gen_fnseed(model_str, prompt, seed=1):
+    if model_str == 'NA':
+        return None
     try:
-        model_functions[model_idx] = gr.Interface.load(f"models/{model_path}", live=False, preprocess=True, postprocess=False)
-    except Exception as error:
-        def the_fn(txt):
-            return None
-        model_functions[model_idx] = gr.Interface(fn=the_fn, inputs=["text"], outputs=["image"])
-    model_idx+=1
-
-
-def send_it_idx(idx):
-    def send_it_fn(prompt):
-        output = (model_functions.get(str(idx)) or model_functions.get(str(1)))(prompt)
-        return output
-    return send_it_fn
-
-def get_prompts(prompt_text):
-    return prompt_text
-
-def clear_it(val):
-    if int(val) != 0:
-        val = 0
-    else:
-        val = 0
-        pass
-    return val
-
-def all_task_end(cnt,t_stamp):
-    to = t_stamp + 360
-    et = time.time()
-    if et > to and t_stamp != 0:
-        d = gr.update(value=0)
-        tog = gr.update(value=1)
-        #print(f'to: {to}  et: {et}')
-    else:
-        if cnt != 0:
-            d = gr.update(value=et)
-        else:
-            d = gr.update(value=0)
-        tog = gr.update(value=0)
-        #print (f'passing:  to: {to}  et: {et}')
-        pass
-    return d, tog
-
-def all_task_start():
-    print("\n\n\n\n\n\n\n")
-    t = time.gmtime()
-    t_stamp = time.time()
-    current_time = time.strftime("%H:%M:%S", t)
-    return gr.update(value=t_stamp), gr.update(value=t_stamp), gr.update(value=0)
-
-def clear_fn():
-    nn = len(models)
-    return tuple([None, *[None for _ in range(nn)]])
-
-
-
-with gr.Blocks(title="SD Models") as my_interface:
-    with gr.Column(scale=12):
-        # with gr.Row():
-        #     gr.Markdown("""- Primary prompt: 你想画的内容(英文单词，如 a cat, 加英文逗号效果更好；点 Improve 按钮进行完善)\n- Real prompt: 完善后的提示词，出现后再点右边的 Run 按钮开始运行""")
+        loop = asyncio.new_event_loop()
+        result = loop.run_until_complete(infer(model_str, prompt, seed, inference_timeout))
+    except (Exception, asyncio.CancelledError) as e:
+        print(e)
+        print(f"Task aborted: {model_str}")
+        result = None
+        with lock:
+            image = "https://huggingface.co/spaces/Yntec/ToyWorld/resolve/main/error.png"
+        result = image
+    finally:
+        loop.close()
+    return result
+
+with gr.Blocks(theme='Yntec/HaleyCH_Theme_Blue_Teal') as demo:
+    with gr.Tab('🤗 2024 is the year this space was launched not the number of models! (they were brand new back then!) 🤗'): 
+        txt_input = gr.Textbox(label='Your prompt:', lines=4)
+        gen_button = gr.Button('Generate up to 9 images in up to 3 minutes total')
+        with gr.Row():
+            seed = gr.Slider(label="Use a seed to replicate the same image later (maximum 666666666)", minimum=0, maximum=MAX_SEED, step=1, value=starting_seed, scale=3)
+            seed_rand = gr.Button("Randomize Seed 🎲", size="sm", variant="secondary", scale=1)    
+        seed_rand.click(randomize_seed, None, [seed], queue=False)
+        #stop_button = gr.Button('Stop', variant = 'secondary', interactive = False)
+               
+        gen_button.click(lambda s: gr.update(interactive = True), None)
+        gr.HTML(
+        """
+            <div style="text-align: center; max-width: 1200px; margin: 0 auto;">
+              <div>
+                <body>
+                <div class="center"><p style="margin-bottom: 10px; color: #000000;">Scroll down to see more images and select models.</p>
+                </div>
+                </body>
+              </div>
+            </div>
+        """
+               )
         with gr.Row():
-            with gr.Row(scale=6):
-                primary_prompt=gr.Textbox(label="Prompt (2024 is the year this space was released, not the amount of models)", value="")
-                # real_prompt=gr.Textbox(label="Real prompt")
-            with gr.Row(scale=6):
-                # improve_prompts_btn=gr.Button("Improve")
-                with gr.Row():
-                    run=gr.Button("Run",variant="primary")
-                    clear_btn=gr.Button("Clear")
+            output = [gr.Image(label = m, min_width=480) 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):
+                gen_event = gr.on(triggers=[gen_button.click, txt_input.submit], fn=gen_fnseed,
+                            inputs=[m, txt_input, seed], outputs=[o], concurrency_limit=None, queue=False)
+                #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 = 'Untick the models you will not be using', value=default_models, interactive=True)
+            #model_choice = gr.CheckboxGroup(models, label = f'Choose up to {num_models} different models from the 2 available! Untick them to only use one!', value = default_models, multiselect = True, max_choices = num_models, interactive = True, filterable = False)
+            model_choice.change(update_imgbox, model_choice, output)
+            model_choice.change(extend_choices, model_choice, current_models)
         with gr.Row():
-            sd_outputs = {}
-            model_idx = 1
-            for model_path in models:
-                with gr.Column(scale=3, min_width=320):
-                    with gr.Box():
-                        sd_outputs[model_idx] = gr.Image(label=model_path)
-                    pass
-                model_idx += 1
-                pass
-            pass
-
-        with gr.Row(visible=False):
-            start_box=gr.Number(interactive=False)
-            end_box=gr.Number(interactive=False)
-            tog_box=gr.Textbox(value=0,interactive=False)
-
-        start_box.change(
-            all_task_end,
-            [start_box, end_box],
-            [start_box, tog_box],
-            every=1,
-            show_progress=True)
-
-        primary_prompt.submit(all_task_start, None, [start_box, end_box, tog_box])
-        run.click(all_task_start, None, [start_box, end_box, tog_box])
-        runs_dict = {}
-        model_idx = 1
-        for model_path in models:
-            runs_dict[model_idx] = run.click(model_functions[model_idx], inputs=[primary_prompt], outputs=[sd_outputs[model_idx]])
-            model_idx += 1
-            pass
-        pass
-
-        # improve_prompts_btn_clicked=improve_prompts_btn.click(
-        #     get_prompts,
-        #     inputs=[primary_prompt],
-        #     outputs=[primary_prompt],
-        #     cancels=list(runs_dict.values()))
-        clear_btn.click(
-            clear_fn,
-            None,
-            [primary_prompt, *list(sd_outputs.values())],
-            cancels=[*list(runs_dict.values())])
-        tog_box.change(
-            clear_it,
-            tog_box,
-            tog_box,
-            cancels=[*list(runs_dict.values())])
-
-my_interface.queue(concurrency_count=600, status_update_rate=1)
-my_interface.launch(inline=True, show_api=False)
- 
+            gr.HTML(
+    """
+        <div class="footer">
+        <p> For more than a hundred times more models (that's not a typo) check out <a href="https://huggingface.co/spaces/Yntec/ToyWorld">Toy World</a>!</a>
+        </p>
+    """
+)
+
+demo.queue(default_concurrency_limit=200, max_size=200)
+demo.launch(show_api=False, max_threads=400)