zhiweili commited on
Commit
8f95526
Β·
1 Parent(s): 7378271

use nexfort_compile

Browse files
Files changed (1) hide show
  1. app_onediff.py +5 -2
app_onediff.py CHANGED
@@ -14,7 +14,10 @@ os.system("python3 -m pip --no-cache-dir install --pre nexfort -f https://github
14
  os.system("python3 -m pip install --pre onediff")
15
 
16
  import nexfort as flow
17
- from onediff.infer_compiler import oneflow_compile
 
 
 
18
 
19
  BASE_MODEL = "stabilityai/sdxl-turbo"
20
  device = "cuda"
@@ -39,7 +42,7 @@ base_pipe.scheduler = DDPMScheduler.from_pretrained(
39
  subfolder="scheduler",
40
  )
41
 
42
- base_pipe.unet = oneflow_compile(base_pipe.unet)
43
 
44
  def create_demo() -> gr.Blocks:
45
 
 
14
  os.system("python3 -m pip install --pre onediff")
15
 
16
  import nexfort as flow
17
+ from onediff.compile import onediff_compile
18
+
19
+ def nexfort_compile(torch_module: torch.nn.Module, *, options=None) -> onediff.DeployableModule:
20
+ return onediff_compile(torch_module, backend="nexfort", options=options)
21
 
22
  BASE_MODEL = "stabilityai/sdxl-turbo"
23
  device = "cuda"
 
42
  subfolder="scheduler",
43
  )
44
 
45
+ base_pipe.unet = nexfort_compile(base_pipe.unet)
46
 
47
  def create_demo() -> gr.Blocks:
48