Update app.py

app.py

@@ -31,12 +31,12 @@ 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
-from numba.cuda import jit as cuda
+from numba import njit as cpu1, jit as cpu2, cuda
+from numba.cuda import jit as gpu
 
 # optimization:
 
-# @cuda(cache=True)
+# @gpu(cache=True)
 # @cpu1(cache=True,nopython=True,parallel=True)
 # @cpu2(cache=True,nopython=True,parallel=True)
 # @cpu1(cache=True)
@@ -126,12 +126,17 @@ function custom(){
 
 # functionality
 
-@cuda(cache=True)
+@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(cmd, assert_success=False, capture_output=False, env=None, dry_run=False):
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     if dry_run:
         print(f"--> {cmd}")
         result = 1
@@ -145,12 +150,17 @@ def run(cmd, assert_success=False, capture_output=False, env=None, dry_run=False
 
     return result
     
-@cuda(cache=True)
+@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(text,lang):
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     if text == None or lang == None:
         return ""       
     text = re.sub(f'[{string.punctuation}]', '', re.sub('[\s+]', ' ', text)).lower().strip()
@@ -191,22 +201,32 @@ def translate(text,lang):
     print(ret)
     return ret
     
-@cuda(cache=True)
+@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(length):
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     characters = string.ascii_letters + string.digits
     return ''.join(random.choice(characters) for _ in range(length))
     
-@cuda(cache=True)
+@gpu(cache=True)
 # @cpu1(cache=True,nopython=True,parallel=True)
 # @cpu2(cache=True,nopython=True,parallel=True)
 # @cpu1(cache=True)
 # @cpu2(cache=True)
 @spaces.GPU(duration=65)
 def Piper(image,positive,negative,motion):
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     global last_motion
     global ip_loaded
 
@@ -242,38 +262,48 @@ def Piper(image,positive,negative,motion):
         num_frames=(fps*time)
     )
 
-@cuda(cache=True)
+@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(pm):
-        print("infer: started")
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
+    print("infer: started")
     
-        p1 = pm["p"]
-        name = generate_random_string(12)+".png"
+    p1 = pm["p"]
+    name = generate_random_string(12)+".png"
     
-        neg = pm["n"]
-        if neg != "":
-            neg = f"{neg} where in the image"
-
-        _do = ['photographed', 'realistic', 'dynamic poze', 'deep field', 'reasonable', "natural", 'rough', 'best quality', 'focused', "highly detailed"]
-        if p1 != "":
-            _do.append(f"a new {p1} content in the image")
-        posi = ", ".join(_do)
-
-        if pm["i"] == None:
-            return None
-        out = Piper(pm["i"],posi,neg,pm["m"])
-        export_to_gif(out.frames[0],name,fps=fps)
-        return name
-
-@cuda(cache=True)
+    neg = pm["n"]
+    if neg != "":
+        neg = f"{neg} where in the image"
+
+    _do = ['photographed', 'realistic', 'dynamic poze', 'deep field', 'reasonable', "natural", 'rough', 'best quality', 'focused', "highly detailed"]
+    if p1 != "":
+        _do.append(f"a new {p1} content in the image")
+    posi = ", ".join(_do)
+
+    if pm["i"] == None:
+        return None
+    out = Piper(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(i,m,p1,p2,*result):    
+def handle(i,m,p1,p2,*result):
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     p1_en = translate(p1,"english")
     p2_en = translate(p2,"english")
     pm = {"p":p1_en,"n":p2_en,"m":m,"i":i}
@@ -287,12 +317,17 @@ def handle(i,m,p1,p2,*result):
         ret.append(infer,pm)
     return ret
 
-@cuda(cache=True)
+@gpu(cache=True)
 # @cpu1(cache=True,nopython=True,parallel=True)
 # @cpu2(cache=True,nopython=True,parallel=True)
 # @cpu1(cache=True)
 # @cpu2(cache=True)
 def ui():
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     with gr.Blocks(theme=gr.themes.Soft(),css=css,js=js) as demo:
         with gr.Column(elem_id="col-container"):
             gr.Markdown(f"""
@@ -345,12 +380,17 @@ def ui():
         )
         demo.queue().launch()
 
-@cuda(cache=True)
+@gpu(cache=True)
 # @cpu1(cache=True,nopython=True,parallel=True)
 # @cpu2(cache=True,nopython=True,parallel=True)
 # @cpu1(cache=True)
 # @cpu2(cache=True)
 def pre():
+    tx = cuda.threadIdx.x
+    bx = cuda.blockIdx.x
+    dx = cuda.blockDim.x
+    pos = tx + bx * dx
+
     pipe = AnimateDiffPipeline.from_pretrained(base, vae=vae, motion_adapter=adapter, torch_dtype=dtype).to(device)
     pipe.scheduler = DDIMScheduler(
         clip_sample=False,
@@ -368,7 +408,8 @@ def pre():
 # entry
 
 if __name__ == "__main__":
-    pre()
-    ui()
+
+    pre[32,32]()
+    ui[32,32]()
 
 # end