@@ -4,7 +4,7 @@ The `train_text_to_image.py` script shows how to fine-tune stable diffusion mode
...
@@ -4,7 +4,7 @@ The `train_text_to_image.py` script shows how to fine-tune stable diffusion mode
___Note___:
___Note___:
___This script is experimental. The script fine-tunes the whole model and often times the model overfits and runs into issues like catastrophic forgetting. It's recommended to try different hyperparamters to get the best result on your dataset.___
___This script is experimental. The script fine-tunes the whole model and often times the model overfits and runs into issues like catastrophic forgetting. It's recommended to try different hyperparameters to get the best result on your dataset.___
The `train_text_to_image_sdxl.py` script shows how to fine-tune Stable Diffusion XL (SDXL) on your own dataset.
The `train_text_to_image_sdxl.py` script shows how to fine-tune Stable Diffusion XL (SDXL) on your own dataset.
🚨 This script is experimental. The script fine-tunes the whole model and often times the model overfits and runs into issues like catastrophic forgetting. It's recommended to try different hyperparamters to get the best result on your dataset. 🚨
🚨 This script is experimental. The script fine-tunes the whole model and often times the model overfits and runs into issues like catastrophic forgetting. It's recommended to try different hyperparameters to get the best result on your dataset. 🚨
# distributed under the License is distributed on an "AS IS" BASIS,
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# See the License for the specific language governing permissions and
# limitations under the License.
importargparse
importargparse
importlogging
importlogging
...
@@ -395,7 +396,7 @@ def parse_args():
...
@@ -395,7 +396,7 @@ def parse_args():
"--prediction_type",
"--prediction_type",
type=str,
type=str,
default=None,
default=None,
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.",
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediction_type` is chosen.",
)
)
parser.add_argument(
parser.add_argument(
"--hub_model_id",
"--hub_model_id",
...
@@ -635,7 +636,7 @@ def main():
...
@@ -635,7 +636,7 @@ def main():
ema_unet.to(accelerator.device)
ema_unet.to(accelerator.device)
delload_model
delload_model
foriinrange(len(models)):
for_inrange(len(models)):
# pop models so that they are not loaded again
# pop models so that they are not loaded again
model=models.pop()
model=models.pop()
...
@@ -810,7 +811,7 @@ def main():
...
@@ -810,7 +811,7 @@ def main():
ifargs.use_ema:
ifargs.use_ema:
ema_unet.to(accelerator.device)
ema_unet.to(accelerator.device)
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
# For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
# as these weights are only used for inference, keeping weights in full precision is not required.
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
#
...
@@ -293,7 +294,7 @@ def parse_args():
...
@@ -293,7 +294,7 @@ def parse_args():
"--prediction_type",
"--prediction_type",
type=str,
type=str,
default=None,
default=None,
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.",
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediction_type` is chosen.",
)
)
parser.add_argument(
parser.add_argument(
"--hub_model_id",
"--hub_model_id",
...
@@ -454,7 +455,7 @@ def main():
...
@@ -454,7 +455,7 @@ def main():
vae.requires_grad_(False)
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
text_encoder.requires_grad_(False)
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
# For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
# as these weights are only used for inference, keeping weights in full precision is not required.
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.",
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediction_type` is chosen.",
)
)
parser.add_argument(
parser.add_argument(
"--hub_model_id",
"--hub_model_id",
...
@@ -585,7 +585,7 @@ def main(args):
...
@@ -585,7 +585,7 @@ def main(args):
text_encoder_two.requires_grad_(False)
text_encoder_two.requires_grad_(False)
unet.requires_grad_(False)
unet.requires_grad_(False)
# For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
# For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
# as these weights are only used for inference, keeping weights in full precision is not required.
weight_dtype=torch.float32
weight_dtype=torch.float32
ifaccelerator.mixed_precision=="fp16":
ifaccelerator.mixed_precision=="fp16":
...
@@ -648,7 +648,7 @@ def main(args):
...
@@ -648,7 +648,7 @@ def main(args):
defsave_model_hook(models,weights,output_dir):
defsave_model_hook(models,weights,output_dir):
ifaccelerator.is_main_process:
ifaccelerator.is_main_process:
# there are only two options here. Either are just the unet attn processor layers
# there are only two options here. Either are just the unet attn processor layers
# or there are the unet and text encoder atten layers
# or there are the unet and text encoder attn layers
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.",
help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediction_type` is chosen.",
)
)
parser.add_argument(
parser.add_argument(
"--hub_model_id",
"--hub_model_id",
...
@@ -683,7 +683,7 @@ def main(args):
...
@@ -683,7 +683,7 @@ def main(args):
# Set unet as trainable.
# Set unet as trainable.
unet.train()
unet.train()
# For mixed precision training we cast all non-trainable weigths to half-precision
# For mixed precision training we cast all non-trainable weights to half-precision
# as these weights are only used for inference, keeping weights in full precision is not required.
# as these weights are only used for inference, keeping weights in full precision is not required.