fix "Stage 1 stops after generating video frame"

2023-05-31 16:02:22 +09:00 · 2023-05-31 16:02:22 +09:00 · 44c18e0ebf
parent 04388c4c19
commit 44c18e0ebf
1 changed files with 3 additions and 3 deletions
--- a/stage1.py
+++ b/stage1.py
@ -147,9 +147,9 @@ def create_mask_clipseg(input_dir, output_dir, clipseg_mask_prompt, clipseg_excl
 def create_mask_transparent_background(input_dir, output_dir, tb_use_fast_mode, tb_use_jit, st1_mask_threshold):
    fast_str = " --fast" if tb_use_fast_mode else ""
    jit_str = " --jit" if tb_use_jit else ""
-    venv = os.environ['VIRTUAL_ENV']
-    if not venv:
-        venv = "venv"
+    venv = "venv"
+    if 'VIRTUAL_ENV' in os.environ:
+        venv = os.environ['VIRTUAL_ENV']
    bin_path = os.path.join(venv, "Scripts")
    bin_path = os.path.join(bin_path, "transparent-background")