oneflow

Dockerfile

@@ -15,7 +15,7 @@ RUN apt-get update && apt-get install --no-install-recommends -y \
     google-perftools \
     ca-certificates curl gnupg \
     && apt-get clean && rm -rf /var/lib/apt/lists/*
-    
+
 WORKDIR /code
 
 RUN mkdir -p /etc/apt/keyrings 
@@ -36,7 +36,7 @@ ENV HOME=/home/user \
     PYTHONUNBUFFERED=1 \
     SYSTEM=spaces
 
-RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
+RUN pip3 install --no-cache-dir --upgrade --pre -r /code/requirements.txt
 
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app

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()
 

build-run.sh

@@ -12,5 +12,9 @@ cd ../
 if [ -z ${PIPELINE+x} ]; then
     PIPELINE="controlnet"
 fi
+if [ -z ${COMPILE+x} ]; then
+    COMPILE="--sfast"
+fi
 echo -e "\033[1;32m\npipeline: $PIPELINE \033[0m"
-python3 run.py --port 7860 --host 0.0.0.0 --pipeline $PIPELINE --sfast
+echo -e "\033[1;32m\ncompile: $COMPILE \033[0m"
+python3 run.py --port 7860 --host 0.0.0.0 --pipeline $PIPELINE $COMPILE

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/controlnetSDTurbo.py

@@ -182,6 +182,7 @@ class Pipeline:
             ).to(device)
 
         if args.sfast:
+            print("\nRunning sfast compile\n")
             from sfast.compilers.stable_diffusion_pipeline_compiler import (
                 compile,
                 CompilationConfig,
@@ -193,6 +194,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.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)

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,14 @@ 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.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
         if device.type != "mps":

requirements.txt

@@ -1,7 +1,7 @@
-git+https://github.com/huggingface/diffusers@2d94c7838e273c40920ffd6d24d724357add7f2d
+diffusers==0.24.0
 transformers==4.35.2
 --extra-index-url https://download.pytorch.org/whl/cu121;
-torch==2.1.0
+torch
 fastapi==0.104.1
 uvicorn[standard]==0.24.0.post1
 Pillow==10.1.0
@@ -11,4 +11,6 @@ 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
+oneflow @ https://oneflow-pro.oss-cn-beijing.aliyuncs.com/branch/community/cu121/794a56cc787217f46b21f5cbc84f65295664b82c/oneflow-0.9.1%2Bcu121.git.794a56c-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
+git+https://github.com/Oneflow-Inc/onediff.git@main#egg=onediff