Decouple soft HPU check from hard device validation to support HPU migration

Signed-off-by: Daniel Socek <[email protected]>

Author: dsocek

src/diffusers/pipelines/pipeline_utils.py

@@ -447,7 +447,16 @@ def module_is_offloaded(module):
             )
 
         # Enable generic support for Intel Gaudi accelerator using GPU/HPU migration
-        if kwargs.pop("hpu_migration", True) and is_hpu_available():
+        if device_type == "hpu" and kwargs.pop("hpu_migration", True) and is_hpu_available():
+            os.environ["PT_HPU_GPU_MIGRATION"] = "1"
+            logger.debug("Environment variable set: PT_HPU_GPU_MIGRATION=1")
+
+            import habana_frameworks.torch  # noqa: F401
+
+            # HPU hardware check
+            if not (hasattr(torch, "hpu") and torch.hpu.is_available()):
+                raise ValueError("You are trying to call `.to('hpu')` but HPU device is unavailable.")
+
             os.environ["PT_HPU_MAX_COMPOUND_OP_SIZE"] = "1"
             logger.debug("Environment variable set: PT_HPU_MAX_COMPOUND_OP_SIZE=1")
 

src/diffusers/utils/import_utils.py

@@ -337,19 +337,7 @@ def is_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()
+    return all(importlib.util.find_spec(lib) for lib in ("habana_frameworks", "habana_frameworks.torch"))
 
 
 # docstyle-ignore