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