Skip to content

Commit f3a519f

Browse files
faaanyfaaany
committed
update
1 parent 2f3ad32 commit f3a519f

File tree

4 files changed

+31
-37
lines changed

4 files changed

+31
-37
lines changed

‎src/diffusers/utils/testing_utils.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1085,8 +1085,8 @@ def _is_torch_fp64_available(device):
10851085
"default": torch.manual_seed,
10861086
}
10871087
BACKEND_RESET_PEAK_MEMORY_STATS = {
1088-
"cuda": torch.cuda.reset_peak_memory_stats(),
1089-
"xpu": torch.xpu.reset_peak_memory_stats(),
1088+
"cuda": torch.cuda.reset_peak_memory_stats,
1089+
"xpu": torch.xpu.reset_peak_memory_stats,
10901090
"default": None,
10911091
}
10921092

‎tests/pipelines/deepfloyd_if/test_if_inpainting.py‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
import gc
1716
import random
1817
import unittest
1918

@@ -24,9 +23,10 @@
2423
from diffusers.utils.import_utils import is_xformers_available
2524
from diffusers.utils.testing_utils import (
2625
floats_tensor,
26+
flush_memory,
2727
load_numpy,
2828
require_accelerator,
29-
require_torch_gpu,
29+
require_torch_accelerator,
3030
skip_mps,
3131
slow,
3232
torch_device,
@@ -99,30 +99,26 @@ def test_inference_batch_single_identical(self):
9999

100100

101101
@slow
102-
@require_torch_gpu
102+
@require_torch_accelerator
103103
class IFInpaintingPipelineSlowTests(unittest.TestCase):
104104
def setUp(self):
105105
# clean up the VRAM before each test
106106
super().setUp()
107-
gc.collect()
108-
torch.cuda.empty_cache()
107+
flush_memory(torch_device, gc_collect=True)
109108

110109
def tearDown(self):
111110
# clean up the VRAM after each test
112111
super().tearDown()
113-
gc.collect()
114-
torch.cuda.empty_cache()
112+
flush_memory(torch_device, gc_collect=True)
115113

116114
def test_if_inpainting(self):
117115
pipe = IFInpaintingPipeline.from_pretrained(
118116
"DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16
119117
)
120118
pipe.unet.set_attn_processor(AttnAddedKVProcessor())
121-
pipe.enable_model_cpu_offload()
119+
pipe.enable_model_cpu_offload(device=torch_device)
122120

123-
torch.cuda.empty_cache()
124-
torch.cuda.reset_max_memory_allocated()
125-
torch.cuda.reset_peak_memory_stats()
121+
flush_memory(torch_device, reset_mem_stats=True)
126122

127123
image = floats_tensor((1, 3, 64, 64), rng=random.Random(0)).to(torch_device)
128124
mask_image = floats_tensor((1, 3, 64, 64), rng=random.Random(1)).to(torch_device)

‎tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
import gc
1716
import random
1817
import unittest
1918

@@ -24,9 +23,10 @@
2423
from diffusers.utils.import_utils import is_xformers_available
2524
from diffusers.utils.testing_utils import (
2625
floats_tensor,
26+
flush_memory,
2727
load_numpy,
2828
require_accelerator,
29-
require_torch_gpu,
29+
require_torch_accelerator,
3030
skip_mps,
3131
slow,
3232
torch_device,
@@ -101,31 +101,27 @@ def test_inference_batch_single_identical(self):
101101

102102

103103
@slow
104-
@require_torch_gpu
104+
@require_torch_accelerator
105105
class IFInpaintingSuperResolutionPipelineSlowTests(unittest.TestCase):
106106
def setUp(self):
107107
# clean up the VRAM before each test
108108
super().setUp()
109-
gc.collect()
110-
torch.cuda.empty_cache()
109+
flush_memory(torch_device, gc_collect=True)
111110

112111
def tearDown(self):
113112
# clean up the VRAM after each test
114113
super().tearDown()
115-
gc.collect()
116-
torch.cuda.empty_cache()
114+
flush_memory(torch_device, gc_collect=True)
117115

118116
def test_if_inpainting_superresolution(self):
119117
pipe = IFInpaintingSuperResolutionPipeline.from_pretrained(
120118
"DeepFloyd/IF-II-L-v1.0", variant="fp16", torch_dtype=torch.float16
121119
)
122120
pipe.unet.set_attn_processor(AttnAddedKVProcessor())
123-
pipe.enable_model_cpu_offload()
121+
pipe.enable_model_cpu_offload(device=torch_device)
124122

125123
# Super resolution test
126-
torch.cuda.empty_cache()
127-
torch.cuda.reset_max_memory_allocated()
128-
torch.cuda.reset_peak_memory_stats()
124+
flush_memory(torch_device, reset_mem_stats=True)
129125

130126
generator = torch.Generator(device="cpu").manual_seed(0)
131127

@@ -147,7 +143,10 @@ def test_if_inpainting_superresolution(self):
147143

148144
assert image.shape == (256, 256, 3)
149145

150-
mem_bytes = torch.cuda.max_memory_allocated()
146+
if torch_device == "cuda":
147+
mem_bytes = torch.cuda.max_memory_allocated()
148+
elif torch_device == "xpu":
149+
mem_bytes = torch.xpu.max_memory_allocated()
151150
assert mem_bytes < 12 * 10**9
152151

153152
expected_image = load_numpy(

‎tests/pipelines/deepfloyd_if/test_if_superresolution.py‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
import gc
1716
import random
1817
import unittest
1918

@@ -24,9 +23,10 @@
2423
from diffusers.utils.import_utils import is_xformers_available
2524
from diffusers.utils.testing_utils import (
2625
floats_tensor,
26+
flush_memory,
2727
load_numpy,
2828
require_accelerator,
29-
require_torch_gpu,
29+
require_torch_accelerator,
3030
skip_mps,
3131
slow,
3232
torch_device,
@@ -94,31 +94,27 @@ def test_inference_batch_single_identical(self):
9494

9595

9696
@slow
97-
@require_torch_gpu
97+
@require_torch_accelerator
9898
class IFSuperResolutionPipelineSlowTests(unittest.TestCase):
9999
def setUp(self):
100100
# clean up the VRAM before each test
101101
super().setUp()
102-
gc.collect()
103-
torch.cuda.empty_cache()
102+
flush_memory(torch_device, gc_collect=True)
104103

105104
def tearDown(self):
106105
# clean up the VRAM after each test
107106
super().tearDown()
108-
gc.collect()
109-
torch.cuda.empty_cache()
107+
flush_memory(torch_device, gc_collect=True)
110108

111109
def test_if_superresolution(self):
112110
pipe = IFSuperResolutionPipeline.from_pretrained(
113111
"DeepFloyd/IF-II-L-v1.0", variant="fp16", torch_dtype=torch.float16
114112
)
115113
pipe.unet.set_attn_processor(AttnAddedKVProcessor())
116-
pipe.enable_model_cpu_offload()
114+
pipe.enable_model_cpu_offload(device=torch_device)
117115

118116
# Super resolution test
119-
torch.cuda.empty_cache()
120-
torch.cuda.reset_max_memory_allocated()
121-
torch.cuda.reset_peak_memory_stats()
117+
flush_memory(torch_device, reset_mem_stats=True)
122118

123119
image = floats_tensor((1, 3, 64, 64), rng=random.Random(0)).to(torch_device)
124120
generator = torch.Generator(device="cpu").manual_seed(0)
@@ -134,7 +130,10 @@ def test_if_superresolution(self):
134130

135131
assert image.shape == (256, 256, 3)
136132

137-
mem_bytes = torch.cuda.max_memory_allocated()
133+
if torch_device == "cuda":
134+
mem_bytes = torch.cuda.max_memory_allocated()
135+
elif torch_device == "xpu":
136+
mem_bytes = torch.xpu.max_memory_allocated()
138137
assert mem_bytes < 12 * 10**9
139138

140139
expected_image = load_numpy(

0 commit comments

Comments
 (0)