kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update train_text_to_image_lora.py (#2464)

Browse Source 
* Update train_text_to_image_lora.py

* Update train_text_to_image_lora.py
This commit is contained in:
Haofan Wang
2023-02-23 13:38:21 +08:00
committed by GitHub
parent 39a3c77e0d
commit 589faa8c88
1 changed files with 14 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/text_to_image/train_text_to_image_lora.py
+14 -13
View File
@@ -448,19 +448,6 @@ def main():
vae.to(accelerator.device, dtype=weight_dtype)
text_encoder.to(accelerator.device, dtype=weight_dtype)
if args.enable_xformers_memory_efficient_attention:
if is_xformers_available():
import xformers
xformers_version = version.parse(xformers.__version__)
if xformers_version == version.parse("0.0.16"):
logger.warn(
"xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
)
unet.enable_xformers_memory_efficient_attention()
else:
raise ValueError("xformers is not available. Make sure it is installed correctly")
# now we will add new LoRA weights to the attention layers
# It's important to realize here how many attention weights will be added and of which sizes
# The sizes of the attention layers consist only of two different variables:
@@ -492,6 +479,20 @@ def main():
)
unet.set_attn_processor(lora_attn_procs)
if args.enable_xformers_memory_efficient_attention:
if is_xformers_available():
import xformers
xformers_version = version.parse(xformers.__version__)
if xformers_version == version.parse("0.0.16"):
logger.warn(
"xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
)
unet.enable_xformers_memory_efficient_attention()
else:
raise ValueError("xformers is not available. Make sure it is installed correctly")
lora_layers = AttnProcsLayers(unet.attn_processors)
# Enable TF32 for faster training on Ampere GPUs,