sayakpaul HF staff commited on
Commit
16d8d11
1 Parent(s): 2de4e1c

fix for PR: 8623

Browse files
Files changed (1) hide show
  1. my_pipeline.py +4 -13
my_pipeline.py CHANGED
@@ -25,12 +25,7 @@ from diffusers.loaders import (
25
  TextualInversionLoaderMixin,
26
  )
27
  from diffusers.models import AutoencoderKL, UNet2DConditionModel
28
- from diffusers.models.attention_processor import (
29
- AttnProcessor2_0,
30
- LoRAAttnProcessor2_0,
31
- LoRAXFormersAttnProcessor,
32
- XFormersAttnProcessor,
33
- )
34
  from diffusers.models.lora import adjust_lora_scale_text_encoder
35
  from diffusers.schedulers import KarrasDiffusionSchedulers
36
  from diffusers.utils import (
@@ -135,6 +130,7 @@ class MyPipeline(
135
  watermark output images. If not defined, it will default to True if the package is installed, otherwise no
136
  watermarker will be used.
137
  """
 
138
  model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
139
  _optional_components = ["tokenizer", "tokenizer_2", "text_encoder", "text_encoder_2"]
140
 
@@ -571,12 +567,7 @@ class MyPipeline(
571
  self.vae.to(dtype=torch.float32)
572
  use_torch_2_0_or_xformers = isinstance(
573
  self.vae.decoder.mid_block.attentions[0].processor,
574
- (
575
- AttnProcessor2_0,
576
- XFormersAttnProcessor,
577
- LoRAXFormersAttnProcessor,
578
- LoRAAttnProcessor2_0,
579
- ),
580
  )
581
  # if xformers or torch_2_0 is used attention block does not need
582
  # to be in float32 which can save lots of memory
@@ -971,4 +962,4 @@ class MyPipeline(
971
  # Offload all models
972
  self.maybe_free_model_hooks()
973
 
974
- return (image,)
 
25
  TextualInversionLoaderMixin,
26
  )
27
  from diffusers.models import AutoencoderKL, UNet2DConditionModel
28
+ from diffusers.models.attention_processor import AttnProcessor2_0, XFormersAttnProcessor
 
 
 
 
 
29
  from diffusers.models.lora import adjust_lora_scale_text_encoder
30
  from diffusers.schedulers import KarrasDiffusionSchedulers
31
  from diffusers.utils import (
 
130
  watermark output images. If not defined, it will default to True if the package is installed, otherwise no
131
  watermarker will be used.
132
  """
133
+
134
  model_cpu_offload_seq = "text_encoder->text_encoder_2->unet->vae"
135
  _optional_components = ["tokenizer", "tokenizer_2", "text_encoder", "text_encoder_2"]
136
 
 
567
  self.vae.to(dtype=torch.float32)
568
  use_torch_2_0_or_xformers = isinstance(
569
  self.vae.decoder.mid_block.attentions[0].processor,
570
+ (AttnProcessor2_0, XFormersAttnProcessor),
 
 
 
 
 
571
  )
572
  # if xformers or torch_2_0 is used attention block does not need
573
  # to be in float32 which can save lots of memory
 
962
  # Offload all models
963
  self.maybe_free_model_hooks()
964
 
965
+ return (image,)