From e6bc9eebff1a9c67597d2bf656de6f76c8c08e91 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Sun, 23 Jun 2024 11:30:26 +0800
Subject: [PATCH] Create __init__.py

add
---
 streamv2v/acceleration/sfast/__init__.py | 33 ++++++++++++++++++++++++
 1 file changed, 33 insertions(+)
 create mode 100644 streamv2v/acceleration/sfast/__init__.py

diff --git a/streamv2v/acceleration/sfast/__init__.py b/streamv2v/acceleration/sfast/__init__.py
new file mode 100644
index 0000000..cffb81e
--- /dev/null
+++ b/streamv2v/acceleration/sfast/__init__.py
@@ -0,0 +1,33 @@
+from typing import Optional
+
+from sfast.compilers.stable_diffusion_pipeline_compiler import CompilationConfig, compile
+
+from ...pipeline import StreamV2V
+
+
+def accelerate_with_stable_fast(
+    stream: StreamV2V,
+    config: Optional[CompilationConfig] = None,
+):
+    if config is None:
+        config = CompilationConfig.Default()
+        # xformers and Triton are suggested for achieving best performance.
+        try:
+            import xformers
+
+            config.enable_xformers = True
+        except ImportError:
+            print("xformers not installed, skip")
+        try:
+            import triton
+
+            config.enable_triton = True
+        except ImportError:
+            print("Triton not installed, skip")
+        # CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.
+        config.enable_cuda_graph = True
+    stream.pipe = compile(stream.pipe, config)
+    stream.unet = stream.pipe.unet
+    stream.vae = stream.pipe.vae
+    stream.text_encoder = stream.pipe.text_encoder
+    return stream