From be4abffe881d05523084005a94dbb8da336d9535 Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Sun, 14 May 2023 21:26:20 -0700 Subject: [PATCH] Fix tests and remove some testing code. --- .../unidiffuser/pipeline_unidiffuser.py | 8 ------- .../pipelines/unidiffuser/test_unidiffuser.py | 21 ++----------------- 2 files changed, 2 insertions(+), 27 deletions(-) diff --git a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py index 97908cc16d..bf7b4b938c 100644 --- a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +++ b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py @@ -1314,9 +1314,6 @@ class UniDiffuserPipeline(DiffusionPipeline): elif mode in ["img2text", "text"]: latents = prompt_embeds - print(f"Initial latents: {latents}") - print(f"Initial latents shape: {latents.shape}") - # 7. Check that shapes of latents and image match the UNet channels. # TODO @@ -1329,7 +1326,6 @@ class UniDiffuserPipeline(DiffusionPipeline): num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): - print(f"Step {i} / timestep {t}") # predict the noise residual # Also applies classifier-free guidance as described in the UniDiffuser paper noise_pred = self._get_noise_pred( @@ -1348,15 +1344,11 @@ class UniDiffuserPipeline(DiffusionPipeline): width, ) - print(f"noise_pred: {noise_pred}") - # TODO: do we need to worry about sigma space stuff for the scheduler? # compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample - print(f"New latents: {latents}") - # call the callback, if provided if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() diff --git a/tests/pipelines/unidiffuser/test_unidiffuser.py b/tests/pipelines/unidiffuser/test_unidiffuser.py index e7834c61a0..57fb497bcf 100644 --- a/tests/pipelines/unidiffuser/test_unidiffuser.py +++ b/tests/pipelines/unidiffuser/test_unidiffuser.py @@ -585,7 +585,7 @@ class UniDiffuserPipelineSlowTests(unittest.TestCase): inputs = self.get_inputs() del inputs["prompt"] sample = pipe(**inputs) - text = sample.images + text = sample.text expected_text_prefix = "Astronaut " assert text[0][:10] == expected_text_prefix @@ -629,23 +629,6 @@ class UniDiffuserPipelineSlowTests(unittest.TestCase): expected_slice = np.array([0.4702, 0.4666, 0.4446, 0.4829, 0.4468, 0.4565, 0.4663, 0.4956, 0.4277]) assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_unidiffuser_default_text2img_v1_fp16_no_cfg(self): - pipe = UniDiffuserPipeline.from_pretrained("dg845/unidiffuser-diffusers", torch_dtype=torch.float16) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - pipe.enable_attention_slicing() - - inputs = self.get_inputs() - del inputs["image"] - inputs["guidance_scale"] = 0.0 - sample = pipe(**inputs) - image = sample.images - assert image.shape == (1, 512, 512, 3) - - image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.4702, 0.4666, 0.4446, 0.4829, 0.4468, 0.4565, 0.4663, 0.4956, 0.4277]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 - def test_unidiffuser_default_img2text_v1_fp16(self): pipe = UniDiffuserPipeline.from_pretrained("dg845/unidiffuser-diffusers", torch_dtype=torch.float16) pipe.to(torch_device) @@ -655,7 +638,7 @@ class UniDiffuserPipelineSlowTests(unittest.TestCase): inputs = self.get_inputs() del inputs["prompt"] sample = pipe(**inputs) - text = sample.images + text = sample.text expected_text_prefix = "Astronaut " assert text[0][:10] == expected_text_prefix