Skip to content

Commit 99a064e

Browse files
linoytsabanlinoytsaban
authored
Merge branch 'main' into sana
2 parents 255498b + 049082e commit 99a064e

File tree

11 files changed

+110
-65
lines changed

11 files changed

+110
-65
lines changed

‎src/diffusers/utils/testing_utils.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -635,10 +635,10 @@ def load_numpy(arry: Union[str, np.ndarray], local_path: Optional[str] = None) -
635635
return arry
636636

637637

638-
def load_pt(url: str, map_location: str):
638+
def load_pt(url: str, map_location: Optional[str] = None, weights_only: Optional[bool] = True):
639639
response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT)
640640
response.raise_for_status()
641-
arry = torch.load(BytesIO(response.content), map_location=map_location)
641+
arry = torch.load(BytesIO(response.content), map_location=map_location, weights_only=weights_only)
642642
return arry
643643

644644

‎tests/pipelines/stable_cascade/test_stable_cascade_decoder.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -304,7 +304,8 @@ def test_stable_cascade_decoder(self):
304304

305305
generator = torch.Generator(device="cpu").manual_seed(0)
306306
image_embedding = load_pt(
307-
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_cascade/image_embedding.pt"
307+
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_cascade/image_embedding.pt",
308+
map_location=torch_device,
308309
)
309310

310311
image = pipe(
@@ -320,4 +321,4 @@ def test_stable_cascade_decoder(self):
320321
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_cascade/stable_cascade_decoder_image.npy"
321322
)
322323
max_diff = numpy_cosine_similarity_distance(image.flatten(), expected_image.flatten())
323-
assert max_diff < 1e-4
324+
assert max_diff < 2e-4

‎tests/pipelines/stable_diffusion_k_diffusion/test_stable_diffusion_k_diffusion.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,26 +20,32 @@
2020
import torch
2121

2222
from diffusers import StableDiffusionKDiffusionPipeline
23-
from diffusers.utils.testing_utils import enable_full_determinism, nightly, require_torch_gpu, torch_device
23+
from diffusers.utils.testing_utils import (
24+
backend_empty_cache,
25+
enable_full_determinism,
26+
nightly,
27+
require_torch_accelerator,
28+
torch_device,
29+
)
2430

2531

2632
enable_full_determinism()
2733

2834

2935
@nightly
30-
@require_torch_gpu
36+
@require_torch_accelerator
3137
class StableDiffusionPipelineIntegrationTests(unittest.TestCase):
3238
def setUp(self):
3339
# clean up the VRAM before each test
3440
super().setUp()
3541
gc.collect()
36-
torch.cuda.empty_cache()
42+
backend_empty_cache(torch_device)
3743

3844
def tearDown(self):
3945
# clean up the VRAM after each test
4046
super().tearDown()
4147
gc.collect()
42-
torch.cuda.empty_cache()
48+
backend_empty_cache(torch_device)
4349

4450
def test_stable_diffusion_1(self):
4551
sd_pipe = StableDiffusionKDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

‎tests/pipelines/stable_diffusion_ldm3d/test_stable_diffusion_ldm3d.py‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,13 @@
2828
StableDiffusionLDM3DPipeline,
2929
UNet2DConditionModel,
3030
)
31-
from diffusers.utils.testing_utils import enable_full_determinism, nightly, require_torch_gpu, torch_device
31+
from diffusers.utils.testing_utils import (
32+
backend_empty_cache,
33+
enable_full_determinism,
34+
nightly,
35+
require_torch_accelerator,
36+
torch_device,
37+
)
3238

3339
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
3440

@@ -205,17 +211,17 @@ def test_stable_diffusion_negative_prompt(self):
205211

206212

207213
@nightly
208-
@require_torch_gpu
214+
@require_torch_accelerator
209215
class StableDiffusionLDM3DPipelineSlowTests(unittest.TestCase):
210216
def setUp(self):
211217
super().setUp()
212218
gc.collect()
213-
torch.cuda.empty_cache()
219+
backend_empty_cache(torch_device)
214220

215221
def tearDown(self):
216222
super().tearDown()
217223
gc.collect()
218-
torch.cuda.empty_cache()
224+
backend_empty_cache(torch_device)
219225

220226
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
221227
generator = torch.Generator(device=generator_device).manual_seed(seed)
@@ -256,17 +262,17 @@ def test_ldm3d_stable_diffusion(self):
256262

257263

