Update app.py

app.py

@@ -27,16 +27,13 @@ from safetensors.torch import load_file, save_file
 from diffusers import DiffusionPipeline, AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, DDIMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel, AutoencoderKL, UNet3DConditionModel
 #import jax
 #import jax.numpy as jnp
-from numba import njit as cpu1, jit as cpu2, cuda
+from numba import jit as cpu, cuda
 from numba.cuda import jit as gpu
 
 # optimization:
 
 # @gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
+# @cpu(cache=True)
 
 # logging
 
@@ -118,18 +115,19 @@ function custom(){
 
 # functionality
 
-@gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
-def run(*args):
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
-
-    cmd=args[0]
+# @gpu(cache=True)
+@cpu(cache=True,nopython=True,parallel=True)
+# @cpu(cache=True)
+def run(cmd):
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+
+    cmd=cmd[pos]
 
     result = subprocess.run(cmd, shell=True, capture_output=True, env=None)
     if result.returncode != 0:
@@ -139,18 +137,20 @@ def run(*args):
         sys.exit()
     return result
     
-@gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
-def translate(*args):
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
-
-    text,lang=args
+# @gpu(cache=True)
+@cpu(cache=True,nopython=True,parallel=True)
+# @cpu(cache=True)
+def translate(args):
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+
+    text=text[pos]
+    lang=lang[pos]
 
     if text == None or lang == None:
         return ""       
@@ -181,7 +181,7 @@ def translate(*args):
     translated = text
     try:
         src_lang = html.xpath('//*[@class="source-language"]')[0].text_content().lower().strip()
-        trgt_lang = html.xpath('//*[@class="target-language"]')[0].text_content().lower().strip()
+        trgt_lang = html.xpath'//*[@class="target-language"]')[0].text_content().lower().strip()
         src_text = html.xpath('//*[@id="tw-source-text"]/*')[0].text_content().lower().strip()
         trgt_text = html.xpath('//*[@id="tw-target-text"]/*')[0].text_content().lower().strip()
         if trgt_lang == lang:
@@ -192,34 +192,39 @@ def translate(*args):
     print(ret)
     return ret
     
-@gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
-def generate_random_string(*args):
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
-
-    length=args[0]
+# @gpu(cache=True)
+@cpu(cache=True,nopython=True,parallel=True)
+# @cpu(cache=True)
+def generate_random_string(length):
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+
+    length=length[pos]
 
     characters = string.ascii_letters + string.digits
     return ''.join(random.choice(characters) for _ in range(length))
     
 @gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
-def Piper(*args):
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
-
-    image,positive,negative,motion=args
+# @cpu(cache=True,nopython=True,parallel=True)
+# @cpu(cache=True)
+def Piper(image,positive,negative,motion):
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+
+    image=image[pos]
+    positive=positive[pos]
+    negative=negative[pos]
+    motion=motion[pos]
 
     global last_motion
     global ip_loaded
@@ -257,22 +262,23 @@ def Piper(*args):
     )
 
 @gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
-def infer(args):
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
-
-    pm = args[0]
+# @cpu(cache=True,nopython=True,parallel=True)
+# @cpu(cache=True)
+def infer(pm):
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+
+    pm = pm[pos]
 
     print("infer: started")
     
     p1 = pm["p"]
-    name = generate_random_string[32,32](12)+".png"
+    name = generate_random_string([12])+".png"
     
     neg = pm["n"]
     if neg != "":
@@ -285,44 +291,50 @@ def infer(args):
 
     if pm["i"] == None:
         return None
-    out = Piper[32,32](pm["i"],posi,neg,pm["m"])
+    out = Piper[32,32]([pm["i"]],[posi],[neg],[pm["m"]])
     export_to_gif(out.frames[0],name,fps=fps)
     return name
 
-@gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
-def handle(*args):
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
-
-    i,m,p1,p2,*result=args
-
-    p1_en = translate[32,32](p1,"english")
-    p2_en = translate[32,32](p2,"english")
+# @gpu(cache=True)
+@cpu(cache=True,nopython=True,parallel=True)
+# @cpu(cache=True)
+def handle(i,m,p1,p2,result):
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+
+    i=i[pos]
+    m=m[pos]
+    p1=p1[pos]
+    p2=p2[pos]
+    result=result[pos]
+
+    p1_en = translate([p1],["english"])
+    p2_en = translate([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[32+ln,32](pm)
+    ret = infer[32+ln,32](arr)
     return ret
 
-@gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
+# @gpu(cache=True)
+# @cpu(cache=True,nopython=True,parallel=True)
+@cpu(cache=True)
 def ui():
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
 
     with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
         with gr.Column(elem_id="col-container"):
@@ -372,20 +384,21 @@ def ui():
                 
         gr.on(
             triggers=[run_button.click, prompt.submit, prompt2.submit],
-            fn=handle[32,32],inputs=[img,motion,prompt,prompt2,*result],outputs=result
+            fn=handle,inputs=[[img],[motion],[prompt],[prompt2],[result]],outputs=result
         )
         demo.queue().launch()
 
-@gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-# @cpu1(cache=True)
-# @cpu2(cache=True)
+# @gpu(cache=True)
+# @cpu(cache=True,nopython=True,parallel=True)
+@cpu(cache=True)
 def pre():
-    tx = cuda.threadIdx.x
-    bx = cuda.blockIdx.x
-    dx = cuda.blockDim.x
-    pos = tx + bx * dx
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
 
     pipe = AnimateDiffPipeline.from_pretrained(base, vae=vae, motion_adapter=adapter, torch_dtype=dtype).to(device)
     pipe.scheduler = DDIMScheduler(
@@ -402,15 +415,21 @@ def pre():
     pipe.enable_free_init(method="butterworth", use_fast_sampling=fast)
 
 # @gpu(cache=True)
-# @cpu1(cache=True,nopython=True,parallel=True)
-# @cpu2(cache=True,nopython=True,parallel=True)
-@cpu1(cache=True)
-# @cpu2(cache=True)
+# @cpu(cache=True,nopython=True,parallel=True)
+@cpu(cache=True)
 def entry():
+    try:
+        tx = cuda.threadIdx.x
+        bx = cuda.blockIdx.x
+        dx = cuda.blockDim.x
+        pos = tx + bx * dx
+    except:
+        pos = 0
+    
     os.chdir(os.path.abspath(os.path.dirname(__file__)))
     mp.set_start_method("spawn", force=True)
-    pre[32,32]()
-    ui[32,32]()
+    pre()
+    ui()
 
 # entry