tsi-org/tango

1# coding=utf-82# Copyright 2023 HuggingFace Inc..3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16 17import logging18import os19import shutil20import subprocess21import sys22import tempfile23import unittest24from typing import List25 26from accelerate.utils import write_basic_config27 28from diffusers import DiffusionPipeline, UNet2DConditionModel29 30 31logging.basicConfig(level=logging.DEBUG)32 33logger = logging.getLogger()34 35 36# These utils relate to ensuring the right error message is received when running scripts37class SubprocessCallException(Exception):38    pass39 40 41def run_command(command: List[str], return_stdout=False):42    """43    Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture44    if an error occurred while running `command`45    """46    try:47        output = subprocess.check_output(command, stderr=subprocess.STDOUT)48        if return_stdout:49            if hasattr(output, "decode"):50                output = output.decode("utf-8")51            return output52    except subprocess.CalledProcessError as e:53        raise SubprocessCallException(54            f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}"55        ) from e56 57 58stream_handler = logging.StreamHandler(sys.stdout)59logger.addHandler(stream_handler)60 61 62class ExamplesTestsAccelerate(unittest.TestCase):63    @classmethod64    def setUpClass(cls):65        super().setUpClass()66        cls._tmpdir = tempfile.mkdtemp()67        cls.configPath = os.path.join(cls._tmpdir, "default_config.yml")68 69        write_basic_config(save_location=cls.configPath)70        cls._launch_args = ["accelerate", "launch", "--config_file", cls.configPath]71 72    @classmethod73    def tearDownClass(cls):74        super().tearDownClass()75        shutil.rmtree(cls._tmpdir)76 77    def test_train_unconditional(self):78        with tempfile.TemporaryDirectory() as tmpdir:79            test_args = f"""80                examples/unconditional_image_generation/train_unconditional.py81                --dataset_name hf-internal-testing/dummy_image_class_data82                --model_config_name_or_path diffusers/ddpm_dummy83                --resolution 6484                --output_dir {tmpdir}85                --train_batch_size 286                --num_epochs 187                --gradient_accumulation_steps 188                --ddpm_num_inference_steps 289                --learning_rate 1e-390                --lr_warmup_steps 591                """.split()92 93            run_command(self._launch_args + test_args, return_stdout=True)94            # save_pretrained smoke test95            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "unet", "diffusion_pytorch_model.bin")))96            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "scheduler", "scheduler_config.json")))97 98    def test_textual_inversion(self):99        with tempfile.TemporaryDirectory() as tmpdir:100            test_args = f"""101                examples/textual_inversion/textual_inversion.py102                --pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-pipe103                --train_data_dir docs/source/en/imgs104                --learnable_property object105                --placeholder_token <cat-toy>106                --initializer_token a107                --resolution 64108                --train_batch_size 1109                --gradient_accumulation_steps 1110                --max_train_steps 2111                --learning_rate 5.0e-04112                --scale_lr113                --lr_scheduler constant114                --lr_warmup_steps 0115                --output_dir {tmpdir}116                """.split()117 118            run_command(self._launch_args + test_args)119            # save_pretrained smoke test120            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "learned_embeds.bin")))121 122    def test_dreambooth(self):123        with tempfile.TemporaryDirectory() as tmpdir:124            test_args = f"""125                examples/dreambooth/train_dreambooth.py126                --pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-pipe127                --instance_data_dir docs/source/en/imgs128                --instance_prompt photo129                --resolution 64130                --train_batch_size 1131                --gradient_accumulation_steps 1132                --max_train_steps 2133                --learning_rate 5.0e-04134                --scale_lr135                --lr_scheduler constant136                --lr_warmup_steps 0137                --output_dir {tmpdir}138                """.split()139 140            run_command(self._launch_args + test_args)141            # save_pretrained smoke test142            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "unet", "diffusion_pytorch_model.bin")))143            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "scheduler", "scheduler_config.json")))144 145    def test_dreambooth_checkpointing(self):146        instance_prompt = "photo"147        pretrained_model_name_or_path = "hf-internal-testing/tiny-stable-diffusion-pipe"148 149        with tempfile.TemporaryDirectory() as tmpdir:150            # Run training script with checkpointing151            # max_train_steps == 5, checkpointing_steps == 2152            # Should create checkpoints at steps 2, 4153 154            initial_run_args = f"""155                examples/dreambooth/train_dreambooth.py156                --pretrained_model_name_or_path {pretrained_model_name_or_path}157                --instance_data_dir docs/source/en/imgs158                --instance_prompt {instance_prompt}159                --resolution 64160                --train_batch_size 1161                --gradient_accumulation_steps 1162                --max_train_steps 5163                --learning_rate 5.0e-04164                --scale_lr165                --lr_scheduler constant166                --lr_warmup_steps 0167                --output_dir {tmpdir}168                --checkpointing_steps=2169                --seed=0170                """.split()171 172            run_command(self._launch_args + initial_run_args)173 174            # check can run the original fully trained output pipeline175            pipe = DiffusionPipeline.from_pretrained(tmpdir, safety_checker=None)176            pipe(instance_prompt, num_inference_steps=2)177 178            # check checkpoint directories exist179            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-2")))180            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-4")))181 182            # check can run an intermediate checkpoint183            unet = UNet2DConditionModel.from_pretrained(tmpdir, subfolder="checkpoint-2/unet")184            pipe = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path, unet=unet, safety_checker=None)185            pipe(instance_prompt, num_inference_steps=2)186 187            # Remove checkpoint 2 so that we can check only later checkpoints exist after resuming188            shutil.rmtree(os.path.join(tmpdir, "checkpoint-2"))189 190            # Run training script for 7 total steps resuming from checkpoint 4191 192            resume_run_args = f"""193                examples/dreambooth/train_dreambooth.py194                --pretrained_model_name_or_path {pretrained_model_name_or_path}195                --instance_data_dir docs/source/en/imgs196                --instance_prompt {instance_prompt}197                --resolution 64198                --train_batch_size 1199                --gradient_accumulation_steps 1200                --max_train_steps 7201                --learning_rate 5.0e-04202                --scale_lr203                --lr_scheduler constant204                --lr_warmup_steps 0205                --output_dir {tmpdir}206                --checkpointing_steps=2207                --resume_from_checkpoint=checkpoint-4208                --seed=0209                """.split()210 211            run_command(self._launch_args + resume_run_args)212 213            # check can run new fully trained pipeline214            pipe = DiffusionPipeline.from_pretrained(tmpdir, safety_checker=None)215            pipe(instance_prompt, num_inference_steps=2)216 217            # check old checkpoints do not exist218            self.assertFalse(os.path.isdir(os.path.join(tmpdir, "checkpoint-2")))219 220            # check new checkpoints exist221            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-4")))222            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-6")))223 224    def test_text_to_image(self):225        with tempfile.TemporaryDirectory() as tmpdir:226            test_args = f"""227                examples/text_to_image/train_text_to_image.py228                --pretrained_model_name_or_path hf-internal-testing/tiny-stable-diffusion-pipe229                --dataset_name hf-internal-testing/dummy_image_text_data230                --resolution 64231                --center_crop232                --random_flip233                --train_batch_size 1234                --gradient_accumulation_steps 1235                --max_train_steps 2236                --learning_rate 5.0e-04237                --scale_lr238                --lr_scheduler constant239                --lr_warmup_steps 0240                --output_dir {tmpdir}241                """.split()242 243            run_command(self._launch_args + test_args)244            # save_pretrained smoke test245            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "unet", "diffusion_pytorch_model.bin")))246            self.assertTrue(os.path.isfile(os.path.join(tmpdir, "scheduler", "scheduler_config.json")))247 248    def test_text_to_image_checkpointing(self):249        pretrained_model_name_or_path = "hf-internal-testing/tiny-stable-diffusion-pipe"250        prompt = "a prompt"251 252        with tempfile.TemporaryDirectory() as tmpdir:253            # Run training script with checkpointing254            # max_train_steps == 5, checkpointing_steps == 2255            # Should create checkpoints at steps 2, 4256 257            initial_run_args = f"""258                examples/text_to_image/train_text_to_image.py259                --pretrained_model_name_or_path {pretrained_model_name_or_path}260                --dataset_name hf-internal-testing/dummy_image_text_data261                --resolution 64262                --center_crop263                --random_flip264                --train_batch_size 1265                --gradient_accumulation_steps 1266                --max_train_steps 5267                --learning_rate 5.0e-04268                --scale_lr269                --lr_scheduler constant270                --lr_warmup_steps 0271                --output_dir {tmpdir}272                --checkpointing_steps=2273                --seed=0274                """.split()275 276            run_command(self._launch_args + initial_run_args)277 278            pipe = DiffusionPipeline.from_pretrained(tmpdir, safety_checker=None)279            pipe(prompt, num_inference_steps=2)280 281            # check checkpoint directories exist282            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-2")))283            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-4")))284 285            # check can run an intermediate checkpoint286            unet = UNet2DConditionModel.from_pretrained(tmpdir, subfolder="checkpoint-2/unet")287            pipe = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path, unet=unet, safety_checker=None)288            pipe(prompt, num_inference_steps=2)289 290            # Remove checkpoint 2 so that we can check only later checkpoints exist after resuming291            shutil.rmtree(os.path.join(tmpdir, "checkpoint-2"))292 293            # Run training script for 7 total steps resuming from checkpoint 4294 295            resume_run_args = f"""296                examples/text_to_image/train_text_to_image.py297                --pretrained_model_name_or_path {pretrained_model_name_or_path}298                --dataset_name hf-internal-testing/dummy_image_text_data299                --resolution 64300                --center_crop301                --random_flip302                --train_batch_size 1303                --gradient_accumulation_steps 1304                --max_train_steps 7305                --learning_rate 5.0e-04306                --scale_lr307                --lr_scheduler constant308                --lr_warmup_steps 0309                --output_dir {tmpdir}310                --checkpointing_steps=2311                --resume_from_checkpoint=checkpoint-4312                --seed=0313                """.split()314 315            run_command(self._launch_args + resume_run_args)316 317            # check can run new fully trained pipeline318            pipe = DiffusionPipeline.from_pretrained(tmpdir, safety_checker=None)319            pipe(prompt, num_inference_steps=2)320 321            # check old checkpoints do not exist322            self.assertFalse(os.path.isdir(os.path.join(tmpdir, "checkpoint-2")))323 324            # check new checkpoints exist325            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-4")))326            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-6")))327 328    def test_text_to_image_checkpointing_use_ema(self):329        pretrained_model_name_or_path = "hf-internal-testing/tiny-stable-diffusion-pipe"330        prompt = "a prompt"331 332        with tempfile.TemporaryDirectory() as tmpdir:333            # Run training script with checkpointing334            # max_train_steps == 5, checkpointing_steps == 2335            # Should create checkpoints at steps 2, 4336 337            initial_run_args = f"""338                examples/text_to_image/train_text_to_image.py339                --pretrained_model_name_or_path {pretrained_model_name_or_path}340                --dataset_name hf-internal-testing/dummy_image_text_data341                --resolution 64342                --center_crop343                --random_flip344                --train_batch_size 1345                --gradient_accumulation_steps 1346                --max_train_steps 5347                --learning_rate 5.0e-04348                --scale_lr349                --lr_scheduler constant350                --lr_warmup_steps 0351                --output_dir {tmpdir}352                --checkpointing_steps=2353                --use_ema354                --seed=0355                """.split()356 357            run_command(self._launch_args + initial_run_args)358 359            pipe = DiffusionPipeline.from_pretrained(tmpdir, safety_checker=None)360            pipe(prompt, num_inference_steps=2)361 362            # check checkpoint directories exist363            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-2")))364            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-4")))365 366            # check can run an intermediate checkpoint367            unet = UNet2DConditionModel.from_pretrained(tmpdir, subfolder="checkpoint-2/unet")368            pipe = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path, unet=unet, safety_checker=None)369            pipe(prompt, num_inference_steps=2)370 371            # Remove checkpoint 2 so that we can check only later checkpoints exist after resuming372            shutil.rmtree(os.path.join(tmpdir, "checkpoint-2"))373 374            # Run training script for 7 total steps resuming from checkpoint 4375 376            resume_run_args = f"""377                examples/text_to_image/train_text_to_image.py378                --pretrained_model_name_or_path {pretrained_model_name_or_path}379                --dataset_name hf-internal-testing/dummy_image_text_data380                --resolution 64381                --center_crop382                --random_flip383                --train_batch_size 1384                --gradient_accumulation_steps 1385                --max_train_steps 7386                --learning_rate 5.0e-04387                --scale_lr388                --lr_scheduler constant389                --lr_warmup_steps 0390                --output_dir {tmpdir}391                --checkpointing_steps=2392                --resume_from_checkpoint=checkpoint-4393                --use_ema394                --seed=0395                """.split()396 397            run_command(self._launch_args + resume_run_args)398 399            # check can run new fully trained pipeline400            pipe = DiffusionPipeline.from_pretrained(tmpdir, safety_checker=None)401            pipe(prompt, num_inference_steps=2)402 403            # check old checkpoints do not exist404            self.assertFalse(os.path.isdir(os.path.join(tmpdir, "checkpoint-2")))405 406            # check new checkpoints exist407            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-4")))408            self.assertTrue(os.path.isdir(os.path.join(tmpdir, "checkpoint-6")))409