test compile pipe

app_onediff.py

@@ -3,6 +3,7 @@ import gradio as gr
 import time
 import torch
 import os
+import json
 
 from diffusers import (
     DDPMScheduler,
@@ -11,16 +12,15 @@ from diffusers import (
 )
 
 os.system("python3 -m pip --no-cache-dir install --pre nexfort -f https://github.com/siliconflow/nexfort_releases/releases/expanded_assets/torch2.4.1_cu121")
-os.system("python3 -m pip install --pre onediff")
+os.system("git clone https://github.com/siliconflow/onediff.git")
+os.system("cd onediff && python3 -m pip install -e .")
+os.system("cd onediff_diffusers_extensions && python3 -m pip install -e .")
 
-import nexfort as flow
-from onediff.infer_compiler import (
-    compile as onediff_compile,
-    DeployableModule,
-)
+from onediffx import compile_pipe
 
-def nexfort_compile(torch_module: torch.nn.Module, *, options=None) -> DeployableModule:
-    return onediff_compile(torch_module, backend="nexfort", options=options)
+def nexfort_compile(torch_module: torch.nn.Module):
+    options = json.loads('{"mode": "max-autotune:cudagraphs", "dynamic": true}')
+    return compile_pipe(torch_module, backend="nexfort", options=options, fuse_qkv_projections=True)
 
 BASE_MODEL = "stabilityai/sdxl-turbo"
 device = "cuda"
@@ -45,7 +45,7 @@ base_pipe.scheduler = DDPMScheduler.from_pretrained(
     subfolder="scheduler",
 )
 
-# base_pipe.unet = nexfort_compile(base_pipe.unet)
+base_pipe = compile_pipe(base_pipe)
 
 def create_demo() -> gr.Blocks: