mirror of
https://github.com/huggingface/diffusers.git
synced 2026-01-27 17:22:53 +03:00
enable 2 test cases on XPU (#11332)
* enable 2 test cases on XPU Signed-off-by: YAO Matrix <matrix.yao@intel.com> * Apply style fixes --------- Signed-off-by: YAO Matrix <matrix.yao@intel.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -1,4 +1,10 @@
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user