Update app.py

app.py

@@ -28,7 +28,7 @@ from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, Eul
 #import jax
 #import jax.numpy as jnp
 from numba import cuda, njit as cpu, void, int64 as int, float64 as float, boolean as bool
-from numba.cuda import autojit as gpu, grid
+from numba.cuda import jit as gpu, grid
 from numba.types import unicode_type as string
 # logging
 
@@ -50,7 +50,6 @@ root.addHandler(handler2)
 
 last_motion=None
 dtype = torch.float16
-result=[]
 device = "cuda"
 #repo = "ByteDance/AnimateDiff-Lightning"
 #ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
@@ -122,7 +121,7 @@ def run(cmd):
         sys.exit()
     return str(result.stdout)
     
-@gpu()
+@gpu(string(string,string),device=True,inline=True)
 def translate(text,lang):
 
     if text == None or lang == None:
@@ -165,60 +164,61 @@ def translate(text,lang):
     print(ret)
     return ret
     
-@gpu()
+@gpu(string(int),device=True,inline=True)
 def generate_random_string(length):
     
     characters = string.ascii_letters + string.digits
     return ''.join(random.choice(characters) for _ in range(length))
     
-@gpu()
-def Piper(image,positive,negative,motion):
+@gpu(void(),device=True,inline=True)
+def Piper():
 
     global last_motion
     global ip_loaded
+    global out
 
-    if last_motion != motion:
+    x = grid(1)
+
+    if last_motion != pinp["motion"]:
         pipe.unload_lora_weights()
-        if motion != "":
-            pipe.load_lora_weights(motion, adapter_name="motion")
+        if pinp["motion"] != "":
+            pipe.load_lora_weights(pinp["motion"], adapter_name="motion")
             pipe.fuse_lora()
             pipe.set_adapters(["motion"], [0.7])
-        last_motion = motion
+        last_motion = pinp["motion"]
 
     pipe.to(device,dtype)
 
-    if negative=="":
-        return pipe(
-            prompt=positive,
+    if pinp["negative"]=="":
+        out[x] = pipe(
+            prompt=pinp["positive"],
             height=height,
             width=width,
-            ip_adapter_image=image.convert("RGB").resize((width,height)),
+            ip_adapter_image=pinp["image"].convert("RGB").resize((width,height)),
             num_inference_steps=step,
             guidance_scale=accu,
             num_frames=(fps*time)
         )
     
-    return pipe(
-        prompt=positive,
-        negative_prompt=negative,
+    out[x] = pipe(
+        prompt=pinp["positive"],
+        negative_prompt=pinp["negative"],
         height=height,
         width=width,
-        ip_adapter_image=image.convert("RGB").resize((width,height)),
+        ip_adapter_image=pinp["image"].convert("RGB").resize((width,height)),
         num_inference_steps=step,
         guidance_scale=accu,
         num_frames=(fps*time)
     )
 
-@gpu()
-def infer(pm):
-    x = grid(1)
-
-    pm = pm[x]
-
-    print("infer: started")
+@gpu(void(),device=True,inline=True)
+def infer():
+    global pinp    
+    global out
+    
+    out =[]
     
     p1 = pm["p"]
-    name = generate_random_string[1,32](12)+".png"
     
     neg = pm["n"]
     if neg != "":
@@ -231,34 +231,37 @@ def infer(pm):
 
     if pm["i"] == None:
         return None
-    out = Piper[1,32](pm["i"],posi,neg,pm["m"])
-    export_to_gif(out.frames[0],name,fps=fps)
-    return name
 
-def handle(i,m,p1,p2,result):
+    pinp={"image":pm["i"],"positive":posi,"negative":neg,"motion":pm["m"]}
 
+    ln = len(result)
+    Piper[ln,32]()
+    for i in range(ln):
+        name = generate_random_string[1,32](12)+".png"
+        export_to_gif(out[i].frames[0],name,fps=fps)
+        out[i] = name
+
+@cpu(string[:](),cache=True,parallel=True)
+def handle():
+    global pm
     p1_en = translate[1,32](p1,"english")
     p2_en = translate[1,32](p2,"english")
     pm = {"p":p1_en,"n":p2_en,"m":m,"i":i}
-    ln = len(result)
-    rng = list(range(ln))
-    arr = [pm for _ in rng]
-    #with Pool(f'{ ln }:ppn=2', queue='productionQ', timelimit='5:00:00', workdir='.') as pool:
-        #return pool.map(infer,arr)
-    ret = infer[ln,32](arr)
-    return ret
+    infer[1,32]()
+    return out
 
-@gpu()
+@gpu(void(),device=True,inline=True)
 def ui():
-
     with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
         with gr.Column(elem_id="col-container"):
             gr.Markdown(f"""
                 # MULTI-LANGUAGE IMAGE GENERATOR
             """)
             with gr.Row():
+                global img
                 img = gr.Image(label="STATIC PHOTO",show_label=True,container=True,type="pil")
             with gr.Row():
+                global prompt
                 prompt = gr.Textbox(
                     elem_id="prompt",
                     placeholder="INCLUDE",
@@ -266,6 +269,7 @@ def ui():
                     max_lines=1
                 )
             with gr.Row():
+                global prompt2
                 prompt2 = gr.Textbox(
                     elem_id="prompt2",
                     placeholder="EXCLUDE",
@@ -273,6 +277,7 @@ def ui():
                     max_lines=1
                 )
             with gr.Row():
+                    global motion
                     motion = gr.Dropdown(
                         label='CAMERA',
                         show_label=True,
@@ -292,20 +297,18 @@ def ui():
                         interactive=True
                     )
             with gr.Row():
-                    run_button = gr.Button("START",elem_classes="btn",scale=0)
+                global run_button
+                run_button = gr.Button("START",elem_classes="btn",scale=0)
             with gr.Row():
-                    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))
-                
-        gr.on(
-            triggers=[run_button.click, prompt.submit, prompt2.submit],
-            fn=handle,inputs=[img,motion,prompt,prompt2,result],outputs=result
-        )
+                global result
+                result = []
+                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))
         demo.queue().launch()
 
-@gpu()
+@gpu(void(),device=True,inline=True)
 def pre():
-
+    global pipe
     pipe = AnimateDiffPipeline.from_pretrained(base, vae=vae, motion_adapter=adapter, torch_dtype=dtype).to(device)
     pipe.scheduler = DDIMScheduler(
         clip_sample=False,
@@ -320,15 +323,28 @@ def pre():
     pipe.enable_vae_slicing()
     pipe.enable_free_init(method="butterworth", use_fast_sampling=fast)
 
+@gpu(void(),device=True,inline=True)
+def events():
+    gr.on(
+        triggers=[
+            run_button.click,
+            prompt.submit,
+            prompt2.submit
+        ],
+        fn=handle
+        output=result
+    )
+
 @cpu(void(),cache=True,parallel=True)
 def entry():
     os.chdir(os.path.abspath(os.path.dirname(__file__)))
     mp.set_start_method("spawn", force=True)
     pre[1,32]()
     ui[1,32]()
+    events[1,32]()
 
 # entry
 
-entry()
+entry[1,32]()
 
 # end