Update app.py

app.py

@@ -9,7 +9,7 @@ import gradio as gr
 import numpy as np
 from lxml.html import fromstring
 from transformers import pipeline
-from torch.multiprocessing import Pool, Process, set_start_method
+from torch import multiprocessing as mp, _dynamo
 #from pathos.multiprocessing import ProcessPool as Pool
 #from pathos.threading import ThreadPool as Pool
 #from diffusers.pipelines.flux import FluxPipeline
@@ -20,10 +20,15 @@ from diffusers import DiffusionPipeline, StableDiffusionXLImg2ImgPipeline
 from diffusers.utils import load_image
 #import jax
 #import jax.numpy as jnp
-import torch._dynamo
 
-set_start_method("spawn", force=True)
-torch._dynamo.config.suppress_errors = True
+def init_pool(q):
+    global b
+    b = q
+
+_dynamo.config.suppress_errors = True
+
+mp.set_start_method("spawn", force=True)
+b=mp.Queue()
 
 #pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1", token=os.getenv("hf_token")).to(device)
 #pipe2 = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to(device)
@@ -126,7 +131,9 @@ def tok(txt):
     print(toks)
     return toks
 
-def infer(p1,p2):
+def infer():
+    p1 = b["a"]
+    p2 = b["b"]
     name = generate_random_string(12)+".png"
     _do = ['beautiful', 'playful', 'photographed', 'realistic', 'dynamic poze', 'deep field', 'reasonable coloring', 'rough texture', 'best quality', 'focused']
     if p1 != "":
@@ -218,29 +225,24 @@ with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
             result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
             result.append(gr.Image(interactive=False,elem_classes="image-container", label="Result", show_label=False, type='filepath', show_share_button=False))
 
-    def _ret(p):
-
-        print(f'Starting!')
-        v = infer(p["a"],p["b"])
-        print(f'Finished!')
-        return v
-        
-    def _rets(p1,p2):
+    def main(p1,p2):
         
         p1_en = translate(p1,"english")
         p2_en = translate(p2,"english")
         
         p = {"a":p1_en,"b":p2_en}
-        
         ln = len(result)
         rng = range(ln)
-        p_arr = [p for _ in rng]
-        pool = Pool(processes=ln)
-        lst = list( pool.imap( _ret, p_arr ) )
-        pool.clear()
-        return lst
+        
+        for _ in rng:
+            b.put(p)
+
+        with mp.Pool(ln, initializer=init_pool, initargs=(b,)) as pool:
+            out = pool.map(infer)
+            pool.clear()
+            return list(out.get())
         
         #return list( _ret(p1_en,p2_en) )
         
-    run_button.click(fn=_rets,inputs=[prompt,prompt2],outputs=result)
-    demo.queue().launch(server_port=7861)
+    run_button.click(fn=main,inputs=[prompt,prompt2],outputs=result)
+    demo.queue().launch()