From c926010fdaf35c710babd617741c886f6079b034 Mon Sep 17 00:00:00 2001
From: salt <alotofcatz@gmail.com>
Date: Tue, 12 Dec 2023 23:11:45 -0800
Subject: [PATCH] change rerender lora to work with storyteller-rust

---
 animation/rerender-lora/install.py           |  1 +
 animation/rerender-lora/rerender.py          | 14 ++++++++++--
 animation/rerender-lora/src/ddim_v_hacked.py |  3 ++-
 animation/rerender-lora/video_blend.py       | 23 ++++++++++++++++++++
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/animation/rerender-lora/install.py b/animation/rerender-lora/install.py
index 9ae6f29..cab4f6d 100644
--- a/animation/rerender-lora/install.py
+++ b/animation/rerender-lora/install.py
@@ -71,6 +71,7 @@ def download_vae():
     download(url, 'models')
 
 
+print("Before installing, make sure you have initialized the submodules! (git submodule init)")
 build_ebsynth()
 download_gmflow_ckpt()
 download_controlnet_canny()
diff --git a/animation/rerender-lora/rerender.py b/animation/rerender-lora/rerender.py
index e1f56e6..4b72fa0 100644
--- a/animation/rerender-lora/rerender.py
+++ b/animation/rerender-lora/rerender.py
@@ -1,6 +1,8 @@
 import argparse
 import os
 import random
+import subprocess
+import sys
 
 import cv2
 import einops
@@ -402,12 +404,20 @@ def postprocess(cfg: RerenderConfig, ne: bool, max_process: int, tmp: bool,
     use_ps = '-ps' if ps else ''
     o_video_cmd = f'--output {o_video}'
 
+    python_executable = sys.executable
+
     cmd = (
-        f'python video_blend.py {video_base_dir} --beg 1 --end {end_frame} '
+        f'{python_executable} video_blend.py {video_base_dir} --beg 1 --end {end_frame} '
         f'--itv {interval} --key {key_dir} {use_e} {o_video_cmd} --fps {fps} '
         f'--n_proc {max_process} {use_tmp} {use_ps}')
     print(cmd)
-    os.system(cmd)
+    completed_process = subprocess.run(cmd, shell=True, stdout=sys.stdout, stderr=sys.stderr)
+
+    # Check if the command was executed successfully
+    if completed_process.returncode == 0:
+        print("Command executed successfully")
+    else:
+        print(f"Command failed with return code {completed_process.returncode}")
 
 
 if __name__ == '__main__':
diff --git a/animation/rerender-lora/src/ddim_v_hacked.py b/animation/rerender-lora/src/ddim_v_hacked.py
index 3172899..f0561ae 100644
--- a/animation/rerender-lora/src/ddim_v_hacked.py
+++ b/animation/rerender-lora/src/ddim_v_hacked.py
@@ -2,6 +2,7 @@
 
 # CrossAttn precision handling
 import os
+import sys
 
 import einops
 import numpy as np
@@ -290,7 +291,7 @@ class DDIMVSampler(object):
             else timesteps.shape[0]
         print(f'Running DDIM Sampling with {total_steps} timesteps')
 
-        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps, file=sys.stdout)
         if controller is not None:
             controller.set_total_step(total_steps)
         if mask is None:
diff --git a/animation/rerender-lora/video_blend.py b/animation/rerender-lora/video_blend.py
index 78587be..fb7419f 100644
--- a/animation/rerender-lora/video_blend.py
+++ b/animation/rerender-lora/video_blend.py
@@ -4,6 +4,7 @@ import platform
 import struct
 import subprocess
 import time
+from collections import namedtuple
 from typing import List
 
 import cv2
@@ -272,6 +273,28 @@ def main(args):
         video_sequence.remove_out_and_tmp()
 
 
+def process_video(name, output=None, fps=30, beg=1, end=101, itv=10, key='keys0', n_proc=8, ps=False, ne=False, tmp=False):
+    """
+    Process a video with the specified parameters.
+
+    :param name: Path to input video.
+    :param output: Path to output video.
+    :param fps: The FPS of output video.
+    :param beg: The index of the first frame to be stylized.
+    :param end: The index of the last frame to be stylized.
+    :param itv: The interval of key frame.
+    :param key: The subfolder name of stylized key frames.
+    :param n_proc: The max process count.
+    :param ps: Use poisson gradient blending if True.
+    :param ne: Do not run ebsynth if True (use previous ebsynth output).
+    :param tmp: Keep temporary output if True.
+    """
+    Args = namedtuple('Args', ['name', 'output', 'fps', 'beg', 'end', 'itv', 'key', 'n_proc', 'ps', 'ne', 'tmp'])
+    args = Args(name, output, fps, beg, end, itv, key, n_proc, ps, ne, tmp)
+
+    main(args)
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('name', type=str, help='Path to input video')