File size: 5,684 Bytes
1239b39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
import json
import os
from typing import Optional, Sequence, Tuple

from src.video_util import get_frame_count


class RerenderConfig:

    def __init__(self):
        ...

    def create_from_parameters(self,
                               input_path: str,
                               output_path: str,
                               prompt: str,
                               work_dir: Optional[str] = None,
                               key_subdir: str = 'keys',
                               frame_count: Optional[int] = None,
                               interval: int = 10,
                               crop: Sequence[int] = (0, 0, 0, 0),
                               sd_model: Optional[str] = None,
                               a_prompt: str = '',
                               n_prompt: str = '',
                               ddim_steps=20,
                               scale=7.5,
                               control_type: str = 'HED',
                               control_strength=1,
                               seed: int = -1,
                               image_resolution: int = 512,
                               x0_strength: float = -1,
                               style_update_freq: int = 10,
                               cross_period: Tuple[float, float] = (0, 1),
                               warp_period: Tuple[float, float] = (0, 0.1),
                               mask_period: Tuple[float, float] = (0.5, 0.8),
                               ada_period: Tuple[float, float] = (1.0, 1.0),
                               mask_strength: float = 0.5,
                               inner_strength: float = 0.9,
                               smooth_boundary: bool = True,
                               color_preserve: bool = True,
                               **kwargs):
        self.input_path = input_path
        self.output_path = output_path
        self.prompt = prompt
        self.work_dir = work_dir
        if work_dir is None:
            self.work_dir = os.path.dirname(output_path)
        self.key_dir = os.path.join(self.work_dir, key_subdir)
        self.first_dir = os.path.join(self.work_dir, 'first')

        # Split video into frames
        if not os.path.isfile(input_path):
            raise FileNotFoundError(f'Cannot find video file {input_path}')
        self.input_dir = os.path.join(self.work_dir, 'video')

        self.frame_count = frame_count
        if frame_count is None:
            self.frame_count = get_frame_count(self.input_path)
        self.interval = interval
        self.crop = crop
        self.sd_model = sd_model
        self.a_prompt = a_prompt
        self.n_prompt = n_prompt
        self.ddim_steps = ddim_steps
        self.scale = scale
        self.control_type = control_type
        if self.control_type == 'canny':
            self.canny_low = kwargs.get('canny_low', 100)
            self.canny_high = kwargs.get('canny_high', 200)
        else:
            self.canny_low = None
            self.canny_high = None
        self.control_strength = control_strength
        self.seed = seed
        self.image_resolution = image_resolution
        self.x0_strength = x0_strength
        self.style_update_freq = style_update_freq
        self.cross_period = cross_period
        self.mask_period = mask_period
        self.warp_period = warp_period
        self.ada_period = ada_period
        self.mask_strength = mask_strength
        self.inner_strength = inner_strength
        self.smooth_boundary = smooth_boundary
        self.color_preserve = color_preserve

        os.makedirs(self.input_dir, exist_ok=True)
        os.makedirs(self.work_dir, exist_ok=True)
        os.makedirs(self.key_dir, exist_ok=True)
        os.makedirs(self.first_dir, exist_ok=True)

    def create_from_path(self, cfg_path: str):
        with open(cfg_path, 'r') as fp:
            cfg = json.load(fp)
        kwargs = dict()

        def append_if_not_none(key):
            value = cfg.get(key, None)
            if value is not None:
                kwargs[key] = value

        kwargs['input_path'] = cfg['input']
        kwargs['output_path'] = cfg['output']
        kwargs['prompt'] = cfg['prompt']
        append_if_not_none('work_dir')
        append_if_not_none('key_subdir')
        append_if_not_none('frame_count')
        append_if_not_none('interval')
        append_if_not_none('crop')
        append_if_not_none('sd_model')
        append_if_not_none('a_prompt')
        append_if_not_none('n_prompt')
        append_if_not_none('ddim_steps')
        append_if_not_none('scale')
        append_if_not_none('control_type')
        if kwargs.get('control_type', '') == 'canny':
            append_if_not_none('canny_low')
            append_if_not_none('canny_high')
        append_if_not_none('control_strength')
        append_if_not_none('seed')
        append_if_not_none('image_resolution')
        append_if_not_none('x0_strength')
        append_if_not_none('style_update_freq')
        append_if_not_none('cross_period')
        append_if_not_none('warp_period')
        append_if_not_none('mask_period')
        append_if_not_none('ada_period')
        append_if_not_none('mask_strength')
        append_if_not_none('inner_strength')
        append_if_not_none('smooth_boundary')
        append_if_not_none('color_perserve')
        self.create_from_parameters(**kwargs)

    @property
    def use_warp(self):
        return self.warp_period[0] <= self.warp_period[1]

    @property
    def use_mask(self):
        return self.mask_period[0] <= self.mask_period[1]

    @property
    def use_ada(self):
        return self.ada_period[0] <= self.ada_period[1]