258264
@nightly
259-
@require_torch_gpu
265+
@require_torch_accelerator
260266
class StableDiffusionPipelineNightlyTests(unittest.TestCase):
261267
def setUp(self):
262268
super().setUp()
263269
gc.collect()
264-
torch.cuda.empty_cache()
270+
backend_empty_cache(torch_device)
265271

266272
def tearDown(self):
267273
super().tearDown()
268274
gc.collect()
269-
torch.cuda.empty_cache()
275+
backend_empty_cache(torch_device)
270276

271277
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
272278
generator = torch.Generator(device=generator_device).manual_seed(seed)

‎tests/pipelines/stable_diffusion_sag/test_stable_diffusion_sag.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,13 @@
2929
StableDiffusionSAGPipeline,
3030
UNet2DConditionModel,
3131
)
32-
from diffusers.utils.testing_utils import enable_full_determinism, nightly, require_torch_gpu, torch_device
32+
from diffusers.utils.testing_utils import (
33+
backend_empty_cache,
34+
enable_full_determinism,
35+
nightly,
36+
require_torch_accelerator,
37+
torch_device,
38+
)
3339

3440
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
3541
from ..test_pipelines_common import (
@@ -162,19 +168,19 @@ def test_encode_prompt_works_in_isolation(self):
162168

163169

164170
@nightly
165-
@require_torch_gpu
171+
@require_torch_accelerator
166172
class StableDiffusionPipelineIntegrationTests(unittest.TestCase):
167173
def setUp(self):
168174
# clean up the VRAM before each test
169175
super().setUp()
170176
gc.collect()
171-
torch.cuda.empty_cache()
177+
backend_empty_cache(torch_device)
172178

173179
def tearDown(self):
174180
# clean up the VRAM after each test
175181
super().tearDown()
176182
gc.collect()
177-
torch.cuda.empty_cache()
183+
backend_empty_cache(torch_device)
178184

179185
def test_stable_diffusion_1(self):
180186
sag_pipe = StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")

‎tests/pipelines/stable_unclip/test_stable_unclip.py‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,17 @@
1313
UNet2DConditionModel,
1414
)
1515
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
16-
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, nightly, require_torch_gpu, torch_device
16+
from diffusers.utils.testing_utils import (
17+
backend_empty_cache,
18+
backend_max_memory_allocated,
19+
backend_reset_max_memory_allocated,
20+
backend_reset_peak_memory_stats,
21+
enable_full_determinism,
22+
load_numpy,
23+
nightly,
24+
require_torch_accelerator,
25+
torch_device,
26+
)
1727

