oneflow

app.py

@@ -6,17 +6,9 @@ from app_init import init_app
 from user_queue import user_data
 from util import get_pipeline_class
 
-
 print("DEVICE:", device)
 print("TORCH_DTYPE:", torch_dtype)
-print("PIPELINE:", args.pipeline)
-print("SAFETY_CHECKER:", args.safety_checker)
-print("TORCH_COMPILE:", args.torch_compile)
-print("SFast:", args.sfast)
-print("USE_TAESD:", args.taesd)
-print("COMPEL:", args.compel)
-print("DEBUG:", args.debug)
-
+args.pretty_print()
 
 app = FastAPI()
 

config.py

@@ -17,9 +17,16 @@ class Args(NamedTuple):
     ssl_certfile: str
     ssl_keyfile: str
     sfast: bool
+    oneflow: bool = False
     compel: bool = False
     debug: bool = False
 
+    def pretty_print(self):
+        print("\n")
+        for field, value in self._asdict().items():
+            print(f"{field}: {value}")
+        print("\n")
+
 
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
@@ -109,6 +116,12 @@ parser.add_argument(
     default=False,
     help="Enable Stable Fast",
 )
+parser.add_argument(
+    "--oneflow",
+    action="store_true",
+    default=False,
+    help="Enable OneFlow",
+)
 parser.set_defaults(taesd=USE_TAESD)
 
 args = Args(**vars(parser.parse_args()))

pipelines/controlnet.py

@@ -175,6 +175,7 @@ class Pipeline:
             ).to(device)
 
         if args.sfast:
+            print("\nRunning sfast compile\n")
             from sfast.compilers.stable_diffusion_pipeline_compiler import (
                 compile,
                 CompilationConfig,
@@ -186,6 +187,15 @@ class Pipeline:
             config.enable_cuda_graph = True
             self.pipe = compile(self.pipe, config=config)
 
+        if args.oneflow:
+            print("\nRunning oneflow compile\n")
+            from onediff.infer_compiler import oneflow_compile
+
+            self.pipe.unet = oneflow_compile(self.pipe.unet)
+            self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder)
+            self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder)
+            self.pipe.controlnet = oneflow_compile(self.pipe.controlnet)
+
         self.canny_torch = SobelOperator(device=device)
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)

pipelines/img2imgSDTurbo.py

@@ -109,6 +109,7 @@ class Pipeline:
             ).to(device)
 
         if args.sfast:
+            print("\nRunning sfast compile\n")
             from sfast.compilers.stable_diffusion_pipeline_compiler import (
                 compile,
                 CompilationConfig,
@@ -120,6 +121,13 @@ class Pipeline:
             config.enable_cuda_graph = True
             self.pipe = compile(self.pipe, config=config)
 
+        if args.oneflow:
+            print("\nRunning oneflow compile\n")
+            from onediff.infer_compiler import oneflow_compile
+
+            self.pipe.vae.encoder = oneflow_compile(self.pipe.vae.encoder)
+            self.pipe.vae.decoder = oneflow_compile(self.pipe.vae.decoder)
+
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
         if device.type != "mps":

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/huggingface/diffusers@2d94c7838e273c40920ffd6d24d724357add7f2d
+git+https://github.com/huggingface/diffusers@22b45304bf85a3c5281753d6b3259ccaf96e5085
 transformers==4.35.2
 --extra-index-url https://download.pytorch.org/whl/cu121;
 torch==2.1.0
@@ -11,4 +11,7 @@ controlnet-aux==0.0.7
 peft==0.6.0
 xformers; sys_platform != 'darwin' or platform_machine != 'arm64'
 markdown2
-stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v0.0.15.post1/stable_fast-0.0.15.post1+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl
+stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v0.0.15.post1/stable_fast-0.0.15.post1+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl
+--extra-index-url https://oneflow-pro.oss-cn-beijing.aliyuncs.com/branch/community/cu121
+oneflow
+git+https://github.com/Oneflow-Inc/onediff.git@main#egg=onediff