Set _torch_version to N/A if torch is disabled.

Author: rasmi

src/diffusers/utils/import_utils.py

@@ -37,7 +37,9 @@
 else:
     import importlib.metadata as importlib_metadata
 try:
-    _package_map = importlib_metadata.packages_distributions()  # load-once to avoid expensive calls
+    _package_map = (
+        importlib_metadata.packages_distributions()
+    )  # load-once to avoid expensive calls
 except Exception:
     _package_map = None
 
@@ -53,12 +55,23 @@
 DIFFUSERS_SLOW_IMPORT = os.environ.get("DIFFUSERS_SLOW_IMPORT", "FALSE").upper()
 DIFFUSERS_SLOW_IMPORT = DIFFUSERS_SLOW_IMPORT in ENV_VARS_TRUE_VALUES
 
-STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt}
-
-_is_google_colab = "google.colab" in sys.modules or any(k.startswith("COLAB_") for k in os.environ)
+STR_OPERATION_TO_FUNC = {
+    ">": op.gt,
+    ">=": op.ge,
+    "==": op.eq,
+    "!=": op.ne,
+    "<=": op.le,
+    "<": op.lt,
+}
+
+_is_google_colab = "google.colab" in sys.modules or any(
+    k.startswith("COLAB_") for k in os.environ
+)
 
 
-def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[bool, str]:
+def _is_package_available(
+    pkg_name: str, get_dist_name: bool = False
+) -> Tuple[bool, str]:
     global _package_map
     pkg_exists = importlib.util.find_spec(pkg_name) is not None
     pkg_version = "N/A"
@@ -69,11 +82,16 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
             try:
                 # Fallback for Python < 3.10
                 for dist in importlib_metadata.distributions():
-                    _top_level_declared = (dist.read_text("top_level.txt") or "").split()
+                    _top_level_declared = (
+                        dist.read_text("top_level.txt") or ""
+                    ).split()
                     _infered_opt_names = {
-                        f.parts[0] if len(f.parts) > 1 else inspect.getmodulename(f) for f in (dist.files or [])
+                        f.parts[0] if len(f.parts) > 1 else inspect.getmodulename(f)
+                        for f in (dist.files or [])
                     } - {None}
-                    _top_level_inferred = filter(lambda name: "." not in name, _infered_opt_names)
+                    _top_level_inferred = filter(
+                        lambda name: "." not in name, _infered_opt_names
+                    )
                     for pkg in _top_level_declared or _top_level_inferred:
                         _package_map[pkg].append(dist.metadata["Name"])
             except Exception as _:
@@ -99,16 +117,22 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
 else:
     logger.info("Disabling PyTorch because USE_TORCH is set")
     _torch_available = False
+    _torch_version = "N/A"
 
 _jax_version = "N/A"
 _flax_version = "N/A"
 if USE_JAX in ENV_VARS_TRUE_AND_AUTO_VALUES:
-    _flax_available = importlib.util.find_spec("jax") is not None and importlib.util.find_spec("flax") is not None
+    _flax_available = (
+        importlib.util.find_spec("jax") is not None
+        and importlib.util.find_spec("flax") is not None
+    )
     if _flax_available:
         try:
             _jax_version = importlib_metadata.version("jax")
             _flax_version = importlib_metadata.version("flax")
-            logger.info(f"JAX version {_jax_version}, Flax version {_flax_version} available.")
+            logger.info(
+                f"JAX version {_jax_version}, Flax version {_flax_version} available."
+            )
         except importlib_metadata.PackageNotFoundError:
             _flax_available = False
 else:
@@ -148,7 +172,9 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
             pass
     _onnx_available = _onnxruntime_version is not None
     if _onnx_available:
-        logger.debug(f"Successfully imported onnxruntime version {_onnxruntime_version}")
+        logger.debug(
+            f"Successfully imported onnxruntime version {_onnxruntime_version}"
+        )
 
 # (sayakpaul): importlib.util.find_spec("opencv-python") returns None even when it's installed.
 # _opencv_available = importlib.util.find_spec("opencv-python") is not None
@@ -183,7 +209,9 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
 _invisible_watermark_available = importlib.util.find_spec("imwatermark") is not None
 try:
     _invisible_watermark_version = importlib_metadata.version("invisible-watermark")
-    logger.debug(f"Successfully imported invisible-watermark version {_invisible_watermark_version}")
+    logger.debug(
+        f"Successfully imported invisible-watermark version {_invisible_watermark_version}"
+    )
 except importlib_metadata.PackageNotFoundError:
     _invisible_watermark_available = False
 
@@ -198,7 +226,9 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
 _wandb_available, _wandb_version = _is_package_available("wandb")
 _tensorboard_available, _tensorboard_version = _is_package_available("tensorboard")
 _compel_available, _compel_version = _is_package_available("compel")
-_sentencepiece_available, _sentencepiece_version = _is_package_available("sentencepiece")
+_sentencepiece_available, _sentencepiece_version = _is_package_available(
+    "sentencepiece"
+)
 _torchsde_available, _torchsde_version = _is_package_available("torchsde")
 _peft_available, _peft_version = _is_package_available("peft")
 _torchvision_available, _torchvision_version = _is_package_available("torchvision")
@@ -214,11 +244,19 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> Tuple[b
 _gguf_available, _gguf_version = _is_package_available("gguf")
 _torchao_available, _torchao_version = _is_package_available("torchao")
 _bitsandbytes_available, _bitsandbytes_version = _is_package_available("bitsandbytes")
-_optimum_quanto_available, _optimum_quanto_version = _is_package_available("optimum", get_dist_name=True)
-_pytorch_retinaface_available, _pytorch_retinaface_version = _is_package_available("pytorch_retinaface")
-_better_profanity_available, _better_profanity_version = _is_package_available("better_profanity")
+_optimum_quanto_available, _optimum_quanto_version = _is_package_available(
+    "optimum", get_dist_name=True
+)
+_pytorch_retinaface_available, _pytorch_retinaface_version = _is_package_available(
+    "pytorch_retinaface"
+)
+_better_profanity_available, _better_profanity_version = _is_package_available(
+    "better_profanity"
+)
 _nltk_available, _nltk_version = _is_package_available("nltk")
-_cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available("cosmos_guardrail")
+_cosmos_guardrail_available, _cosmos_guardrail_version = _is_package_available(
+    "cosmos_guardrail"
+)
 
 
 def is_torch_available():
@@ -374,7 +412,10 @@ def is_cosmos_guardrail_available():
 
 
 def is_hpu_available():
-    return all(importlib.util.find_spec(lib) for lib in ("habana_frameworks", "habana_frameworks.torch"))
+    return all(
+        importlib.util.find_spec(lib)
+        for lib in ("habana_frameworks", "habana_frameworks.torch")
+    )
 
 
 # docstyle-ignore
@@ -560,7 +601,10 @@ def is_hpu_available():
         ("compel", (is_compel_available, COMPEL_IMPORT_ERROR)),
         ("ftfy", (is_ftfy_available, FTFY_IMPORT_ERROR)),
         ("torchsde", (is_torchsde_available, TORCHSDE_IMPORT_ERROR)),
-        ("invisible_watermark", (is_invisible_watermark_available, INVISIBLE_WATERMARK_IMPORT_ERROR)),
+        (
+            "invisible_watermark",
+            (is_invisible_watermark_available, INVISIBLE_WATERMARK_IMPORT_ERROR),
+        ),
         ("peft", (is_peft_available, PEFT_IMPORT_ERROR)),
         ("safetensors", (is_safetensors_available, SAFETENSORS_IMPORT_ERROR)),
         ("bitsandbytes", (is_bitsandbytes_available, BITSANDBYTES_IMPORT_ERROR)),
@@ -569,8 +613,14 @@ def is_hpu_available():
         ("gguf", (is_gguf_available, GGUF_IMPORT_ERROR)),
         ("torchao", (is_torchao_available, TORCHAO_IMPORT_ERROR)),
         ("quanto", (is_optimum_quanto_available, QUANTO_IMPORT_ERROR)),
-        ("pytorch_retinaface", (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR)),
-        ("better_profanity", (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR)),
+        (
+            "pytorch_retinaface",
+            (is_pytorch_retinaface_available, PYTORCH_RETINAFACE_IMPORT_ERROR),
+        ),
+        (
+            "better_profanity",
+            (is_better_profanity_available, BETTER_PROFANITY_IMPORT_ERROR),
+        ),
         ("nltk", (is_nltk_available, NLTK_IMPORT_ERROR)),
     ]
 )
@@ -598,9 +648,10 @@ def requires_backends(obj, backends):
             " --upgrade transformers \n```"
         )
 
-    if name in ["StableDiffusionDepth2ImgPipeline", "StableDiffusionPix2PixZeroPipeline"] and is_transformers_version(
-        "<", "4.26.0"
-    ):
+    if name in [
+        "StableDiffusionDepth2ImgPipeline",
+        "StableDiffusionPix2PixZeroPipeline",
+    ] and is_transformers_version("<", "4.26.0"):
         raise ImportError(
             f"You need to install `transformers>=4.26` in order to use {name}: \n```\n pip install"
             " --upgrade transformers \n```"
@@ -620,7 +671,9 @@ def __getattr__(cls, key):
 
 
 # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319
-def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
+def compare_versions(
+    library_or_version: Union[str, Version], operation: str, requirement_version: str
+):
     """
     Compares a library version to some requirement using a given operation.
 
@@ -633,7 +686,9 @@ def compare_versions(library_or_version: Union[str, Version], operation: str, re
             The version to compare the library version against
     """
     if operation not in STR_OPERATION_TO_FUNC.keys():
-        raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}")
+        raise ValueError(
+            f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}"
+        )
     operation = STR_OPERATION_TO_FUNC[operation]
     if isinstance(library_or_version, str):
         library_or_version = parse(importlib_metadata.version(library_or_version))
@@ -837,15 +892,19 @@ class _LazyModule(ModuleType):
 
     # Very heavily inspired by optuna.integration._IntegrationModule
     # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py
-    def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None):
+    def __init__(
+        self, name, module_file, import_structure, module_spec=None, extra_objects=None
+    ):
         super().__init__(name)
         self._modules = set(import_structure.keys())
         self._class_to_module = {}
         for key, values in import_structure.items():
             for value in values:
                 self._class_to_module[value] = key
         # Needed for autocompletion in an IDE
-        self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values()))
+        self.__all__ = list(import_structure.keys()) + list(
+            chain(*import_structure.values())
+        )
         self.__file__ = module_file
         self.__spec__ = module_spec
         self.__path__ = [os.path.dirname(module_file)]