Fix Colab and Notebook checks for diffusers-cli env (huggingface#8408)

* chore: Update is_google_colab check to use environment variable

* Check Colab with all possible COLAB_* env variables

* Remove unnecessary word

* Make `_is_google_colab` more inclusive

* Revert "Make `_is_google_colab` more inclusive"

This reverts commit 6406db2.

* Make `_is_google_colab` more inclusive.

* chore: Update import_utils.py with notebook check improvement

* Refactor import_utils.py to improve notebook detection for VS Code's notebook

* chore: Remove `is_notebook()` function and related code

---------

Co-authored-by: Sayak Paul <[email protected]>

Author: tolgacangoz

src/diffusers/commands/env.py

@@ -24,7 +24,6 @@
     is_bitsandbytes_available,
     is_flax_available,
     is_google_colab,
-    is_notebook,
     is_peft_available,
     is_safetensors_available,
     is_torch_available,
@@ -107,8 +106,6 @@ def run(self) -> dict:
 
         platform_info = platform.platform()
 
-        is_notebook_str = "Yes" if is_notebook() else "No"
-
         is_google_colab_str = "Yes" if is_google_colab() else "No"
 
         accelerator = "NA"
@@ -123,7 +120,7 @@ def run(self) -> dict:
                 out_str = out_str.decode("utf-8")
 
                 if len(out_str) > 0:
-                    accelerator = out_str.strip() + " VRAM"
+                    accelerator = out_str.strip()
             except FileNotFoundError:
                 pass
         elif platform.system() == "Darwin":  # Mac OS
@@ -155,7 +152,6 @@ def run(self) -> dict:
         info = {
             "🤗 Diffusers version": version,
             "Platform": platform_info,
-            "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})",

src/diffusers/utils/__init__.py

@@ -73,7 +73,6 @@
     is_librosa_available,
     is_matplotlib_available,
     is_note_seq_available,
-    is_notebook,
     is_onnx_available,
     is_peft_available,
     is_peft_version,

src/diffusers/utils/import_utils.py

@@ -321,18 +321,7 @@ def is_timm_available():
 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
+_is_google_colab = "google.colab" in sys.modules or any(k.startswith("COLAB_") for k in os.environ)
 
 
 def is_torch_available():
@@ -443,10 +432,6 @@ def is_bitsandbytes_available():
     return _bitsandbytes_available
 
 
-def is_notebook():
-    return _is_notebook
-
-
 def is_google_colab():
     return _is_google_colab