diff --git a/reqs.txt b/reqs.txt index b0c986f..9049834 100644 --- a/reqs.txt +++ b/reqs.txt @@ -9,3 +9,4 @@ tensorflow; platform_system != "Windows" facexlib==0.3.0 insightface==0.7.3 basicsr==1.4.2 +protobuf==3.19.6 \ No newline at end of file diff --git a/reqs_portable.txt b/reqs_portable.txt index a29f000..ddec4c0 100644 --- a/reqs_portable.txt +++ b/reqs_portable.txt @@ -12,3 +12,4 @@ qrcode[pil] rembg==2.0.50 tensorflow==2.10.1; platform_system == "Windows" facexlib==0.3.0 +protobuf==3.19.6 \ No newline at end of file diff --git a/interpolate_frames.py b/scripts/interpolate_frames.py similarity index 100% rename from interpolate_frames.py rename to scripts/interpolate_frames.py diff --git a/utils.py b/utils.py index e8aa4d3..d6ab36f 100644 --- a/utils.py +++ b/utils.py @@ -5,6 +5,21 @@ from pathlib import Path import sys from typing import List from .log import log +import signal +from contextlib import suppress +from queue import Queue, Empty +import subprocess +import threading +import os + +# - detect mode +comfy_mode = None +if os.environ.get("COLAB_GPU"): + comfy_mode = "colab" +elif "python_embeded" in sys.executable: + comfy_mode = "embeded" +elif ".venv" in sys.executable: + comfy_mode = "venv" # region MISC Utilities @@ -24,6 +39,79 @@ def add_path(path, prepend=False): sys.path.append(path) +def enqueue_output(out, queue): + for line in iter(out.readline, b""): + queue.put(line) + out.close() + + +def run_command(cmd): + if isinstance(cmd, str): + shell_cmd = cmd + elif isinstance(cmd, list): + shell_cmd = "" + for arg in cmd: + if isinstance(arg, Path): + arg = arg.as_posix() + shell_cmd += f"{arg} " + else: + raise ValueError( + "Invalid 'cmd' argument. It must be a string or a list of arguments." + ) + + process = subprocess.Popen( + shell_cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + shell=True, + ) + + # Create separate threads to read standard output and standard error streams + stdout_queue = Queue() + stderr_queue = Queue() + stdout_thread = threading.Thread( + target=enqueue_output, args=(process.stdout, stdout_queue) + ) + stderr_thread = threading.Thread( + target=enqueue_output, args=(process.stderr, stderr_queue) + ) + stdout_thread.daemon = True + stderr_thread.daemon = True + stdout_thread.start() + stderr_thread.start() + + interrupted = False + + def signal_handler(signum, frame): + nonlocal interrupted + interrupted = True + print("Command execution interrupted.") + + # Register the signal handler for keyboard interrupts (SIGINT) + signal.signal(signal.SIGINT, signal_handler) + + # Process output from both streams until the process completes or interrupted + while not interrupted and ( + process.poll() is None or not stdout_queue.empty() or not stderr_queue.empty() + ): + with suppress(Empty): + stdout_line = stdout_queue.get_nowait() + if stdout_line.strip() != "": + print(stdout_line.strip()) + with suppress(Empty): + stderr_line = stderr_queue.get_nowait() + if stderr_line.strip() != "": + print(stderr_line.strip()) + return_code = process.returncode + + if return_code == 0 and not interrupted: + print("Command executed successfully!") + else: + if not interrupted: + print(f"Command failed with return code: {return_code}") + + # todo use the requirements library reqs_map = { "onnxruntime": "onnxruntime-gpu==1.15.1", @@ -127,6 +215,7 @@ def tensor2np(tensor: torch.Tensor) -> List[np.ndarray]: # endregion + # region MODEL Utilities def download_antelopev2(): antelopev2_url = "https://drive.google.com/uc?id=18wEUfMNohBJ4K3Ly5wpTejPfDzp-8fI8"