Update app.py

app.py

@@ -10,6 +10,7 @@ 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 = {}
@@ -23,7 +24,6 @@ def load_fn(models):
                 m = gr.Interface(lambda: None, ['text'], ['image'])
             models_load.update({model: m})
 
-
 load_fn(models)
 
 
@@ -31,6 +31,8 @@ num_models = 1
 default_models = models[:num_models]
 inference_timeout = 600
 
+MAX_SEED=3999999999
+
 
 
 def extend_choices(choices):
@@ -41,13 +43,57 @@ 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('') #str(randint(0, 99999999999))
     return models_load[model_str](f'{prompt} {noise}')
 
+
+
+
+
+
+
+
+
+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:
+        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
+    finally:
+        loop.close()
+    return result
+
 def gen_fnsix(model_str, prompt):
     if model_str == 'NA':
         return None
@@ -64,7 +110,7 @@ with gr.Blocks() as demo:
     gr.HTML(
     """
         <div>
-        <p> <center>For negative prompts, seed control, and other features visit John6666's <a href="https://huggingface.co/spaces/John6666/PrintingPress4">Printing Press 4</a>!</center>
+        <p> <center>For negative prompts, Width and Height, and other features visit John6666's <a href="https://huggingface.co/spaces/John6666/PrintingPress4">Printing Press 4</a>!</center>
         </p></div>
     """
 )  
@@ -91,7 +137,41 @@ with gr.Blocks() as demo:
             gr.HTML(
     """
         <div class="footer">
-        <p> Based on the <a href="https://huggingface.co/spaces/derwahnsinn/TestGen">TestGen</a> Space by derwahnsinn, the <a href="https://huggingface.co/spaces/RdnUser77/SpacIO_v1">SpacIO</a> Space by RdnUser77, Omnibus's Maximum Multiplier, <a href="https://huggingface.co/spaces/Yntec/Diffusion60XX">Diffusion60XX</a> and <a href="https://huggingface.co/spaces/Yntec/ToyWorld">Toy World</a>!
+        <p> Based on the <a href="https://huggingface.co/spaces/derwahnsinn/TestGen">TestGen</a> Space by derwahnsinn, the <a href="https://huggingface.co/spaces/RdnUser77/SpacIO_v1">SpacIO</a> Space by RdnUser77, Omnibus's Maximum Multiplier, and <a href="https://huggingface.co/spaces/Yntec/ToyWorld">Toy World</a>!
+        </p>
+    """
+)
+    with gr.Tab('Seed it!'):
+        model_choiceseed = gr.Dropdown(models, label = f'Choose a model from the {len(models)} available! Try clearing the box and typing on it to filter them!', value = models[0], filterable = True)
+        txt_inputseed = gr.Textbox(label = 'Your prompt:')
+        seed = gr.Slider(label="Use a seed to replicate the same image later", info="Max 3999999999", minimum=0, maximum=MAX_SEED, step=1, value=1)
+        
+        max_imagesseed = 1
+        num_imagesseed = gr.Slider(1, max_imagesone, value = max_imagesone, step = 1, label = 'One, because more would make it produce identical images with the seed', visible = False)
+        
+        gen_buttonseed = gr.Button('Generate an image using the seed')
+        #stop_button = gr.Button('Stop', variant = 'secondary', interactive = False)
+        gen_button.click(lambda s: gr.update(interactive = True), None)
+        
+        with gr.Row():
+            outputseed = [gr.Image(label = '') for _ in range(max_imagesseed)]
+
+        for i, o in enumerate(outputseed):
+            img_is = gr.Number(i, visible = False)
+            num_imagesseed.change(lambda i, n: gr.update(visible = (i < n)), [img_is, num_imagesseed], o, show_progress = False)
+            #gen_eventseed = gen_buttonseed.click(lambda i, n, m, t, n1: gen_fnseed(m, t, n1) if (i < n) else None, [img_is, num_imagesseed, model_choiceseed, txt_inputseed, useseed], o, concurrency_limit=None, queue=False)
+
+            gen_eventseed = gr.on(triggers=[gen_buttonseed.click, txt_inputseed.submit],
+                               fn=lambda i, n, m, t, n1: gen_fnseed(m, t, n1) if (i < n) else None,
+                               inputs=[img_is, num_imagesseed, model_choiceseed, txt_inputseed, seed], outputs=[o],
+                                       concurrency_limit=None, queue=False) # Be sure to delete ", queue=False" when activating the stop button
+                        
+            #stop_button.click(lambda s: gr.update(interactive = False), None, stop_button, cancels = [gen_event])
+        with gr.Row():
+            gr.HTML(
+    """
+        <div class="footer">
+        <p> Based on the <a href="https://huggingface.co/spaces/derwahnsinn/TestGen">TestGen</a> Space by derwahnsinn, the <a href="https://huggingface.co/spaces/RdnUser77/SpacIO_v1">SpacIO</a> Space by RdnUser77, Omnibus's Maximum Multiplier, and <a href="https://huggingface.co/spaces/Yntec/ToyWorld">Toy World</a>!
         </p>
     """
 )