Unverified Commit e9aa0925 authored by Adrià Arrufat's avatar Adrià Arrufat Committed by GitHub
Browse files

Rename --only_save_embeds to --save_as_full_pipeline (#3206)

* Set --only_save_embeds to False by default

Due to how the option is named, it makes more sense to behave like this.

* Refactor only_save_embeds to save_as_full_pipeline
parent 36f43ea7
...@@ -176,10 +176,9 @@ def parse_args(): ...@@ -176,10 +176,9 @@ def parse_args():
help="Save learned_embeds.bin every X updates steps.", help="Save learned_embeds.bin every X updates steps.",
) )
parser.add_argument( parser.add_argument(
"--only_save_embeds", "--save_as_full_pipeline",
action="store_true", action="store_true",
default=True, help="Save the complete stable diffusion pipeline.",
help="Save only the embeddings for the new concept.",
) )
parser.add_argument( parser.add_argument(
"--num_vectors", "--num_vectors",
...@@ -900,11 +899,11 @@ def main(): ...@@ -900,11 +899,11 @@ def main():
# Create the pipeline using the trained modules and save it. # Create the pipeline using the trained modules and save it.
accelerator.wait_for_everyone() accelerator.wait_for_everyone()
if accelerator.is_main_process: if accelerator.is_main_process:
if args.push_to_hub and args.only_save_embeds: if args.push_to_hub and not args.save_as_full_pipeline:
logger.warn("Enabling full model saving because --push_to_hub=True was specified.") logger.warn("Enabling full model saving because --push_to_hub=True was specified.")
save_full_model = True save_full_model = True
else: else:
save_full_model = not args.only_save_embeds save_full_model = args.save_as_full_pipeline
if save_full_model: if save_full_model:
pipeline = StableDiffusionPipeline.from_pretrained( pipeline = StableDiffusionPipeline.from_pretrained(
args.pretrained_model_name_or_path, args.pretrained_model_name_or_path,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment