Update app.py

app.py

@@ -27,9 +27,9 @@ 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 cuda
+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.types import unicode_type as string
 # logging
 
 warnings.filterwarnings("ignore")
@@ -101,6 +101,7 @@ footer {
     display: flex;
 }
 """
+
 js="""
 function custom(){
     document.querySelector("div#prompt input").setAttribute("maxlength","38")
@@ -110,11 +111,9 @@ function custom(){
 
 # functionality
 
+@cpu(string(string),cache=True,parallel=True)
 def run(cmd):
-    #x = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
-    #y = cuda.threadIdx.y + cuda.blockIdx.y * cuda.blockDim.y
-    #z = cuda.threadIdx.z + cuda.blockIdx.z * cuda.blockDim.z
-
+    
     result = subprocess.run(cmd, shell=True, capture_output=True, env=None)
     if result.returncode != 0:
         logging.error(
@@ -126,11 +125,6 @@ def run(cmd):
 @gpu()
 def translate(text,lang):
 
-    
-    
-    text=text[]
-    lang=lang[pos]
-
     if text == None or lang == None:
         return ""       
     text = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', text)).lower().strip()
@@ -171,35 +165,14 @@ def translate(text,lang):
     print(ret)
     return ret
     
-@cpu(cache=True)
+@gpu()
 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)
+@gpu()
 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
@@ -236,22 +209,16 @@ def Piper(image,positive,negative,motion):
         num_frames=(fps*time)
     )
 
-@gpu(cache=True)
+@gpu()
 def infer(pm):
-    try:
-        tx = cuda.threadIdx.x
-        bx = cuda.blockIdx.x
-        dx = cuda.blockDim.x
-        pos = tx + bx * dx
-    except:
-        pos = 0
+    x = grid(1)
 
-    pm = pm[pos]
+    pm = pm[x]
 
     print("infer: started")
     
     p1 = pm["p"]
-    name = generate_random_string([12])+".png"
+    name = generate_random_string[1,32](12)+".png"
     
     neg = pm["n"]
     if neg != "":
@@ -264,46 +231,25 @@ def infer(pm):
 
     if pm["i"] == None:
         return None
-    out = Piper[32,32]([pm["i"]],[posi],[neg],[pm["m"]])
+    out = Piper[1,32](pm["i"],posi,neg,pm["m"])
     export_to_gif(out.frames[0],name,fps=fps)
     return name
 
-@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"])
+    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[32+ln,32](arr)
+    ret = infer[ln,32](arr)
     return ret
 
-@cpu(cache=True)
+@gpu()
 def ui():
-    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"):
@@ -353,19 +299,12 @@ def ui():
                 
         gr.on(
             triggers=[run_button.click, prompt.submit, prompt2.submit],
-            fn=handle,inputs=[[img],[motion],[prompt],[prompt2],[result]],outputs=result
+            fn=handle,inputs=[img,motion,prompt,prompt2,result],outputs=result
         )
         demo.queue().launch()
 
-@cpu(cache=True)
+@gpu()
 def pre():
-    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(
@@ -381,20 +320,12 @@ def pre():
     pipe.enable_vae_slicing()
     pipe.enable_free_init(method="butterworth", use_fast_sampling=fast)
 
-@cpu(cache=True)
+@cpu(void(),cache=True,parallel=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()
-    ui()
+    pre[1,32]()
+    ui[1,32]()
 
 # entry