mirror of
https://github.com/huggingface/diffusers.git
synced 2026-10-29 11:12:00 +03:00
f7fd76adcd
* tests: add cp backend and attention backend tests. * up * up * up * fix ring for flash and flash_3 * generate. * Apply suggestions from code review Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> * up * up --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
23 lines
712 B
Python
23 lines
712 B
Python
import torch
|
|
|
|
from diffusers.models.attention_dispatch import AttentionBackendName
|
|
|
|
|
|
_BF16_REQUIRED_BACKENDS = {
|
|
AttentionBackendName._NATIVE_CUDNN,
|
|
AttentionBackendName.FLASH_HUB,
|
|
AttentionBackendName._FLASH_3_HUB,
|
|
}
|
|
|
|
|
|
def _maybe_cast_to_bf16(backend, model, inputs_dict):
|
|
"""Cast model and floating-point inputs to bfloat16 when the backend requires it."""
|
|
if not backend or backend not in _BF16_REQUIRED_BACKENDS:
|
|
return model, inputs_dict
|
|
model = model.to(dtype=torch.bfloat16)
|
|
inputs_dict = {
|
|
k: v.to(dtype=torch.bfloat16) if isinstance(v, torch.Tensor) and v.is_floating_point() else v
|
|
for k, v in inputs_dict.items()
|
|
}
|
|
return model, inputs_dict
|