From cf55dcf0ff9e7a79150ac0741849e8a7befcbec6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tolga=20Cang=C3=B6z?=
 <46008593+tolgacangoz@users.noreply.github.com>
Date: Tue, 23 Jul 2024 15:34:20 +0300
Subject: [PATCH] Fix Colab and Notebook checks for `diffusers-cli env` (#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 6406db21ac46139c20ccae20ccb44877eec87ba7.

* 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 <spsayakpaul@gmail.com>
---
 src/diffusers/commands/env.py       |  6 +-----
 src/diffusers/utils/__init__.py     |  1 -
 src/diffusers/utils/import_utils.py | 17 +----------------
 3 files changed, 2 insertions(+), 22 deletions(-)

diff --git a/src/diffusers/commands/env.py b/src/diffusers/commands/env.py
index 7a6b598469..d0af30bf1c 100644
--- a/src/diffusers/commands/env.py
+++ b/src/diffusers/commands/env.py
@@ -24,7 +24,6 @@ from ..utils import (
     is_bitsandbytes_available,
     is_flax_available,
     is_google_colab,
-    is_notebook,
     is_peft_available,
     is_safetensors_available,
     is_torch_available,
@@ -107,8 +106,6 @@ class EnvironmentCommand(BaseDiffusersCLICommand):
 
         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 @@ class EnvironmentCommand(BaseDiffusersCLICommand):
                 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 @@ class EnvironmentCommand(BaseDiffusersCLICommand):
         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})",
diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py
index 1612ca5ae4..f41edfcda3 100644
--- a/src/diffusers/utils/__init__.py
+++ b/src/diffusers/utils/__init__.py
@@ -73,7 +73,6 @@ from .import_utils import (
     is_librosa_available,
     is_matplotlib_available,
     is_note_seq_available,
-    is_notebook,
     is_onnx_available,
     is_peft_available,
     is_peft_version,
diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py
index 475d6e415e..44477df2e2 100644
--- a/src/diffusers/utils/import_utils.py
+++ b/src/diffusers/utils/import_utils.py
@@ -321,18 +321,7 @@ try:
 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