tolerance of 5e-3

piercus · Feb 21, 2024 · c40e555 · c40e555
1 parent 8ab4e8b
commit c40e555
Showing 1 changed file with 13 additions and 11 deletions.
diff --git a/tests/e2e/test_diffusion.py b/tests/e2e/test_diffusion.py
@@ -776,41 +776,42 @@ def test_diffusion_batch2(
     step = sd15.steps[0]
 
     manual_seed(2)
-    x_b2 = torch.randn(2, 4, 64, 64, device=test_device)
+    rand_b2 = torch.randn(2, 4, 64, 64, device=test_device)
 
     x_b2 = sd15(
-        x_b2,
+        rand_b2,
         step=step,
         clip_text_embedding=clip_text_embedding_b2,
         condition_scale=7.5,
     )
 
     assert x_b2.shape == (2, 4, 64, 64)
 
-    x_1 = x_b2[0:1]
+    rand_1 = rand_b2[0:1]
     clip_text_embedding_1 = sd15.compute_clip_text_embedding(
         text=[prompt1], negative_text=[negative_prompt1]
     )
     x_1 = sd15(
-        x_1,
+        rand_1,
         step=step,
         clip_text_embedding=clip_text_embedding_1,
         condition_scale=7.5,
     )
 
-    x_2 = x_b2[1:2]
-    clip_text_embedding_1 = sd15.compute_clip_text_embedding(
+    rand_2 = rand_b2[1:2]
+    clip_text_embedding_2 = sd15.compute_clip_text_embedding(
         text=[prompt2], negative_text=[negative_prompt2]
     )
     x_2 = sd15(
-        x_2,
+        rand_2,
         step=step,
         clip_text_embedding=clip_text_embedding_2,
         condition_scale=7.5,
     )
 
-    assert torch.allclose(x_b2[0], x_1[0])
-    assert torch.allclose(x_b2[1], x_2[0])
+    # The 5e-3 tolerance is detailed in https://github.com/finegrain-ai/refiners/pull/263#issuecomment-1956404911
+    assert torch.allclose(x_b2[0], x_1[0], atol=5e-3, rtol=0), f"Batch 2 and batch1 output should be the same and are distant of {torch.max((x_b2[0] - x_1[0]).abs()).item()}"
+    assert torch.allclose(x_b2[1], x_2[0], atol=5e-3, rtol=0), f"Batch 2 and batch1 output should be the same and are distant of {torch.max((x_b2[1] - x_2[0]).abs()).item()}"
 
 @no_grad()
 def test_diffusion_std_random_init_euler(
@@ -1373,8 +1374,9 @@ def test_diffusion_sdxl_batch2(sdxl_ddim: StableDiffusion_XL) -> None:
         time_ids=time_ids,
     )
 
-    assert torch.allclose(x_b2[0], x_1[0])
-    assert torch.allclose(x_b2[1], x_2[0])
+    # The 5e-3 tolerance is detailed in https://github.com/finegrain-ai/refiners/pull/263#issuecomment-1956404911
+    assert torch.allclose(x_b2[0], x_1[0], atol=5e-3, rtol=0), f"Batch 2 and batch1 output should be the same and are distant of {torch.max((x_b2[0] - x_1[0]).abs()).item()}"
+    assert torch.allclose(x_b2[1], x_2[0], atol=5e-3, rtol=0), f"Batch 2 and batch1 output should be the same and are distant of {torch.max((x_b2[1] - x_2[0]).abs()).item()}"
 
 
 @no_grad()