Unverified Commit eef3d659 authored by Yao Matrix's avatar Yao Matrix Committed by GitHub
Browse files

enable 2 test cases on XPU (#11332)



* enable 2 test cases on XPU
Signed-off-by: default avatarYAO Matrix <matrix.yao@intel.com>

* Apply style fixes

---------
Signed-off-by: default avatarYAO Matrix <matrix.yao@intel.com>
Co-authored-by: default avatargithub-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: default avatarDhruv Nair <dhruv.nair@gmail.com>
parent ee6ad51d
......@@ -523,13 +523,15 @@ class SlowBnb8bitTests(Base8bitTests):
torch_dtype=torch.float16,
device_map=torch_device,
)
# CUDA device placement works.
device = torch_device if torch_device != "rocm" else "cuda"
pipeline_8bit = DiffusionPipeline.from_pretrained(
self.model_name,
transformer=transformer_8bit,
text_encoder_3=text_encoder_3_8bit,
torch_dtype=torch.float16,
).to("cuda")
).to(device)
# Check if inference works.
_ = pipeline_8bit("table", max_sequence_length=20, num_inference_steps=2)
......
from diffusers.utils import is_torch_available
from diffusers.utils.testing_utils import (
backend_empty_cache,
backend_max_memory_allocated,
backend_reset_peak_memory_stats,
torch_device,
)
if is_torch_available():
......@@ -30,9 +36,9 @@ if is_torch_available():
@torch.no_grad()
@torch.inference_mode()
def get_memory_consumption_stat(model, inputs):
torch.cuda.reset_peak_memory_stats()
torch.cuda.empty_cache()
backend_reset_peak_memory_stats(torch_device)
backend_empty_cache(torch_device)
model(**inputs)
max_memory_mem_allocated = torch.cuda.max_memory_allocated()
return max_memory_mem_allocated
max_mem_allocated = backend_max_memory_allocated(torch_device)
return max_mem_allocated
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment