Update app.py

app.py

@@ -77,13 +77,13 @@ def process_image(
         processing_res=processing_res,
         batch_size=1 if processing_res == 0 else 0,
         show_progress_bar=False,
+        mode=mode,
     )
 
     depth_pred = pipe_out.pred_np
     depth_colored = pipe_out.pred_colored
 
     np.save(path_out_fp32, depth_pred)
-    
     depth_colored.save(path_out_vis)
     
     if mode == 'depth':
@@ -98,10 +98,10 @@ def process_image(
         [path_out_16bit, path_out_fp32, path_out_vis],
     )
 
-def run_demo_server(pipe):
-    process_pipe_depth = spaces.GPU(functools.partial(process_image, pipe, mode='depth'))
-    process_pipe_normal = spaces.GPU(functools.partial(process_image, pipe, mode='normal'))
-    process_pipe_dis = spaces.GPU(functools.partial(process_image, pipe, mode='dis'))
+def run_demo_server(pipe_depth, pipe_normal, pipe_dis):
+    process_pipe_depth = spaces.GPU(functools.partial(process_image, pipe_depth, mode='depth'))
+    process_pipe_normal = spaces.GPU(functools.partial(process_image, pipe_normal, mode='normal'))
+    process_pipe_dis = spaces.GPU(functools.partial(process_image, pipe_dis, mode='dis'))
     gradio_theme = gr.themes.Default()
 
     with gr.Blocks(
@@ -336,7 +336,7 @@ def run_demo_server(pipe):
             preprocess=False,
             queue=False,
         ).success(
-            fn=process_pipe_normal,
+            fn=process_pipe_image,
             inputs=[
                 image_input,
                 image_processing_res,
@@ -383,17 +383,28 @@ def main():
     unet_depth_v1 = UNet2DConditionModel.from_pretrained('guangkaixu/GenPercept', subfolder="unet_depth_v1").to(dtype)
     empty_text_embed = torch.from_numpy(np.load("./empty_text_embed.npy")).to(device, dtype)[None] # [1, 77, 1024]
     
-    pipe = GenPerceptPipeline(vae=vae,
-                              unet=unet_depth_v1,
-                              empty_text_embed=empty_text_embed)
+    pipe_depth = GenPerceptPipeline(vae=vae,
+                                    unet=unet_depth_v1,
+                                    empty_text_embed=empty_text_embed)
+    pipe_normal = GenPerceptPipeline(vae=vae,
+                                     unet=unet_normal_v1,
+                                     empty_text_embed=empty_text_embed)
+    pipe_dis = GenPerceptPipeline(vae=vae,
+                                  unet=unet_dis_v1,
+                                  empty_text_embed=empty_text_embed)
     try:
         import xformers
-        pipe.enable_xformers_memory_efficient_attention()
+        pipe_depth.enable_xformers_memory_efficient_attention()
+        pipe_normal.enable_xformers_memory_efficient_attention()
+        pipe_dis.enable_xformers_memory_efficient_attention()
     except:
         pass  # run without xformers
 
-    pipe = pipe.to(device)
-    run_demo_server(pipe)
+    pipe_depth = pipe_depth.to(device)
+    pipe_normal = pipe_normal.to(device)
+    pipe_dis = pipe_dis.to(device)
+    
+    run_demo_server(pipe_depth, pipe_normal, pipe_dis)
 
 
 if __name__ == "__main__":