Upload app.py

app.py

@@ -42,6 +42,7 @@ def update_imgbox_gallery(choices):
 
 
 async def infer(model_str, prompt, timeout):
+    from PIL import Image
     noise = ""
     rand = randint(1, 500)
     for i in range(rand):
@@ -52,10 +53,14 @@ async def infer(model_str, prompt, timeout):
         result = await asyncio.wait_for(task, timeout=timeout)
     except (Exception, asyncio.TimeoutError) as e:
         print(e)
-        print(f"Task timed-out: {model_str}")
+        print(f"Task timed out: {model_str}")
         if not task.done(): task.cancel()
         result = None
-    return result
+    if task.done() and result is not None:
+        with lock:
+            image = Image.open(result).convert('RGBA')
+        return image
+    return None
 
 
 def gen_fn(model_str, prompt):
@@ -73,21 +78,6 @@ def gen_fn(model_str, prompt):
     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:
@@ -113,25 +103,26 @@ def gen_fn_gallery(model_str, prompt, gallery):
 
 
 CSS="""
-.output { width=120px; height=120px; !important; }
-.gallery { width=100%; height=512px; max_height=512px !important; }
-.guide { text-align: center; max-width: 1200px; margin: 0 auto; !important; }
+#container { max-width: 1200px; margin: 0 auto; !important; }
+.output { width=112px; height=112px; !important; }
+.gallery { width=100%; min_height=768px; !important; }
+.guide { text-align: center; !important; }
 """
 
-with gr.Blocks(fill_width=True, css=CSS) as demo:
+with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=CSS) as demo:
     with gr.Tab('The Dream'):
         with gr.Column(scale=2):
             txt_input = gr.Textbox(label='Your prompt:', lines=4)
             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)
+                gen_button.click(lambda: gr.update(interactive = True), None, stop_button)
             gr.Markdown("Scroll down to see more images and select models.", elem_classes="guide")
 
         with gr.Column(scale=1):
             with gr.Group():
                 with gr.Row():
-                    output = [gr.Image(label=m, show_download_button=True, elem_classes="output", interactive=False, show_share_button=True, min_width=80, visible=True) for m in default_models]
+                    output = [gr.Image(label=m, show_download_button=True, elem_classes="output", interactive=False, min_width=80, show_share_button=False, 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]
@@ -139,12 +130,11 @@ with gr.Blocks(fill_width=True, css=CSS) as demo:
         with gr.Column(scale=2):
             gallery = gr.Gallery(label="Output", show_download_button=True, elem_classes="gallery",
                                 interactive=False, show_share_button=True, container=True, format="png",
-                                object_fit="cover", columns=2) # , height=768
-                            
+                                preview=True, object_fit="cover", columns=2, rows=2) 
+
         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])
             gen_event = gr.on(triggers=[gen_button.click, txt_input.submit], fn=gen_fn, inputs=[m, txt_input], outputs=[o])
             o.change(add_gallery, [o, m, gallery], [gallery])
             stop_button.click(lambda: gr.update(interactive = False), None, stop_button, cancels = [gen_event])
@@ -155,7 +145,40 @@ with gr.Blocks(fill_width=True, css=CSS) as demo:
                 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)
-        gr.Markdown("Based on the [TestGen](https://huggingface.co/spaces/derwahnsinn/TestGen) Space by derwahnsinn, the [SpacIO](https://huggingface.co/spaces/RdnUser77/SpacIO_v1) Space by RdnUser77 and Omnibus's Maximum Multiplier!")
+
+    with gr.Tab('Single model'):
+        with gr.Column(scale=2):
+            model_choice2 = gr.Dropdown(models, label='Choose model', value=models[0])
+            txt_input2 = gr.Textbox(label='Your prompt:', lines=4)
+            
+            max_images = 6
+            num_images = gr.Slider(1, max_images, value=max_images, step=1, label='Number of images')
+            
+            with gr.Row():
+                gen_button2 = gr.Button('Generate', scale=2)
+                stop_button2 = gr.Button('Stop', variant='secondary', interactive=False, scale=1)
+                gen_button2.click(lambda: gr.update(interactive = True), None, stop_button2)
+
+        with gr.Column(scale=1):
+            with gr.Group():
+                with gr.Row():
+                    output2 = [gr.Image(label = '', show_download_button=True, elem_classes="output", interactive=False, min_width=80, visible=True, show_share_button=False, show_label=False) for _ in range(max_images)]
+                    #output2 = [gr.Image(label = '') 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, container=True, format="png",
+                                preview=True, object_fit="cover", columns=2, rows=2) 
+
+        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)
+            #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)
+            gen_event2 = gr.on(triggers=[gen_button2.click, txt_input2.submit], fn=lambda i, n, m, t: gen_fn(m, t) if (i < n) else None, inputs=[img_i, num_images, model_choice2, txt_input2], outputs=[o])
+            o.change(add_gallery, [o, model_choice2, gallery2], [gallery2])
+            stop_button2.click(lambda: gr.update(interactive = False), None, stop_button2, cancels = [gen_event2])
+
+    gr.Markdown("Based on the [TestGen](https://huggingface.co/spaces/derwahnsinn/TestGen) Space by derwahnsinn, the [SpacIO](https://huggingface.co/spaces/RdnUser77/SpacIO_v1) Space by RdnUser77 and Omnibus's Maximum Multiplier!")
 
 demo.queue()
 demo.launch()