From 066d0bb6551f7ab0524b8ab31b358d780325b910 Mon Sep 17 00:00:00 2001 From: "Takahashi, Akihiro" Date: Wed, 10 Jan 2024 13:28:22 -0800 Subject: [PATCH] Add SDXL Turbo support --- examples/stable-diffusion/README.md | 7 ++++-- .../text_to_image_generation.py | 16 ++++++++++--- tests/test_diffusers.py | 23 +++++++++++++++++-- 3 files changed, 39 insertions(+), 7 deletions(-) mode change 100644 => 100755 examples/stable-diffusion/text_to_image_generation.py diff --git a/examples/stable-diffusion/README.md b/examples/stable-diffusion/README.md index f14ffc7b87..3c8d3d170f 100644 --- a/examples/stable-diffusion/README.md +++ b/examples/stable-diffusion/README.md @@ -131,7 +131,7 @@ python text_to_image_generation.py \ --use_habana \ --use_hpu_graphs \ --gaudi_config Habana/stable-diffusion \ - --bf16 + --bf16 ``` > HPU graphs are recommended when generating images by batches to get the fastest possible generations. @@ -192,7 +192,10 @@ python text_to_image_generation.py \ --use_habana \ --use_hpu_graphs \ --gaudi_config Habana/stable-diffusion \ - --bf16 + --bf16 \ + --num_inference_steps 1 \ + --guidance_scale 0.0 \ + --timestep_spacing trailing ``` > HPU graphs are recommended when generating images by batches to get the fastest possible generations. diff --git a/examples/stable-diffusion/text_to_image_generation.py b/examples/stable-diffusion/text_to_image_generation.py old mode 100644 new mode 100755 index 972a497191..088279d706 --- a/examples/stable-diffusion/text_to_image_generation.py +++ b/examples/stable-diffusion/text_to_image_generation.py @@ -61,6 +61,13 @@ def main(): help="Name of scheduler", ) + parser.add_argument( + "--timestep_spacing", + default="linspace", + choices=["linspace", "leading", "trailing"], + type=str, + help="The way the timesteps should be scaled.", + ) # Pipeline arguments parser.add_argument( "--prompts", @@ -207,14 +214,17 @@ def main(): logger.setLevel(logging.INFO) # Initialize the scheduler and the generation pipeline + kwargs = {"timestep_spacing": args.timestep_spacing} if args.scheduler == "euler_discrete": - scheduler = GaudiEulerDiscreteScheduler.from_pretrained(args.model_name_or_path, subfolder="scheduler") + scheduler = GaudiEulerDiscreteScheduler.from_pretrained( + args.model_name_or_path, subfolder="scheduler", **kwargs + ) elif args.scheduler == "euler_ancestral_discrete": scheduler = GaudiEulerAncestralDiscreteScheduler.from_pretrained( - args.model_name_or_path, subfolder="scheduler" + args.model_name_or_path, subfolder="scheduler", **kwargs ) else: - scheduler = GaudiDDIMScheduler.from_pretrained(args.model_name_or_path, subfolder="scheduler") + scheduler = GaudiDDIMScheduler.from_pretrained(args.model_name_or_path, subfolder="scheduler", **kwargs) kwargs = { "scheduler": scheduler, diff --git a/tests/test_diffusers.py b/tests/test_diffusers.py index b79bb27fa5..57e270d9c6 100644 --- a/tests/test_diffusers.py +++ b/tests/test_diffusers.py @@ -818,7 +818,7 @@ class GaudiStableDiffusionXLPipelineTester(TestCase): Tests the StableDiffusionXLPipeline for Gaudi. """ - def get_dummy_components(self, time_cond_proj_dim=None): + def get_dummy_components(self, time_cond_proj_dim=None, timestep_spacing="leading"): torch.manual_seed(0) unet = UNet2DConditionModel( block_out_channels=(2, 4), @@ -844,7 +844,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): beta_end=0.012, steps_offset=1, beta_schedule="scaled_linear", - timestep_spacing="leading", + timestep_spacing=timestep_spacing, ) torch.manual_seed(0) vae = AutoencoderKL( @@ -933,6 +933,25 @@ def test_stable_diffusion_xl_euler_ancestral(self): expected_slice = np.array([0.4675, 0.5173, 0.4611, 0.4067, 0.5250, 0.4674, 0.5446, 0.5094, 0.4791]) self.assertLess(np.abs(image_slice.flatten() - expected_slice).max(), 1e-2) + def test_stable_diffusion_xl_turbo_euler_ancestral(self): + device = "cpu" # ensure determinism for the device-dependent torch.Generator + components = self.get_dummy_components(timestep_spacing="trailing") + gaudi_config = GaudiConfig(use_torch_autocast=False) + + sd_pipe = GaudiStableDiffusionXLPipeline(use_habana=True, gaudi_config=gaudi_config, **components) + sd_pipe.scheduler = GaudiEulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) + + sd_pipe.set_progress_bar_config(disable=None) + + inputs = self.get_dummy_inputs(device) + image = sd_pipe(**inputs).images[0] + + image_slice = image[-3:, -3:, -1] + + self.assertEqual(image.shape, (64, 64, 3)) + expected_slice = np.array([0.4675, 0.5173, 0.4611, 0.4067, 0.5250, 0.4674, 0.5446, 0.5094, 0.4791]) + self.assertLess(np.abs(image_slice.flatten() - expected_slice).max(), 1e-2) + @parameterized.expand(["pil", "np", "latent"]) def test_stable_diffusion_xl_output_types(self, output_type): components = self.get_dummy_components()