1828
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
1929
from ..test_pipelines_common import (
@@ -190,19 +200,19 @@ def test_encode_prompt_works_in_isolation(self):
190200

191201

192202
@nightly
193-
@require_torch_gpu
203+
@require_torch_accelerator
194204
class StableUnCLIPPipelineIntegrationTests(unittest.TestCase):
195205
def setUp(self):
196206
# clean up the VRAM before each test
197207
super().setUp()
198208
gc.collect()
199-
torch.cuda.empty_cache()
209+
backend_empty_cache(torch_device)
200210

201211
def tearDown(self):
202212
# clean up the VRAM after each test
203213
super().tearDown()
204214
gc.collect()
205-
torch.cuda.empty_cache()
215+
backend_empty_cache(torch_device)
206216

207217
def test_stable_unclip(self):
208218
expected_image = load_numpy(
@@ -226,9 +236,9 @@ def test_stable_unclip(self):
226236
assert_mean_pixel_difference(image, expected_image)
227237

228238
def test_stable_unclip_pipeline_with_sequential_cpu_offloading(self):
229-
torch.cuda.empty_cache()
230-
torch.cuda.reset_max_memory_allocated()
231-
torch.cuda.reset_peak_memory_stats()
239+
backend_empty_cache(torch_device)
240+
backend_reset_max_memory_allocated(torch_device)
241+
backend_reset_peak_memory_stats(torch_device)
232242

233243
pipe = StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l", torch_dtype=torch.float16)
234244
pipe.set_progress_bar_config(disable=None)
@@ -242,6 +252,6 @@ def test_stable_unclip_pipeline_with_sequential_cpu_offloading(self):
242252
output_type="np",
243253
)
244254

245-
mem_bytes = torch.cuda.max_memory_allocated()
255+
mem_bytes = backend_max_memory_allocated(torch_device)
246256
# make sure that less than 7 GB is allocated
247257
assert mem_bytes < 7 * 10**9

‎tests/pipelines/stable_unclip/test_stable_unclip_img2img.py‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,16 @@
1818
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
1919
from diffusers.utils.import_utils import is_xformers_available
2020
from diffusers.utils.testing_utils import (
21+
backend_empty_cache,
22+
backend_max_memory_allocated,
23+
backend_reset_max_memory_allocated,
24+
backend_reset_peak_memory_stats,
2125
enable_full_determinism,
2226
floats_tensor,
2327
load_image,
2428
load_numpy,
2529
nightly,
26-
require_torch_gpu,
30+
require_torch_accelerator,
2731
skip_mps,
2832
torch_device,
2933
)
@@ -213,19 +217,19 @@ def test_encode_prompt_works_in_isolation(self):
213217

214218

215219
@nightly
216-
@require_torch_gpu
220+
@require_torch_accelerator
217221
class StableUnCLIPImg2ImgPipelineIntegrationTests(unittest.TestCase):
218222
def setUp(self):
219223
# clean up the VRAM before each test
220224
super().setUp()
221225
gc.collect()
222-
torch.cuda.empty_cache()
226+
backend_empty_cache(torch_device)
223227

224228
def tearDown(self):
225229
# clean up the VRAM after each test
226230
super().tearDown()
227231
gc.collect()
228-
torch.cuda.empty_cache()
232+
backend_empty_cache(torch_device)
229233

230234
def test_stable_unclip_l_img2img(self):
231235
input_image = load_image(
@@ -286,9 +290,9 @@ def test_stable_unclip_img2img_pipeline_with_sequential_cpu_offloading(self):
286290
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png"
287291
)
288292

289-
torch.cuda.empty_cache()
290-
torch.cuda.reset_max_memory_allocated()
291-
torch.cuda.reset_peak_memory_stats()
293+
backend_empty_cache(torch_device)
294+
backend_reset_max_memory_allocated(torch_device)
295+
backend_reset_peak_memory_stats(torch_device)
292296

293297
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
294298
"fusing/stable-unclip-2-1-h-img2img", torch_dtype=torch.float16
@@ -304,6 +308,6 @@ def test_stable_unclip_img2img_pipeline_with_sequential_cpu_offloading(self):
304308
output_type="np",
305309
)
306310

307-
mem_bytes = torch.cuda.max_memory_allocated()
311+
mem_bytes = backend_max_memory_allocated(torch_device)
308312
# make sure that less than 7 GB is allocated
309313
assert mem_bytes < 7 * 10**9

‎tests/pipelines/text_to_video_synthesis/test_text_to_video_zero.py‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,37 +19,44 @@
1919
import torch
2020

2121
from diffusers import DDIMScheduler, TextToVideoZeroPipeline
22-
from diffusers.utils.testing_utils import load_pt, nightly, require_torch_gpu
22+
from diffusers.utils.testing_utils import (
23+
backend_empty_cache,
24+
load_pt,
25+
nightly,
26+
require_torch_accelerator,
27+
torch_device,
28+
)
2329

2430
from ..test_pipelines_common import assert_mean_pixel_difference
2531

2632

2733
@nightly
28-
@require_torch_gpu
34+
@require_torch_accelerator
2935
class TextToVideoZeroPipelineSlowTests(unittest.TestCase):
3036
def setUp(self):
3137
# clean up the VRAM before each test
3238
super().setUp()
3339
gc.collect()
34-
torch.cuda.empty_cache()
40+
backend_empty_cache(torch_device)
3541

3642
def tearDown(self):
3743
# clean up the VRAM after each test
3844
super().tearDown()
3945
gc.collect()
40-
torch.cuda.empty_cache()
46+
backend_empty_cache(torch_device)
4147

4248
def test_full_model(self):
4349
model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
44-
pipe = TextToVideoZeroPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")
50+
pipe = TextToVideoZeroPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(torch_device)
4551
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
46-
generator = torch.Generator(device="cuda").manual_seed(0)
52+
generator = torch.Generator(device="cpu").manual_seed(0)
4753

4854
prompt = "A bear is playing a guitar on Times Square"
4955
result = pipe(prompt=prompt, generator=generator).images
5056

5157
expected_result = load_pt(
52-
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text-to-video/A bear is playing a guitar on Times Square.pt"
58+
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text-to-video/A bear is playing a guitar on Times Square.pt",
59+
weights_only=False,
5360
)
5461

5562
assert_mean_pixel_difference(result, expected_result)

0 commit comments

Comments
 (0)