From 58fefbeb6541ce2ca090728945dabc739fe8668d Mon Sep 17 00:00:00 2001
From: aifartist <hexexpert5@gmail.com>
Date: Mon, 18 Mar 2024 16:34:50 -0700
Subject: [PATCH] Partial support for onediff

---
 modules/sd_models_compile.py | 37 +++++++++++++++++++++++++++++++++++-
 modules/shared.py            |  2 +-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/modules/sd_models_compile.py b/modules/sd_models_compile.py
index 76bbdd970..50d472030 100644
--- a/modules/sd_models_compile.py
+++ b/modules/sd_models_compile.py
@@ -136,6 +136,39 @@ def optimize_openvino(sd_model):
     return sd_model
 
 
+def compile_onediff(sd_model):
+    try:
+        from onediff.infer_compiler import oneflow_compile
+
+    except Exception as e:
+        shared.log.warning(f"Model compile using onediff/oneflow: {e}")
+        return sd_model
+
+    try:
+        t0 = time.time()
+        # For some reason compiling the text_encoder, when it is used by
+        # the 'compel' package which sdnext uses, it becomes 100 times
+        # slower as if it is recompiling every time.
+        #sd_model.text_encoder = oneflow_compile(sd_model.text_encoder)
+        #if hasattr(sd_model, 'text_endcoder_2'):
+        #    sd_model.text_encoder_2 = oneflow_compile(sd_model.text_encoder_2)
+        sd_model.unet = oneflow_compile(sd_model.unet)
+        sd_model.vae.encoder = oneflow_compile(sd_model.vae.encoder)
+        sd_model.vae.decoder = oneflow_compile(sd_model.vae.decoder)
+        # How are Loras, Adaptors, and other things compiled
+
+        # DW: I'm unclear whether this is also a problem with onediff
+        # as it was for sfast.
+        setup_logging() # compile messes with logging so reset is needed
+        if shared.opts.cuda_compile_precompile:
+            sd_model("dummy prompt")
+        t1 = time.time()
+        shared.log.info(f"Model compile: task=onediff/oneflow time={t1-t0:.2f}")
+    except Exception as e:
+        shared.log.info(f"Model compile: task=onediff/oneflow error: {e}")
+    return sd_model
+
+
 def compile_stablefast(sd_model):
     try:
         import sfast.compilers.stable_diffusion_pipeline_compiler as sf
@@ -277,7 +310,9 @@ def compile_diffusers(sd_model):
         shared.log.warning('Model compile enabled but no backend specified')
         return sd_model
     shared.log.info(f"Model compile: pipeline={sd_model.__class__.__name__} mode={shared.opts.cuda_compile_mode} backend={shared.opts.cuda_compile_backend} fullgraph={shared.opts.cuda_compile_fullgraph} compile={shared.opts.cuda_compile}")
-    if shared.opts.cuda_compile_backend == 'stable-fast':
+    if shared.opts.cuda_compile_backend == 'onediff':
+        sd_model = compile_onediff(sd_model)
+    elif shared.opts.cuda_compile_backend == 'stable-fast':
         sd_model = compile_stablefast(sd_model)
     elif shared.opts.cuda_compile_backend == 'deep-cache':
         sd_model = compile_deepcache(sd_model)
diff --git a/modules/shared.py b/modules/shared.py
index fd85dcc38..0e82e220a 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -376,7 +376,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
 
     "cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
     "cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"]}),
-    "cuda_compile_backend": OptionInfo("none" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx', 'stable-fast', 'deep-cache', 'olive-ai']}),
+    "cuda_compile_backend": OptionInfo("none" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx', 'onediff', 'stable-fast', 'deep-cache', 'olive-ai']}),
     "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}),
     "cuda_compile_fullgraph": OptionInfo(True if not cmd_opts.use_openvino else False, "Model compile fullgraph"),
     "cuda_compile_precompile": OptionInfo(False, "Model compile precompile"),