Unverified Commit 801a5e21 authored by Dhruv Nair's avatar Dhruv Nair Committed by GitHub
Browse files

Cleanup Pass on flaky slow tests for Stable Diffusion (#4455)

* lower num inference steps and precision checkk

* fix flaky inpaint tests

* remove unsued imports

* set unet default attn processor
parent 1edd0deb
...@@ -731,6 +731,7 @@ class StableDiffusionPipelineSlowTests(unittest.TestCase): ...@@ -731,6 +731,7 @@ class StableDiffusionPipelineSlowTests(unittest.TestCase):
def test_stable_diffusion_attention_slicing(self): def test_stable_diffusion_attention_slicing(self):
torch.cuda.reset_peak_memory_stats() torch.cuda.reset_peak_memory_stats()
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16) pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe.unet.set_default_attn_processor()
pipe = pipe.to(torch_device) pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None) pipe.set_progress_bar_config(disable=None)
...@@ -746,6 +747,7 @@ class StableDiffusionPipelineSlowTests(unittest.TestCase): ...@@ -746,6 +747,7 @@ class StableDiffusionPipelineSlowTests(unittest.TestCase):
# disable slicing # disable slicing
pipe.disable_attention_slicing() pipe.disable_attention_slicing()
pipe.unet.set_default_attn_processor()
inputs = self.get_inputs(torch_device, dtype=torch.float16) inputs = self.get_inputs(torch_device, dtype=torch.float16)
image = pipe(**inputs).images image = pipe(**inputs).images
...@@ -1065,9 +1067,9 @@ class StableDiffusionPipelineCkptTests(unittest.TestCase): ...@@ -1065,9 +1067,9 @@ class StableDiffusionPipelineCkptTests(unittest.TestCase):
pipe.to("cuda") pipe.to("cuda")
generator = torch.Generator(device="cpu").manual_seed(0) generator = torch.Generator(device="cpu").manual_seed(0)
image = pipe("a turtle", num_inference_steps=5, generator=generator, output_type="np").images[0] image = pipe("a turtle", num_inference_steps=2, generator=generator, output_type="np").images[0]
assert np.max(np.abs(image - image_ckpt)) < 1e-4 assert np.max(np.abs(image - image_ckpt)) < 1e-3
@nightly @nightly
......
...@@ -592,6 +592,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -592,6 +592,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
"runwayml/stable-diffusion-inpainting", safety_checker=None "runwayml/stable-diffusion-inpainting", safety_checker=None
) )
pipe.vae = vae pipe.vae = vae
pipe.unet.set_default_attn_processor()
pipe.to(torch_device) pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None) pipe.set_progress_bar_config(disable=None)
pipe.enable_attention_slicing() pipe.enable_attention_slicing()
...@@ -601,9 +602,9 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -601,9 +602,9 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
image_slice = image[0, 253:256, 253:256, -1].flatten() image_slice = image[0, 253:256, 253:256, -1].flatten()
assert image.shape == (1, 512, 512, 3) assert image.shape == (1, 512, 512, 3)
expected_slice = np.array([0.0521, 0.0606, 0.0602, 0.0446, 0.0495, 0.0434, 0.1175, 0.1290, 0.1431]) expected_slice = np.array([0.0522, 0.0604, 0.0596, 0.0449, 0.0493, 0.0427, 0.1186, 0.1289, 0.1442])
assert np.abs(expected_slice - image_slice).max() < 6e-4 assert np.abs(expected_slice - image_slice).max() < 1e-3
def test_stable_diffusion_inpaint_fp16(self): def test_stable_diffusion_inpaint_fp16(self):
vae = AsymmetricAutoencoderKL.from_pretrained( vae = AsymmetricAutoencoderKL.from_pretrained(
...@@ -612,6 +613,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -612,6 +613,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
pipe = StableDiffusionInpaintPipeline.from_pretrained( pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, safety_checker=None "runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, safety_checker=None
) )
pipe.unet.set_default_attn_processor()
pipe.vae = vae pipe.vae = vae
pipe.to(torch_device) pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None) pipe.set_progress_bar_config(disable=None)
...@@ -631,6 +633,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -631,6 +633,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
pipe = StableDiffusionInpaintPipeline.from_pretrained( pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", safety_checker=None "runwayml/stable-diffusion-inpainting", safety_checker=None
) )
pipe.unet.set_default_attn_processor()
pipe.vae = vae pipe.vae = vae
pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config) pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config)
pipe.to(torch_device) pipe.to(torch_device)
...@@ -642,7 +645,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -642,7 +645,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
image_slice = image[0, 253:256, 253:256, -1].flatten() image_slice = image[0, 253:256, 253:256, -1].flatten()
assert image.shape == (1, 512, 512, 3) assert image.shape == (1, 512, 512, 3)
expected_slice = np.array([0.0976, 0.1071, 0.1119, 0.1363, 0.1260, 0.1150, 0.3745, 0.3586, 0.3340]) expected_slice = np.array([0.0966, 0.1083, 0.1148, 0.1422, 0.1318, 0.1197, 0.3702, 0.3537, 0.3288])
assert np.abs(expected_slice - image_slice).max() < 5e-3 assert np.abs(expected_slice - image_slice).max() < 5e-3
...@@ -651,6 +654,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -651,6 +654,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
pipe = StableDiffusionInpaintPipeline.from_pretrained( pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", safety_checker=None "runwayml/stable-diffusion-inpainting", safety_checker=None
) )
pipe.unet.set_default_attn_processor()
pipe.vae = vae pipe.vae = vae
pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config) pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to(torch_device) pipe.to(torch_device)
...@@ -660,10 +664,8 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -660,10 +664,8 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
inputs = self.get_inputs(torch_device) inputs = self.get_inputs(torch_device)
image = pipe(**inputs).images image = pipe(**inputs).images
image_slice = image[0, 253:256, 253:256, -1].flatten() image_slice = image[0, 253:256, 253:256, -1].flatten()
assert image.shape == (1, 512, 512, 3) assert image.shape == (1, 512, 512, 3)
expected_slice = np.array([0.8909, 0.8620, 0.9024, 0.8501, 0.8558, 0.9074, 0.8790, 0.7540, 0.9003]) expected_slice = np.array([0.8931, 0.8683, 0.8965, 0.8501, 0.8592, 0.9118, 0.8734, 0.7463, 0.8990])
assert np.abs(expected_slice - image_slice).max() < 6e-3 assert np.abs(expected_slice - image_slice).max() < 6e-3
def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self): def test_stable_diffusion_inpaint_with_sequential_cpu_offloading(self):
...@@ -722,6 +724,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -722,6 +724,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
pipe = StableDiffusionInpaintPipeline.from_pretrained( pipe = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", safety_checker=None "runwayml/stable-diffusion-inpainting", safety_checker=None
) )
pipe.unet.set_default_attn_processor()
pipe.vae = vae pipe.vae = vae
pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config) pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to(torch_device) pipe.to(torch_device)
...@@ -743,6 +746,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -743,6 +746,7 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
vae = AsymmetricAutoencoderKL.from_pretrained("cross-attention/asymmetric-autoencoder-kl-x-1-5") vae = AsymmetricAutoencoderKL.from_pretrained("cross-attention/asymmetric-autoencoder-kl-x-1-5")
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", safety_checker=None) pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", safety_checker=None)
pipe.vae = vae pipe.vae = vae
pipe.unet.set_default_attn_processor()
pipe.to(torch_device) pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None) pipe.set_progress_bar_config(disable=None)
pipe.enable_attention_slicing() pipe.enable_attention_slicing()
...@@ -753,9 +757,8 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te ...@@ -753,9 +757,8 @@ class StableDiffusionInpaintPipelineAsymmetricAutoencoderKLSlowTests(unittest.Te
image_slice = image[0, 253:256, 253:256, -1].flatten() image_slice = image[0, 253:256, 253:256, -1].flatten()
assert image.shape == (1, 512, 512, 3) assert image.shape == (1, 512, 512, 3)
expected_slice = np.array([0.3312, 0.4052, 0.4103, 0.4153, 0.4347, 0.4154, 0.4932, 0.4920, 0.4431]) expected_slice = np.array([0.3296, 0.4041, 0.4097, 0.4145, 0.4342, 0.4152, 0.4927, 0.4931, 0.4430])
assert np.abs(expected_slice - image_slice).max() < 1e-3
assert np.abs(expected_slice - image_slice).max() < 6e-4
def test_download_local(self): def test_download_local(self):
vae = AsymmetricAutoencoderKL.from_pretrained( vae = AsymmetricAutoencoderKL.from_pretrained(
......
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