disable test_conversion_when_using_device_map (huggingface#7620)

* disable test * update --------- Co-authored-by: yiyixuxu <yixu310@gmail,com>
EMinsight · Apr 9, 2024 · a341b53 · a341b53
1 parent 8e46d97
commit a341b53
Showing 1 changed file with 39 additions and 40 deletions.
diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py
@@ -1,10 +1,7 @@
-import tempfile
 import unittest
 
-import numpy as np
 import torch
 
-from diffusers import DiffusionPipeline
 from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor
 
 
@@ -80,40 +77,42 @@ def test_only_cross_attention(self):
 
 class DeprecatedAttentionBlockTests(unittest.TestCase):
     def test_conversion_when_using_device_map(self):
-        pipe = DiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None)
-
-        pre_conversion = pipe(
-            "foo",
-            num_inference_steps=2,
-            generator=torch.Generator("cpu").manual_seed(0),
-            output_type="np",
-        ).images
-
-        # the initial conversion succeeds
-        pipe = DiffusionPipeline.from_pretrained(
-            "hf-internal-testing/tiny-stable-diffusion-pipe", device_map="sequential", safety_checker=None
-        )
-
-        conversion = pipe(
-            "foo",
-            num_inference_steps=2,
-            generator=torch.Generator("cpu").manual_seed(0),
-            output_type="np",
-        ).images
-
-        with tempfile.TemporaryDirectory() as tmpdir:
-            # save the converted model
-            pipe.save_pretrained(tmpdir)
-
-            # can also load the converted weights
-            pipe = DiffusionPipeline.from_pretrained(tmpdir, device_map="sequential", safety_checker=None)
-
-        after_conversion = pipe(
-            "foo",
-            num_inference_steps=2,
-            generator=torch.Generator("cpu").manual_seed(0),
-            output_type="np",
-        ).images
-
-        self.assertTrue(np.allclose(pre_conversion, conversion, atol=1e-5))
-        self.assertTrue(np.allclose(conversion, after_conversion, atol=1e-5))
+        # To-DO for Sayak: enable this test again and to test `device_map='balanced'` once we have this in accelerate https://github.com/huggingface/accelerate/pull/2641
+        pass
+        # pipe = DiffusionPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-pipe", safety_checker=None)
+
+        # pre_conversion = pipe(
+        #     "foo",
+        #     num_inference_steps=2,
+        #     generator=torch.Generator("cpu").manual_seed(0),
+        #     output_type="np",
+        # ).images
+
+        # # the initial conversion succeeds
+        # pipe = DiffusionPipeline.from_pretrained(
+        #     "hf-internal-testing/tiny-stable-diffusion-pipe", device_map="sequential", safety_checker=None
+        # )
+
+        # conversion = pipe(
+        #     "foo",
+        #     num_inference_steps=2,
+        #     generator=torch.Generator("cpu").manual_seed(0),
+        #     output_type="np",
+        # ).images
+
+        # with tempfile.TemporaryDirectory() as tmpdir:
+        #     # save the converted model
+        #     pipe.save_pretrained(tmpdir)
+
+        #     # can also load the converted weights
+        #     pipe = DiffusionPipeline.from_pretrained(tmpdir, device_map="sequential", safety_checker=None)
+
+        # after_conversion = pipe(
+        #     "foo",
+        #     num_inference_steps=2,
+        #     generator=torch.Generator("cpu").manual_seed(0),
+        #     output_type="np",
+        # ).images
+
+        # self.assertTrue(np.allclose(pre_conversion, conversion, atol=1e-5))
+        # self.assertTrue(np.allclose(conversion, after_conversion, atol=1e-5))