added rich console patch for Google Colab

This commit is contained in:
k4yt3x 2022-09-04 21:41:57 +00:00
parent 102340e2be
commit db0b87597d

View File

@ -27,7 +27,7 @@ __ __ _ _ ___ __ __
Name: Video2X Name: Video2X
Creator: K4YT3X Creator: K4YT3X
Date Created: February 24, 2018 Date Created: February 24, 2018
Last Modified: August 28, 2022 Last Modified: September 4, 2022
Editor: BrianPetkovsek Editor: BrianPetkovsek
Last Modified: June 17, 2019 Last Modified: June 17, 2019
@ -41,6 +41,7 @@ Last Modified: March 23, 2020
import ctypes import ctypes
import math import math
import os
import signal import signal
import sys import sys
import time import time
@ -164,8 +165,14 @@ class Video2X:
original_stdout = sys.stdout original_stdout = sys.stdout
original_stderr = sys.stderr original_stderr = sys.stderr
# detect if the program is running in Google Colab
# COLAB_GPU should be set in Colab runtime environments
in_colab = os.environ.get("COLAB_GPU") is not None
# create console for rich's Live display # create console for rich's Live display
console = Console() # force_jupyter must be True if the script is running in Colab
# other wise the console will not print properly
console = Console(force_jupyter=True if in_colab is True else None)
# redirect STDOUT and STDERR to console # redirect STDOUT and STDERR to console
sys.stdout = FileProxy(console, sys.stdout) sys.stdout = FileProxy(console, sys.stdout)