diff --git a/src/diffusers/pipelines/pipeline_utils.py b/src/diffusers/pipelines/pipeline_utils.py index 381def950169..7c35393b3405 100644 --- a/src/diffusers/pipelines/pipeline_utils.py +++ b/src/diffusers/pipelines/pipeline_utils.py @@ -58,6 +58,7 @@ _is_valid_type, is_accelerate_available, is_accelerate_version, + is_hpu_available, is_torch_npu_available, is_torch_version, is_transformers_version, @@ -445,6 +446,11 @@ def module_is_offloaded(module): f"It seems like you have activated model offloading by calling `enable_model_cpu_offload`, but are now manually moving the pipeline to GPU. It is strongly recommended against doing so as memory gains from offloading are likely to be lost. Offloading automatically takes care of moving the individual components {', '.join(self.components.keys())} to GPU when needed. To make sure offloading works as expected, you should consider moving the pipeline back to CPU: `pipeline.to('cpu')` or removing the move altogether if you use offloading." ) + # Enable generic support for Intel Gaudi accelerator using GPU/HPU migration + if kwargs.pop("hpu_migration", True) and is_hpu_available(): + os.environ["PT_HPU_MAX_COMPOUND_OP_SIZE"] = "1" + logger.debug("Environment variable set: PT_HPU_MAX_COMPOUND_OP_SIZE=1") + module_names, _ = self._get_signature_keys(self) modules = [getattr(self, n, None) for n in module_names] modules = [m for m in modules if isinstance(m, torch.nn.Module)] diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index ed89955ba5a5..3e52f0b70201 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -71,6 +71,7 @@ is_gguf_version, is_google_colab, is_hf_hub_version, + is_hpu_available, is_inflect_available, is_invisible_watermark_available, is_k_diffusion_available, diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index e8d9429f6204..717fe03dbb77 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -336,6 +336,22 @@ def is_timm_available(): return _timm_available +def is_hpu_available(): + if ( + importlib.util.find_spec("habana_frameworks") is None + or importlib.util.find_spec("habana_frameworks.torch") is None + ): + return False + + os.environ["PT_HPU_GPU_MIGRATION"] = "1" + logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1") + + import habana_frameworks.torch # noqa: F401 + import torch + + return hasattr(torch, "hpu") and torch.hpu.is_available() + + # docstyle-ignore FLAX_IMPORT_ERROR = """ {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the