From 729752620b99c6231759723376d88c3fb262fca0 Mon Sep 17 00:00:00 2001 From: damian Date: Thu, 26 Jan 2023 17:27:33 +0100 Subject: [PATCH] trying out JPPhoto's patch on vast.ai --- ldm/invoke/generator/diffusers_pipeline.py | 13 ++++++++++--- ldm/models/diffusion/shared_invokeai_diffusion.py | 5 ++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/ldm/invoke/generator/diffusers_pipeline.py b/ldm/invoke/generator/diffusers_pipeline.py index b7ad925c8c..971a6b4604 100644 --- a/ldm/invoke/generator/diffusers_pipeline.py +++ b/ldm/invoke/generator/diffusers_pipeline.py @@ -304,6 +304,13 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): textual_inversion_manager=self.textual_inversion_manager ) + self._enable_memory_efficient_attention() + + + def _enable_memory_efficient_attention(self): + """ + if xformers is available, use it, otherwise use sliced attention. + """ if is_xformers_available() and not Globals.disable_xformers: self.enable_xformers_memory_efficient_attention() else: @@ -315,7 +322,6 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): else: self.enable_attention_slicing(slice_size='auto') - def image_from_embeddings(self, latents: torch.Tensor, num_inference_steps: int, conditioning_data: ConditioningData, *, @@ -360,6 +366,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): self.scheduler.set_timesteps(num_inference_steps, device=self.unet.device) timesteps = self.scheduler.timesteps infer_latents_from_embeddings = GeneratorToCallbackinator(self.generate_latents_from_embeddings, PipelineIntermediateState) + self._enable_memory_efficient_attention() result: PipelineIntermediateState = infer_latents_from_embeddings( latents, timesteps, conditioning_data, noise=noise, @@ -380,8 +387,8 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): additional_guidance = [] extra_conditioning_info = conditioning_data.extra with self.invokeai_diffuser.custom_attention_context(extra_conditioning_info=extra_conditioning_info, - step_count=len(self.scheduler.timesteps), - do_attention_map_saving=False): + step_count=len(self.scheduler.timesteps) + ): yield PipelineIntermediateState(run_id=run_id, step=-1, timestep=self.scheduler.num_train_timesteps, latents=latents) diff --git a/ldm/models/diffusion/shared_invokeai_diffusion.py b/ldm/models/diffusion/shared_invokeai_diffusion.py index 0c91df9528..10ad328575 100644 --- a/ldm/models/diffusion/shared_invokeai_diffusion.py +++ b/ldm/models/diffusion/shared_invokeai_diffusion.py @@ -60,13 +60,12 @@ class InvokeAIDiffuserComponent: @contextmanager def custom_attention_context(self, extra_conditioning_info: Optional[ExtraConditioningInfo], - step_count: int, - do_attention_map_saving: bool): + step_count: int): do_swap = extra_conditioning_info is not None and extra_conditioning_info.wants_cross_attention_control old_attn_processor = None if do_swap: old_attn_processor = self.setup_cross_attention_control(extra_conditioning_info, - step_count=step_count) + step_count=step_count) try: yield None finally: