Expansion proposal of `diffusers-cli env` (#7403)

* Expand `diffusers-cli env`

* SafeTensors -> Safetensors
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Move `safetensors_version = "not installed"` to `else`

* Update `safetensors_version` checking

* Add GPU detection for Linux, Mac OS, and Windows

* Add accelerator detection to environment command

* Add is_peft_version to import_utils

* Update env.py

* Add `huggingface_hub` reference

* Add `transformers` reference

* Add reference for `huggingface_hub`

* Fix print statement in env.py for unusual OS

* Up

* Fix platform information in env.py

* up

* Fix import order in env.py

* ruff

* make style

* Fix platform system check in env.py

* Fix run method return type in env.py

* 🤗



* No need f-string

* Remove location info

* Remove accelerate config

* Refactor env.py to remove accelerate config

* feat: Add support for `bitsandbytes` library in environment command

---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>

src/diffusers/commands/env.py

@@ -13,12 +13,24 @@
 # limitations under the License.
 
 import platform
+import subprocess
 from argparse import ArgumentParser
 
 import huggingface_hub
 
 from .. import __version__ as version
-from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
+from ..utils import (
+    is_accelerate_available,
+    is_bitsandbytes_available,
+    is_flax_available,
+    is_google_colab,
+    is_notebook,
+    is_peft_available,
+    is_safetensors_available,
+    is_torch_available,
+    is_transformers_available,
+    is_xformers_available,
+)
 from . import BaseDiffusersCLICommand
 
 
@@ -28,13 +40,19 @@ def info_command_factory(_):
 
 class EnvironmentCommand(BaseDiffusersCLICommand):
     @staticmethod
-    def register_subcommand(parser: ArgumentParser):
+    def register_subcommand(parser: ArgumentParser) -> None:
         download_parser = parser.add_parser("env")
         download_parser.set_defaults(func=info_command_factory)
 
-    def run(self):
+    def run(self) -> dict:
         hub_version = huggingface_hub.__version__
 
+        safetensors_version = "not installed"
+        if is_safetensors_available():
+            import safetensors
+
+            safetensors_version = safetensors.__version__
+
         pt_version = "not installed"
         pt_cuda_available = "NA"
         if is_torch_available():
@@ -43,6 +61,20 @@ class EnvironmentCommand(BaseDiffusersCLICommand):
             pt_version = torch.__version__
             pt_cuda_available = torch.cuda.is_available()
 
+        flax_version = "not installed"
+        jax_version = "not installed"
+        jaxlib_version = "not installed"
+        jax_backend = "NA"
+        if is_flax_available():
+            import flax
+            import jax
+            import jaxlib
+
+            flax_version = flax.__version__
+            jax_version = jax.__version__
+            jaxlib_version = jaxlib.__version__
+            jax_backend = jax.lib.xla_bridge.get_backend().platform
+
         transformers_version = "not installed"
         if is_transformers_available():
             import transformers
@@ -55,21 +87,87 @@ class EnvironmentCommand(BaseDiffusersCLICommand):
 
             accelerate_version = accelerate.__version__
 
+        peft_version = "not installed"
+        if is_peft_available():
+            import peft
+
+            peft_version = peft.__version__
+
+        bitsandbytes_version = "not installed"
+        if is_bitsandbytes_available():
+            import bitsandbytes
+
+            bitsandbytes_version = bitsandbytes.__version__
+
         xformers_version = "not installed"
         if is_xformers_available():
             import xformers
 
             xformers_version = xformers.__version__
 
+        is_notebook_str = "Yes" if is_notebook() else "No"
+
+        is_google_colab_str = "Yes" if is_google_colab() else "No"
+
+        accelerator = "NA"
+        if platform.system() in {"Linux", "Windows"}:
+            try:
+                sp = subprocess.Popen(
+                    ["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader"],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                )
+                out_str, _ = sp.communicate()
+                out_str = out_str.decode("utf-8")
+
+                if len(out_str) > 0:
+                    accelerator = out_str.strip() + " VRAM"
+            except FileNotFoundError:
+                pass
+        elif platform.system() == "Darwin":  # Mac OS
+            try:
+                sp = subprocess.Popen(
+                    ["system_profiler", "SPDisplaysDataType"],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                )
+                out_str, _ = sp.communicate()
+                out_str = out_str.decode("utf-8")
+
+                start = out_str.find("Chipset Model:")
+                if start != -1:
+                    start += len("Chipset Model:")
+                    end = out_str.find("\n", start)
+                    accelerator = out_str[start:end].strip()
+
+                    start = out_str.find("VRAM (Total):")
+                    if start != -1:
+                        start += len("VRAM (Total):")
+                        end = out_str.find("\n", start)
+                        accelerator += " VRAM: " + out_str[start:end].strip()
+            except FileNotFoundError:
+                pass
+        else:
+            print("It seems you are running an unusual OS. Could you fill in the accelerator manually?")
+
         info = {
-            "`diffusers` version": version,
-            "Platform": platform.platform(),
+            "🤗 Diffusers version": version,
+            "Platform": f"{platform.freedesktop_os_release().get('PRETTY_NAME', None)} - {platform.platform()}",
+            "Running on a notebook?": is_notebook_str,
+            "Running on Google Colab?": is_google_colab_str,
             "Python version": platform.python_version(),
             "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
+            "Flax version (CPU?/GPU?/TPU?)": f"{flax_version} ({jax_backend})",
+            "Jax version": jax_version,
+            "JaxLib version": jaxlib_version,
             "Huggingface_hub version": hub_version,
             "Transformers version": transformers_version,
             "Accelerate version": accelerate_version,
+            "PEFT version": peft_version,
+            "Bitsandbytes version": bitsandbytes_version,
+            "Safetensors version": safetensors_version,
             "xFormers version": xformers_version,
+            "Accelerator": accelerator,
             "Using GPU in script?": "<fill in>",
             "Using distributed or parallel set-up in script?": "<fill in>",
         }
@@ -80,5 +178,5 @@ class EnvironmentCommand(BaseDiffusersCLICommand):
         return info
 
     @staticmethod
-    def format_dict(d):
+    def format_dict(d: dict) -> str:
         return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"

src/diffusers/utils/__init__.py

@@ -58,18 +58,22 @@ from .import_utils import (
     get_objects_from_module,
     is_accelerate_available,
     is_accelerate_version,
+    is_bitsandbytes_available,
     is_bs4_available,
     is_flax_available,
     is_ftfy_available,
+    is_google_colab,
     is_inflect_available,
     is_invisible_watermark_available,
     is_k_diffusion_available,
     is_k_diffusion_version,
     is_librosa_available,
     is_note_seq_available,
+    is_notebook,
     is_onnx_available,
     is_peft_available,
     is_peft_version,
+    is_safetensors_available,
     is_scipy_available,
     is_tensorboard_available,
     is_torch_available,

src/diffusers/utils/import_utils.py

@@ -295,6 +295,26 @@ try:
 except importlib_metadata.PackageNotFoundError:
     _torchvision_available = False
 
+_bitsandbytes_available = importlib.util.find_spec("bitsandbytes") is not None
+try:
+    _bitsandbytes_version = importlib_metadata.version("bitsandbytes")
+    logger.debug(f"Successfully imported bitsandbytes version {_bitsandbytes_version}")
+except importlib_metadata.PackageNotFoundError:
+    _bitsandbytes_available = False
+
+# Taken from `huggingface_hub`.
+_is_notebook = False
+try:
+    shell_class = get_ipython().__class__  # type: ignore # noqa: F821
+    for parent_class in shell_class.__mro__:  # e.g. "is subclass of"
+        if parent_class.__name__ == "ZMQInteractiveShell":
+            _is_notebook = True  # Jupyter notebook, Google colab or qtconsole
+            break
+except NameError:
+    pass  # Probably standard Python interpreter
+
+_is_google_colab = "google.colab" in sys.modules
+
 
 def is_torch_available():
     return _torch_available
@@ -392,6 +412,22 @@ def is_torchvision_available():
     return _torchvision_available
 
 
+def is_safetensors_available():
+    return _safetensors_available
+
+
+def is_bitsandbytes_available():
+    return _bitsandbytes_available
+
+
+def is_notebook():
+    return _is_notebook
+
+
+def is_google_colab():
+    return _is_google_colab
+
+
 # docstyle-ignore
 FLAX_IMPORT_ERROR = """
 {0} requires the FLAX library but it was not found in your environment. Checkout the instructions on the
@@ -499,6 +535,20 @@ INVISIBLE_WATERMARK_IMPORT_ERROR = """
 {0} requires the invisible-watermark library but it was not found in your environment. You can install it with pip: `pip install invisible-watermark>=0.2.0`
 """
 
+# docstyle-ignore
+PEFT_IMPORT_ERROR = """
+{0} requires the peft library but it was not found in your environment. You can install it with pip: `pip install peft`
+"""
+
+# docstyle-ignore
+SAFETENSORS_IMPORT_ERROR = """
+{0} requires the safetensors library but it was not found in your environment. You can install it with pip: `pip install safetensors`
+"""
+
+# docstyle-ignore
+BITSANDBYTES_IMPORT_ERROR = """
+{0} requires the bitsandbytes library but it was not found in your environment. You can install it with pip: `pip install bitsandbytes`
+"""
 
 BACKENDS_MAPPING = OrderedDict(
     [
@@ -520,6 +570,9 @@ BACKENDS_MAPPING = OrderedDict(
         ("ftfy", (is_ftfy_available, FTFY_IMPORT_ERROR)),
         ("torchsde", (is_torchsde_available, TORCHSDE_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)),
     ]
 )