storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-20 07:45:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
Aryan
2024-12-21 12:24:33 +01:00
parent 66189c804b
commit 2da8cb5798
2 changed files with 75 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
finetrainers/trainer.py
+16
View File
@@ -224,6 +224,16 @@ class Trainer:
self._set_components(condition_components)
self._move_components_to_device()
# TODO(aryan): refactor later. for now only lora is supported
components_to_disable_grads = [
self.text_encoder,
self.text_encoder_2,
self.text_encoder_3,
]
for component in components_to_disable_grads:
if component is not None:
component.requires_grad_(False)
if self.args.caption_dropout_p > 0 and self.args.caption_dropout_technique == "empty":
logger.warning(
"Caption dropout is not supported with precomputation yet. This will be supported in the future."
@@ -275,6 +285,12 @@ class Trainer:
self._set_components(latent_components)
self._move_components_to_device()
# TODO(aryan): refactor later
components_to_disable_grads = [self.vae]
for component in components_to_disable_grads:
if component is not None:
component.requires_grad_(False)
if self.vae is not None:
if self.args.enable_slicing:
self.vae.enable_slicing()