add rerender lora

2026-10-26 08:48:03 +00:00 · 2023-12-11 12:50:04 -08:00
parent 0db3bc3be5
commit 0ce2e8ae01
34 changed files with 6442 additions and 0 deletions
@@ -0,0 +1,189 @@
+import os
+import shutil
+
+
+class VideoSequence:
+
+    def __init__(self,
+                 base_dir,
+                 beg_frame,
+                 end_frame,
+                 interval,
+                 input_subdir='videos',
+                 key_subdir='keys0',
+                 tmp_subdir='tmp',
+                 input_format='frame%04d.jpg',
+                 key_format='%04d.jpg',
+                 out_subdir_format='out_%d',
+                 blending_out_subdir='blend',
+                 output_format='%04d.jpg'):
+        if (end_frame - beg_frame) % interval != 0:
+            end_frame -= (end_frame - beg_frame) % interval
+
+        self.__base_dir = base_dir
+        self.__input_dir = os.path.join(base_dir, input_subdir)
+        self.__key_dir = os.path.join(base_dir, key_subdir)
+        self.__tmp_dir = os.path.join(base_dir, tmp_subdir)
+        self.__input_format = input_format
+        self.__blending_out_dir = os.path.join(base_dir, blending_out_subdir)
+        self.__key_format = key_format
+        self.__out_subdir_format = out_subdir_format
+        self.__output_format = output_format
+        self.__beg_frame = beg_frame
+        self.__end_frame = end_frame
+        self.__interval = interval
+        self.__n_seq = (end_frame - beg_frame) // interval
+        self.__make_out_dirs()
+        os.makedirs(self.__tmp_dir, exist_ok=True)
+
+    @property
+    def beg_frame(self):
+        return self.__beg_frame
+
+    @property
+    def end_frame(self):
+        return self.__end_frame
+
+    @property
+    def n_seq(self):
+        return self.__n_seq
+
+    @property
+    def interval(self):
+        return self.__interval
+
+    @property
+    def blending_dir(self):
+        return os.path.abspath(self.__blending_out_dir)
+
+    def remove_out_and_tmp(self):
+        for i in range(self.n_seq + 1):
+            out_dir = self.__get_out_subdir(i)
+            shutil.rmtree(out_dir)
+        shutil.rmtree(self.__tmp_dir)
+
+    def get_input_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            id_list = list(range(end_id, beg_id, -1))
+        path_dir = [
+            os.path.join(self.__input_dir, self.__input_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_output_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        out_subdir = self.__get_out_subdir(i)
+        path_dir = [
+            os.path.join(out_subdir, self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_temporal_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        tmp_dir = self.__get_tmp_out_subdir(i)
+        path_dir = [
+            os.path.join(tmp_dir, 'temporal_' + self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_edge_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        tmp_dir = self.__get_tmp_out_subdir(i)
+        path_dir = [
+            os.path.join(tmp_dir, 'edge_' + self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_pos_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id))
+        else:
+            i += 1
+            id_list = list(range(end_id, beg_id, -1))
+        tmp_dir = self.__get_tmp_out_subdir(i)
+        path_dir = [
+            os.path.join(tmp_dir, 'pos_' + self.__output_format % id)
+            for id in id_list
+        ]
+        return path_dir
+
+    def get_flow_sequence(self, i, is_forward=True):
+        beg_id = self.get_sequence_beg_id(i)
+        end_id = self.get_sequence_beg_id(i + 1)
+        if is_forward:
+            id_list = list(range(beg_id, end_id - 1))
+            path_dir = [
+                os.path.join(self.__tmp_dir, 'flow_f_%04d.npy' % id)
+                for id in id_list
+            ]
+        else:
+            id_list = list(range(end_id, beg_id + 1, -1))
+            path_dir = [
+                os.path.join(self.__tmp_dir, 'flow_b_%04d.npy' % id)
+                for id in id_list
+            ]
+
+        return path_dir
+
+    def get_input_img(self, i):
+        return os.path.join(self.__input_dir, self.__input_format % i)
+
+    def get_key_img(self, i):
+        sequence_beg_id = self.get_sequence_beg_id(i)
+        return os.path.join(self.__key_dir,
+                            self.__key_format % sequence_beg_id)
+
+    def get_blending_img(self, i):
+        return os.path.join(self.__blending_out_dir, self.__output_format % i)
+
+    def get_sequence_beg_id(self, i):
+        return i * self.__interval + self.__beg_frame
+
+    def __get_out_subdir(self, i):
+        dir_id = self.get_sequence_beg_id(i)
+        out_subdir = os.path.join(self.__base_dir,
+                                  self.__out_subdir_format % dir_id)
+        return out_subdir
+
+    def __get_tmp_out_subdir(self, i):
+        dir_id = self.get_sequence_beg_id(i)
+        tmp_out_subdir = os.path.join(self.__tmp_dir,
+                                      self.__out_subdir_format % dir_id)
+        return tmp_out_subdir
+
+    def __make_out_dirs(self):
+        os.makedirs(self.__base_dir, exist_ok=True)
+        os.makedirs(self.__blending_out_dir, exist_ok=True)
+        for i in range(self.__n_seq + 1):
+            out_subdir = self.__get_out_subdir(i)
+            tmp_subdir = self.__get_tmp_out_subdir(i)
+            os.makedirs(out_subdir, exist_ok=True)
+            os.makedirs(tmp_subdir, exist_ok=True)