File tree Expand file tree Collapse file tree
Expand file tree Collapse file tree Original file line number Diff line number Diff line change @@ -748,12 +748,16 @@ def __call__(
748748 # make sure the VAE is in float32 mode, as it overflows in float16
749749 self .vae .to (dtype = torch .float32 )
750750
751- use_torch_2_0_or_xformers = self .vae .decoder .mid_block .attentions [0 ].processor in [
752- AttnProcessor2_0 ,
753- XFormersAttnProcessor ,
754- LoRAXFormersAttnProcessor ,
755- LoRAAttnProcessor2_0 ,
756- ]
751+ use_torch_2_0_or_xformers = isinstance (
752+ self .vae .decoder .mid_block .attentions [0 ].processor ,
753+ (
754+ AttnProcessor2_0 ,
755+ XFormersAttnProcessor ,
756+ LoRAXFormersAttnProcessor ,
757+ LoRAAttnProcessor2_0 ,
758+ ),
759+ )
760+
757761 # if xformers or torch_2_0 is used attention block does not need
758762 # to be in float32 which can save lots of memory
759763 if not use_torch_2_0_or_xformers :
Original file line number Diff line number Diff line change @@ -786,15 +786,18 @@ def __call__(
786786 # make sure the VAE is in float32 mode, as it overflows in float16
787787 self .vae .to (dtype = torch .float32 )
788788
789- use_torch_2_0_or_xformers = self .vae .decoder .mid_block .attentions [0 ].processor in [
790- AttnProcessor2_0 ,
791- XFormersAttnProcessor ,
792- LoRAXFormersAttnProcessor ,
793- LoRAAttnProcessor2_0 ,
794- ]
789+ use_torch_2_0_or_xformers = isinstance (
790+ self .vae .decoder .mid_block .attentions [0 ].processor ,
791+ (
792+ AttnProcessor2_0 ,
793+ XFormersAttnProcessor ,
794+ LoRAXFormersAttnProcessor ,
795+ LoRAAttnProcessor2_0 ,
796+ ),
797+ )
795798 # if xformers or torch_2_0 is used attention block does not need
796799 # to be in float32 which can save lots of memory
797- if not use_torch_2_0_or_xformers :
800+ if use_torch_2_0_or_xformers :
798801 self .vae .post_quant_conv .to (latents .dtype )
799802 self .vae .decoder .conv_in .to (latents .dtype )
800803 self .vae .decoder .mid_block .to (latents .dtype )
Original file line number Diff line number Diff line change @@ -860,15 +860,18 @@ def __call__(
860860 # make sure the VAE is in float32 mode, as it overflows in float16
861861 self .vae .to (dtype = torch .float32 )
862862
863- use_torch_2_0_or_xformers = self .vae .decoder .mid_block .attentions [0 ].processor in [
864- AttnProcessor2_0 ,
865- XFormersAttnProcessor ,
866- LoRAXFormersAttnProcessor ,
867- LoRAAttnProcessor2_0 ,
868- ]
863+ use_torch_2_0_or_xformers = isinstance (
864+ self .vae .decoder .mid_block .attentions [0 ].processor ,
865+ (
866+ AttnProcessor2_0 ,
867+ XFormersAttnProcessor ,
868+ LoRAXFormersAttnProcessor ,
869+ LoRAAttnProcessor2_0 ,
870+ ),
871+ )
869872 # if xformers or torch_2_0 is used attention block does not need
870873 # to be in float32 which can save lots of memory
871- if not use_torch_2_0_or_xformers :
874+ if use_torch_2_0_or_xformers :
872875 self .vae .post_quant_conv .to (latents .dtype )
873876 self .vae .decoder .conv_in .to (latents .dtype )
874877 self .vae .decoder .mid_block .to (latents .dtype )
You can’t perform that action at this time.
0 commit comments