test compile pipe

app_onediff.py

@@ -60,6 +60,7 @@ base_pipe = AutoPipelineForText2Image.from_pretrained(
 base_pipe.to(device)
 
 # base_pipe = nexfort_compile(base_pipe)
+compiled_pipe = None
 
 from gfpgan.utils import GFPGANer
 from basicsr.archs.srvgg_arch import SRVGGNetCompact
@@ -79,23 +80,6 @@ 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,
@@ -104,7 +88,9 @@ def create_demo() -> gr.Blocks:
         run_task_time = 0
         time_cost_str = ''
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        generated_image = base_pipe(
+        if compiled_pipe is None:
+            compiled_pipe = nexfort_compile(base_pipe)
+        generated_image = compiled_pipe(
             prompt=prompt,
             num_inference_steps=steps,
         ).images[0]
@@ -139,7 +125,7 @@ def create_demo() -> gr.Blocks:
         g_btn.click(
             fn=compile_pipe,
             inputs=[prompt, steps],
-            outputs=[time_cost],
+            outputs=[generated_image, time_cost],
         )
 
     return demo