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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193 | """
(c) CC BY-SA 4.0, Tremeschin
Advanced example of parallel rendering, with multiple parameters per rendered
file, limiting the number of concurrent processes, thread and vram leaks safe
Warning: This WILL use A LOT OF RAM depending on concurrency and image size
Warning: This WILL use A LOT OF CPU for the video encoding, if enough GPU
• For more information, visit https://brokensrc.dev/depthflow
"""
import itertools
import math
import os
import time
from abc import abstractmethod
from pathlib import Path
from threading import Thread
from typing import List, Self, Type
from attr import Factory, define
from DepthFlow.Animation import Actions, Target
from DepthFlow.Scene import DepthScene
from DepthFlow.State import DepthState
from dotmap import DotMap
from Broken.Externals.Depthmap import BaseEstimator, DepthAnythingV2
from Broken.Externals.Upscaler import NoUpscaler, UpscalerBase, Upscayl
def combinations(**options):
"""Returns a dictionary of key='this' of itertools.product"""
for combination in itertools.product(*options.values()):
yield DotMap(zip(options.keys(), combination))
# Note: You can also use your own subclassing like Custom.py!
class YourScene(DepthScene):
def update(self):
self.state.offset_x = 0.3 * math.sin(self.cycle)
self.state.isometric = 1
# ------------------------------------------------------------------------------------------------ #
@define
class DepthManager:
estimator: BaseEstimator = Factory(DepthAnythingV2)
"""A **shared** estimator for all threads"""
upscaler: UpscalerBase = Factory(NoUpscaler)
"""The upscaler to use for all threads"""
threads: List[Thread] = Factory(list)
"""List of running threads"""
concurrency: int = int(os.getenv("WORKERS", 4))
"""Maximum concurrent render workers (high memory usage)"""
outputs: List[Path] = Factory(list)
"""List of all rendered videos on this session"""
def __attrs_post_init__(self):
self.estimator.load_torch()
self.estimator.load_model()
self.upscaler.load_model()
# # Allow for using with statements
def __enter__(self) -> Self:
self.outputs = list()
return self
def __exit__(self, *ignore) -> None:
self.join()
# # User methods
def parallax(self, scene: Type[DepthScene], image: Path) -> None:
self.estimator.estimate(image)
# Limit the maximum concurrent threads, nice pattern 😉
while len(self.threads) >= self.concurrency:
self.threads = list(filter(lambda x: x.is_alive(), self.threads))
time.sleep(0.05)
# Create and add a new running worker, daemon so it dies with the main thread
thread = Thread(target=self._worker, args=(scene, image), daemon=True)
self.threads.append(thread)
thread.start()
@abstractmethod
def filename(self, data: DotMap) -> Path:
"""Find the output path (Default: same path as image, 'Render' folder)"""
return (data.image.parent / "Render") / ("_".join((
data.image.stem,
f"v{data.variation or 0}",
f"{data.render.time}s",
f"{data.render.height}p{data.render.fps or ''}",
)) + ".mp4")
@abstractmethod
def animate(self, data: DotMap) -> None:
"""Add preset system's animations to each export"""
data.scene.animation.add(Actions.State(
vignette_enable=True,
blur_enable=True,
))
data.scene.animation.add(Actions.Set(target=Target.Isometric, value=0.4))
data.scene.animation.add(Actions.Set(target=Target.Height, value=0.10))
data.scene.animation.add(Actions.Circle(
intensity=0.5,
))
@abstractmethod
def variants(self, image: Path) -> DotMap:
return DotMap(
render=combinations(
height=(1080, 1440),
time=(5, 10),
fps=(60,),
)
)
# # Internal methods
def _worker(self, scene: Type[DepthScene], image: Path):
# Note: Share an estimator between threads to avoid memory leaks
scene = scene(backend="headless")
scene.estimator = self.estimator
scene.set_upscaler(self.upscaler)
scene.input(image=image)
# Note: We reutilize the Scene to avoid re-creation!
# Render multiple lengths, or framerates, anything
for data in combinations(**self.variants(image)):
data.update(scene=scene, image=image)
# Find or set common parameters
output = self.filename(data)
scene.animation.clear()
self.animate(data)
# Make sure the output folder exists
output.parent.mkdir(parents=True, exist_ok=True)
# Render the video
video = scene.main(output=output, **data.render)[0]
self.outputs.append(video)
# Imporant: Free up OpenGL resources
scene.window.destroy()
def join(self):
for thread in self.threads:
thread.join()
# ------------------------------------------------------------------------------------------------ #
# Nice: You can subclass the manager itself 🤯
class YourManager(DepthManager):
def variants(self, image: Path) -> DotMap:
return DotMap(
variation=[0, 1],
render=combinations(
height=[1080],
time=[5],
loop=[2],
fps=[60],
)
)
def animate(self, data: DotMap):
if (data.variation == 0):
data.scene.animation.add(Actions.Orbital())
if (data.variation == 1):
data.scene.animation.add(Actions.Set(target=Target.Isometric, value=0.4))
data.scene.animation.add(Actions.Circle(intensity=0.3))
# ------------------------------------------------------------------------------------------------ #
if (__name__ == "__main__"):
images = Path(os.getenv("IMAGES", "/home/tremeschin/Public/Images"))
# Multiple unique videos per file
# Note: Use Upscayl() for some upscaler!
with DepthManager(upscaler=NoUpscaler()) as manager:
# with YourManager(upscaler=Upscayl()) as manager:
for image in images.glob("*"):
if (image.is_file()):
manager.parallax(DepthScene, image)
for output in manager.outputs:
print(f"• {output}")
|