test compile pipe

app_onediff.py

@@ -37,7 +37,7 @@ os.system("pip show nexfort")
 os.system("pip show onediff")
 os.system("pip show onediffx")
 
-from onediffx import compile_pipe
+from onediffx import compile_pipe, save_pipe
 
 def nexfort_compile(torch_module: torch.nn.Module):
     options = json.loads('{"mode": "max-autotune:cudagraphs", "dynamic": true}')
@@ -79,6 +79,23 @@ face_enhancer = GFPGANer(model_path='GFPGANv1.4.pth', upscale=2, arch='clean', c
 
 def create_demo() -> gr.Blocks:
 
+    @spaces.GPU(duration=120)
+    def compile_pipe(
+        prompt:str,
+        steps:int,
+    ):
+        run_task_time = 0
+        time_cost_str = ''
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        compiled_pipe = nexfort_compile(base_pipe)
+        run_task_time = 0
+        time_cost_str = ''
+        run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
+        if not os.path.exists('cache_pipe'):
+            os.makedirs('cached_pipe')
+        save_pipe(compiled_pipe, dir="cached_pipe")
+        return time_cost_str
+
     @spaces.GPU(duration=30)
     def text_to_image(
         prompt:str,
@@ -120,9 +137,9 @@ def create_demo() -> gr.Blocks:
                 time_cost = gr.Textbox(label="Time Cost", lines=1, interactive=False)
         
         g_btn.click(
-            fn=text_to_image,
+            fn=compile_pipe,
             inputs=[prompt, steps],
-            outputs=[generated_image, time_cost],
+            outputs=[time_cost],
         )
 
     return demo