zhiweili commited on
Commit
c555988
Β·
1 Parent(s): 118e1e2

test compile pipe

Browse files
Files changed (1) hide show
  1. app_onediff.py +9 -9
app_onediff.py CHANGED
@@ -3,6 +3,7 @@ import gradio as gr
3
  import time
4
  import torch
5
  import os
 
6
 
7
  from diffusers import (
8
  DDPMScheduler,
@@ -11,16 +12,15 @@ from diffusers import (
11
  )
12
 
13
  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")
14
- os.system("python3 -m pip install --pre onediff")
 
 
15
 
16
- import nexfort as flow
17
- from onediff.infer_compiler import (
18
- compile as onediff_compile,
19
- DeployableModule,
20
- )
21
 
22
- def nexfort_compile(torch_module: torch.nn.Module, *, options=None) -> DeployableModule:
23
- return onediff_compile(torch_module, backend="nexfort", options=options)
 
24
 
25
  BASE_MODEL = "stabilityai/sdxl-turbo"
26
  device = "cuda"
@@ -45,7 +45,7 @@ base_pipe.scheduler = DDPMScheduler.from_pretrained(
45
  subfolder="scheduler",
46
  )
47
 
48
- # base_pipe.unet = nexfort_compile(base_pipe.unet)
49
 
50
  def create_demo() -> gr.Blocks:
51
 
 
3
  import time
4
  import torch
5
  import os
6
+ import json
7
 
8
  from diffusers import (
9
  DDPMScheduler,
 
12
  )
13
 
14
  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")
15
+ os.system("git clone https://github.com/siliconflow/onediff.git")
16
+ os.system("cd onediff && python3 -m pip install -e .")
17
+ os.system("cd onediff_diffusers_extensions && python3 -m pip install -e .")
18
 
19
+ from onediffx import compile_pipe
 
 
 
 
20
 
21
+ def nexfort_compile(torch_module: torch.nn.Module):
22
+ options = json.loads('{"mode": "max-autotune:cudagraphs", "dynamic": true}')
23
+ return compile_pipe(torch_module, backend="nexfort", options=options, fuse_qkv_projections=True)
24
 
25
  BASE_MODEL = "stabilityai/sdxl-turbo"
26
  device = "cuda"
 
45
  subfolder="scheduler",
46
  )
47
 
48
+ base_pipe = compile_pipe(base_pipe)
49
 
50
  def create_demo() -> gr.Blocks:
51