AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

disable test_conversion_when_using_device_map (#7620)

Browse Source 
* disable test

* update

---------

Co-authored-by: yiyixuxu <yixu310@gmail,com>
This commit is contained in:
YiYi Xu
2024-04-09 09:01:19 -10:00
committed by GitHub
parent 8e46d97cd8
commit a341b536a8
1 changed files with 32 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
tests/models/test_attention_processor.py
View File

@@ -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 @@ class AttnAddedKVProcessorTests(unittest.TestCase):
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)
# 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
# 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
)
# # 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
# 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)
# 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)
# # 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
# 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))
# self.assertTrue(np.allclose(pre_conversion, conversion, atol=1e-5))
# self.assertTrue(np.allclose(conversion, after_conversion, atol=1e-5))