[Tests] Better prints (#1043)

Author: patrickvonplaten

src/diffusers/utils/__init__.py

@@ -40,7 +40,15 @@
 
 
 if is_torch_available():
-    from .testing_utils import floats_tensor, load_image, parse_flag_from_env, require_torch_gpu, slow, torch_device
+    from .testing_utils import (
+        floats_tensor,
+        load_image,
+        parse_flag_from_env,
+        require_torch_gpu,
+        slow,
+        torch_all_close,
+        torch_device,
+    )
 
 
 logger = get_logger(__name__)

src/diffusers/utils/testing_utils.py

@@ -34,6 +34,14 @@
         torch_device = "mps" if (mps_backend_registered and torch.backends.mps.is_available()) else torch_device
 
 
+def torch_all_close(a, b, *args, **kwargs):
+    if not is_torch_available():
+        raise ValueError("PyTorch needs to be installed to use this function.")
+    if not torch.allclose(a, b, *args, **kwargs):
+        assert False, f"Max diff is absolute {(a - b).abs().max()}. Diff tensor is {(a - b).abs()}."
+    return True
+
+
 def get_tests_dir(append_path=None):
     """
     Args:

tests/models/test_models_unet_2d.py

@@ -21,7 +21,7 @@
 import torch
 
 from diffusers import UNet2DConditionModel, UNet2DModel
-from diffusers.utils import floats_tensor, require_torch_gpu, slow, torch_device
+from diffusers.utils import floats_tensor, require_torch_gpu, slow, torch_all_close, torch_device
 from parameterized import parameterized
 
 from ..test_modeling_common import ModelTesterMixin
@@ -156,7 +156,7 @@ def test_from_pretrained_accelerate_wont_change_results(self):
         model_normal_load.eval()
         arr_normal_load = model_normal_load(noise, time_step)["sample"]
 
-        assert torch.allclose(arr_accelerate, arr_normal_load, rtol=1e-3)
+        assert torch_all_close(arr_accelerate, arr_normal_load, rtol=1e-3)
 
     @unittest.skipIf(torch_device != "cuda", "This test is supposed to run on GPU")
     def test_memory_footprint_gets_reduced(self):
@@ -207,7 +207,7 @@ def test_output_pretrained(self):
         expected_output_slice = torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800])
         # fmt: on
 
-        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-3))
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-3))
 
 
 class UNet2DConditionModelTests(ModelTesterMixin, unittest.TestCase):
@@ -287,7 +287,7 @@ def test_gradient_checkpointing(self):
         named_params = dict(model.named_parameters())
         named_params_2 = dict(model_2.named_parameters())
         for name, param in named_params.items():
-            self.assertTrue(torch.allclose(param.grad.data, named_params_2[name].grad.data, atol=5e-5))
+            self.assertTrue(torch_all_close(param.grad.data, named_params_2[name].grad.data, atol=5e-5))
 
 
 class NCSNppModelTests(ModelTesterMixin, unittest.TestCase):
@@ -377,7 +377,7 @@ def test_output_pretrained_ve_mid(self):
         expected_output_slice = torch.tensor([-4836.2231, -6487.1387, -3816.7969, -7964.9253, -10966.2842, -20043.6016, 8137.0571, 2340.3499, 544.6114])
         # fmt: on
 
-        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))
 
     def test_output_pretrained_ve_large(self):
         model = UNet2DModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update")
@@ -402,7 +402,7 @@ def test_output_pretrained_ve_large(self):
         expected_output_slice = torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256])
         # fmt: on
 
-        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))
 
     def test_forward_with_norm_groups(self):
         # not required for this model
@@ -464,7 +464,7 @@ def test_compvis_sd_v1_4(self, seed, timestep, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -490,7 +490,7 @@ def test_compvis_sd_v1_4_fp16(self, seed, timestep, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -515,7 +515,7 @@ def test_compvis_sd_v1_5(self, seed, timestep, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -541,7 +541,7 @@ def test_compvis_sd_v1_5_fp16(self, seed, timestep, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -566,7 +566,7 @@ def test_compvis_sd_inpaint(self, seed, timestep, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -592,4 +592,4 @@ def test_compvis_sd_inpaint_fp16(self, seed, timestep, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)

tests/models/test_models_vae.py

@@ -20,7 +20,7 @@
 
 from diffusers import AutoencoderKL
 from diffusers.modeling_utils import ModelMixin
-from diffusers.utils import floats_tensor, require_torch_gpu, slow, torch_device
+from diffusers.utils import floats_tensor, require_torch_gpu, slow, torch_all_close, torch_device
 from parameterized import parameterized
 
 from ..test_modeling_common import ModelTesterMixin
@@ -131,7 +131,7 @@ def test_output_pretrained(self):
                 [-0.2421, 0.4642, 0.2507, -0.0438, 0.0682, 0.3160, -0.2018, -0.0727, 0.2485]
             )
 
-        self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-2))
+        self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2))
 
 
 @slow
@@ -185,7 +185,7 @@ def test_stable_diffusion(self, seed, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -209,7 +209,7 @@ def test_stable_diffusion_fp16(self, seed, expected_slice):
         output_slice = sample[-1, -2:, :2, -2:].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -231,7 +231,7 @@ def test_stable_diffusion_mode(self, seed, expected_slice):
         output_slice = sample[-1, -2:, -2:, :2].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -254,7 +254,7 @@ def test_stable_diffusion_decode(self, seed, expected_slice):
         output_slice = sample[-1, -2:, :2, -2:].flatten().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -276,7 +276,7 @@ def test_stable_diffusion_decode_fp16(self, seed, expected_slice):
         output_slice = sample[-1, -2:, :2, -2:].flatten().float().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)
 
     @parameterized.expand(
         [
@@ -300,4 +300,4 @@ def test_stable_diffusion_encode_sample(self, seed, expected_slice):
         output_slice = sample[0, -1, -3:, -3:].flatten().cpu()
         expected_output_slice = torch.tensor(expected_slice)
 
-        assert torch.allclose(output_slice, expected_output_slice, atol=1e-4)
+        assert torch_all_close(output_slice, expected_output_slice, atol=1e-4)