From 4b235b874e2403dacbf20d0ace33d24b7a6747c5 Mon Sep 17 00:00:00 2001 From: DGX Date: Wed, 10 Jan 2024 08:58:30 +0000 Subject: [PATCH] compile flag with sfast --- latentblending/blending_engine.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/latentblending/blending_engine.py b/latentblending/blending_engine.py index b06277d..e3373a5 100644 --- a/latentblending/blending_engine.py +++ b/latentblending/blending_engine.py @@ -19,6 +19,7 @@ class BlendingEngine(): def __init__( self, dh: None, + do_compile: bool = False, guidance_scale_mid_damper: float = 0.5, mid_compression_scaler: float = 1.2): r""" @@ -80,6 +81,16 @@ class BlendingEngine(): self.benchmark_speed() self.set_branching() + if do_compile: + print("starting compilation") + from sfast.compilers.diffusion_pipeline_compiler import (compile, CompilationConfig) + self.dh.pipe.enable_xformers_memory_efficient_attention() + config = CompilationConfig.Default() + config.enable_xformers = True + config.enable_triton = True + config.enable_cuda_graph = True + self.dh.pipe = compile(self.dh.pipe, config) + def benchmark_speed(self): @@ -833,19 +844,20 @@ if __name__ == "__main__": duration_transition = 12 # In seconds # Spawn latent blending - lb = LatentBlending(dh) - lb.set_prompt1(prompt1) - lb.set_prompt2(prompt2) - lb.set_negative_prompt(negative_prompt) + be = BlendingEngine(dh) + be.set_prompt1(prompt1) + be.set_prompt2(prompt2) + be.set_negative_prompt(negative_prompt) # Run latent blending t0 = time.time() - lb.run_transition(fixed_seeds=[420, 421]) + be.run_transition(fixed_seeds=[420, 421]) dt = time.time() - t0 + print(f"dt = {dt}") # Save movie fp_movie = f'test.mp4' - lb.write_movie_transition(fp_movie, duration_transition) + be.write_movie_transition(fp_movie, duration_transition)