Spaces:
Running
on
Zero
Running
on
Zero
| # coding=utf-8 | |
| # Copyright 2024 HuggingFace Inc. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import logging | |
| import os | |
| import sys | |
| import tempfile | |
| import safetensors | |
| sys.path.append("..") | |
| from test_examples_utils import ExamplesTestsAccelerate, run_command # noqa: E402 | |
| logging.basicConfig(level=logging.DEBUG) | |
| logger = logging.getLogger() | |
| stream_handler = logging.StreamHandler(sys.stdout) | |
| logger.addHandler(stream_handler) | |
| class DreamBoothLoRASDXLWithEDM(ExamplesTestsAccelerate): | |
| def test_dreambooth_lora_sdxl_with_edm(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| test_args = f""" | |
| examples/dreambooth/train_dreambooth_lora_sdxl.py | |
| --pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-xl-pipe | |
| --do_edm_style_training | |
| --instance_data_dir docs/source/en/imgs | |
| --instance_prompt photo | |
| --resolution 64 | |
| --train_batch_size 1 | |
| --gradient_accumulation_steps 1 | |
| --max_train_steps 2 | |
| --learning_rate 5.0e-04 | |
| --scale_lr | |
| --lr_scheduler constant | |
| --lr_warmup_steps 0 | |
| --output_dir {tmpdir} | |
| """.split() | |
| run_command(self._launch_args + test_args) | |
| # save_pretrained smoke test | |
| self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))) | |
| # make sure the state_dict has the correct naming in the parameters. | |
| lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) | |
| is_lora = all("lora" in k for k in lora_state_dict.keys()) | |
| self.assertTrue(is_lora) | |
| # when not training the text encoder, all the parameters in the state dict should start | |
| # with `"unet"` in their names. | |
| starts_with_unet = all(key.startswith("unet") for key in lora_state_dict.keys()) | |
| self.assertTrue(starts_with_unet) | |
| def test_dreambooth_lora_playground(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| test_args = f""" | |
| examples/dreambooth/train_dreambooth_lora_sdxl.py | |
| --pretrained_model_name_or_path hf-internal-testing/tiny-playground-v2-5-pipe | |
| --instance_data_dir docs/source/en/imgs | |
| --instance_prompt photo | |
| --resolution 64 | |
| --train_batch_size 1 | |
| --gradient_accumulation_steps 1 | |
| --max_train_steps 2 | |
| --learning_rate 5.0e-04 | |
| --scale_lr | |
| --lr_scheduler constant | |
| --lr_warmup_steps 0 | |
| --output_dir {tmpdir} | |
| """.split() | |
| run_command(self._launch_args + test_args) | |
| # save_pretrained smoke test | |
| self.assertTrue(os.path.isfile(os.path.join(tmpdir, "pytorch_lora_weights.safetensors"))) | |
| # make sure the state_dict has the correct naming in the parameters. | |
| lora_state_dict = safetensors.torch.load_file(os.path.join(tmpdir, "pytorch_lora_weights.safetensors")) | |
| is_lora = all("lora" in k for k in lora_state_dict.keys()) | |
| self.assertTrue(is_lora) | |
| # when not training the text encoder, all the parameters in the state dict should start | |
| # with `"unet"` in their names. | |
| starts_with_unet = all(key.startswith("unet") for key in lora_state_dict.keys()) | |
| self.assertTrue(starts_with_unet) | |