File size: 16,828 Bytes
90ee73b
 
 
3f9fd43
90ee73b
f9a691e
 
 
 
626a382
90ee73b
 
 
 
 
f9a691e
0a535f7
 
90ee73b
f3a1f2e
 
 
 
 
7da1ebd
337bc14
90ee73b
 
f3a1f2e
90ee73b
 
f3a1f2e
 
 
 
 
5d48b65
 
 
 
81d0ed5
 
257cbfe
 
 
 
626a382
 
6def2be
 
 
 
 
 
 
 
 
 
 
 
 
 
 
626a382
257cbfe
 
 
 
7e1bff8
5d48b65
7e1bff8
 
 
 
 
 
 
 
b889cf7
 
 
b37b78f
7e1bff8
 
 
 
 
 
 
 
 
 
 
f9a691e
 
 
 
 
d7766f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
76d7206
 
 
 
 
 
14c6c40
 
c21df80
76d7206
c21df80
f2c1d61
c21df80
 
 
 
14c6c40
 
 
 
c21df80
 
 
 
82a3d77
 
930facd
 
 
 
 
82a3d77
930facd
 
 
e54869f
c21df80
 
 
 
 
 
 
f9a691e
 
e54869f
76d7206
 
14c6c40
76d7206
 
 
 
c21df80
e54869f
 
76d7206
 
82a3d77
 
 
 
 
 
 
 
 
 
 
b5711a4
f9a691e
 
e54869f
f9a691e
 
e54869f
7316f37
f9a691e
0a535f7
 
 
09f95d8
 
5114719
 
1514a70
0a535f7
90ee73b
f3a1f2e
87d5fe9
 
 
f3a1f2e
 
 
 
 
90ee73b
337bc14
 
1514a70
a67daee
1514a70
337bc14
 
7da1ebd
 
61902e5
092446c
 
61902e5
7da1ebd
 
 
61902e5
decd441
 
 
 
f9a691e
decd441
 
 
 
 
5d48b65
decd441
 
 
 
5d48b65
decd441
0a535f7
 
decd441
0a535f7
 
decd441
 
 
 
c5fccfb
decd441
 
 
 
 
0a535f7
7e1bff8
0a535f7
 
90ee73b
 
 
61902e5
0a535f7
 
 
7e1bff8
 
0a535f7
 
 
61902e5
0a535f7
 
90ee73b
 
87d5fe9
61902e5
87d5fe9
337bc14
f3a1f2e
 
 
 
 
 
f9a691e
f3a1f2e
092446c
 
 
 
 
f9a691e
092446c
 
 
 
 
 
14c6c40
092446c
337bc14
27f6e5d
337bc14
27f6e5d
a67daee
 
006c2e8
 
 
 
 
 
337bc14
1514a70
337bc14
f3a1f2e
 
87d5fe9
 
 
 
 
0a535f7
87d5fe9
5d48b65
 
f9a691e
0a535f7
 
 
90ee73b
 
 
f9a691e
0a535f7
90ee73b
 
0e4418a
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
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
import gradio as gr
import torch
import os
import base64
import uuid
import tempfile
import numpy as np
import cv2
import subprocess
from DeepCache import DeepCacheSDHelper
from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from PIL import Image


SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')

# Constants
bases = {
    "ToonYou": "frankjoshua/toonyou_beta6",
    "epiCRealism": "emilianJR/epiCRealism"
}
step_loaded = None
base_loaded = "epiCRealism"
motion_loaded = None

# Ensure model and scheduler are initialized in GPU-enabled function
if not torch.cuda.is_available():
    raise NotImplementedError("No GPU detected!")

device = "cuda"
dtype = torch.float16
pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")

# those are AnimateDiff defaults - we don't touch them for now
hardcoded_fps = 10
hardcoded_duration_sec = 1.6

# unfortunately 2 steps isn't good enough for AiTube, we need 4 steps
step = 4
repo = "ByteDance/AnimateDiff-Lightning"
ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
step_loaded = step

# Note Julian: I'm not sure this works well when the pipeline changes dynamically.. to check
#helper = DeepCacheSDHelper(pipe=pipe)
#helper.set_params(
#    # cache_interval means the frequency of feature caching, specified as the number of steps between each cache operation.
#    # with AnimateDiff this seems to have large effects, so we cannot use large values,
#    # even with cache_interval=3 I notice a big degradation in quality
#    cache_interval=2,
#
#    # cache_branch_id identifies which branch of the network (ordered from the shallowest to the deepest layer) is responsible for executing the caching processes.
#    # Note Julian: I should create my own benchmarks for this
#    cache_branch_id=0,
#
#    # Opting for a lower cache_branch_id or a larger cache_interval can lead to faster inference speed at the expense of reduced image quality
#    #(ablation experiments of these two hyperparameters can be found in the paper).
#)
#helper.enable()

# ----------------------------------- VIDEO ENCODING ---------------------------------
# The Diffusers utils hardcode MP4V as a codec which is not supported by all browsers.
# This is a critical issue for AiTube so we are forced to implement our own routine.
# ------------------------------------------------------------------------------------

def export_to_video_file(video_frames, output_video_path=None, fps=hardcoded_fps):
    if output_video_path is None:
        output_video_path = tempfile.NamedTemporaryFile(suffix=".webm").name

    if isinstance(video_frames[0], np.ndarray):
        video_frames = [(frame * 255).astype(np.uint8) for frame in video_frames]
    elif isinstance(video_frames[0], Image.Image):
        video_frames = [np.array(frame) for frame in video_frames]

    # Use VP9 codec - don't freak out: yes, this will throw an exception, but this still works
    # https://stackoverflow.com/a/61116338
    # I suspect there is a bug somewhere and the actual hex code should be different
    fourcc = cv2.VideoWriter_fourcc(*'VP90')
    h, w, c = video_frames[0].shape
    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h), True)

    for frame in video_frames:
        # Ensure the video frame is in the correct color format
        img = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        video_writer.write(img)
    video_writer.release()

    return output_video_path

# ----------------------------- FRAME INTERPOLATION ---------------------------------
# we cannot afford to use AI-based algorithms such as FILM or ST-MFNet,
# those are way too slow for a AiTube which needs things to be as fast as possible
# -----------------------------------------------------------------------------------

# Convert the video to specified frame rate using motion interpolation.
# 
# This filter accepts the following options:
# 
# fps
# 
#     Specify the output frame rate. This can be rational e.g. 60000/1001. Frames are dropped if fps is lower than source fps. Default 60.
# mi_mode
# 
#     Motion interpolation mode. Following values are accepted:
# 
#     ‘dup’
# 
#         Duplicate previous or next frame for interpolating new ones. 
#     ‘blend’
# 
#         Blend source frames. Interpolated frame is mean of previous and next frames. 
#     ‘mci’
# 
#         Motion compensated interpolation. Following options are effective when this mode is selected:
# 
#         ‘mc_mode’
# 
#             Motion compensation mode. Following values are accepted:
# 
#             ‘obmc’
# 
#                 Overlapped block motion compensation. 
#             ‘aobmc’
# 
#                 Adaptive overlapped block motion compensation. Window weighting coefficients are controlled adaptively according to the reliabilities of the neighboring motion vectors to reduce oversmoothing. 
# 
#             Default mode is ‘obmc’.
#         ‘me_mode’
# 
#             Motion estimation mode. Following values are accepted:
# 
#             ‘bidir’
# 
#                 Bidirectional motion estimation. Motion vectors are estimated for each source frame in both forward and backward directions. 
#             ‘bilat’
# 
#                 Bilateral motion estimation. Motion vectors are estimated directly for interpolated frame. 
# 
#             Default mode is ‘bilat’.
#         ‘me’
# 
#             The algorithm to be used for motion estimation. Following values are accepted:
# 
#             ‘esa’
# 
#                 Exhaustive search algorithm. 
#             ‘tss’
# 
#                 Three step search algorithm. 
#             ‘tdls’
# 
#                 Two dimensional logarithmic search algorithm. 
#             ‘ntss’
# 
#                 New three step search algorithm. 
#             ‘fss’
# 
#                 Four step search algorithm. 
#             ‘ds’
# 
#                 Diamond search algorithm. 
#             ‘hexbs’
# 
#                 Hexagon-based search algorithm. 
#             ‘epzs’
# 
#                 Enhanced predictive zonal search algorithm. 
#             ‘umh’
# 
#                 Uneven multi-hexagon search algorithm. 
# 
#             Default algorithm is ‘epzs’.
#         ‘mb_size’
# 
#             Macroblock size. Default 16.
#         ‘search_param’
# 
#             Motion estimation search parameter. Default 32.
#         ‘vsbmc’
# 
#             Enable variable-size block motion compensation. Motion estimation is applied with smaller block sizes at object boundaries in order to make the them less blur. Default is 0 (disabled). 
# 
# scd
# 
#     Scene change detection method. Scene change leads motion vectors to be in random direction. Scene change detection replace interpolated frames by duplicate ones. May not be needed for other modes. Following values are accepted:
# 
#     ‘none’
# 
#         Disable scene change detection. 
#     ‘fdiff’
# 
#         Frame difference. Corresponding pixel values are compared and if it satisfies scd_threshold scene change is detected. 
# 
#     Default method is ‘fdiff’.
# scd_threshold
# 
#     Scene change detection threshold. Default is 5.0. 

def interpolate_video_frames(
    input_file_path,
    output_file_path,
    output_fps=hardcoded_fps,
    desired_duration=hardcoded_duration_sec,
    original_duration=hardcoded_duration_sec,
    output_width=None,
    output_height=None,
    use_cuda=False, # this requires FFmpeg to have been compiled with CUDA support (to try - I'm not sure the Hugging Face image has that by default)
    verbose=False):
        
    scale_factor = desired_duration / original_duration

    filters = []

    # Scaling if dimensions are provided
    # note: upscaling produces disastrous results,
    # it will double the compute time
    # I think that's either because we are not hardware-accelerated,
    # or because of the interpolation done after it, which thus become more computationally intensive
    if output_width and output_height:
        filters.append(f'scale={output_width}:{output_height}')


    # note: from all fact, it looks like using a small macroblock is important for us,
    # since the video resolution is very small (usually 512x288px)
    interpolation_filter = f'minterpolate=mi_mode=mci:mc_mode=obmc:me=hexbs:vsbmc=1:mb_size=4:fps={output_fps}:scd=none,setpts={scale_factor}*PTS'
    #- `mi_mode=mci`: Specifies motion compensated interpolation.
    #- `mc_mode=obmc`: Overlapped block motion compensation is used.
    #- `me=hexbs`: Hexagon-based search (motion estimation method).
    #- `vsbmc=1`: Variable-size block motion compensation is enabled.
    #- `mb_size=4`: Sets the macroblock size.
    #- `fps={output_fps}`: Defines the output frame rate.
    #- `scd=none`: Disables scene change detection entirely. 
    #- `setpts={scale_factor}*PTS`: Adjusts for the stretching of the video duration.

    # Frame interpolation setup
    filters.append(interpolation_filter)

    # Combine all filters into a single filter complex
    filter_complex = ','.join(filters)


    cmd = [
        'ffmpeg',
        '-i', input_file_path,
    ]

    # not supported by the current image, we will have to build it
    if use_cuda:
        cmd.extend(['-hwaccel', 'cuda', '-hwaccel_output_format', 'cuda'])
               
    cmd.extend([
        '-filter:v', filter_complex,
        '-r', str(output_fps),
        output_file_path
    ])
        
    # Adjust the log level based on the verbosity input
    if not verbose:
        cmd.insert(1, '-loglevel')
        cmd.insert(2, 'error')
    
    # Logging for debugging if verbose
    if verbose:
        print("output_fps:", output_fps)
        print("desired_duration:", desired_duration)
        print("original_duration:", original_duration)
        print("cmd:", cmd)

    try:
        subprocess.run(cmd, check=True)
        return output_file_path
    except subprocess.CalledProcessError as e:
        print("Failed to interpolate video. Error:", e)
        return input_file_path  # In case of error, return original path

def generate_image(secret_token, prompt, base, width, height, motion, step, desired_duration, desired_fps):
    if secret_token != SECRET_TOKEN:
        raise gr.Error(
            f'Invalid secret token. Please fork the original space if you want to use it for yourself.')


    global step_loaded
    global base_loaded
    global motion_loaded
    # print(prompt, base, step)

    if step_loaded != step:
        repo = "ByteDance/AnimateDiff-Lightning"
        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
        step_loaded = step

    if base_loaded != base:
        pipe.unet.load_state_dict(torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device), strict=False)
        base_loaded = base

    if motion_loaded != motion:
        pipe.unload_lora_weights()
        if motion != "":
            pipe.load_lora_weights(motion, adapter_name="motion")
            pipe.set_adapters(["motion"], [0.7])
        motion_loaded = motion

    output = pipe(
        prompt=prompt,

        width=width,
        height=height,
        
        guidance_scale=1.0,
        num_inference_steps=step,
    )
    
    video_uuid = str(uuid.uuid4()).replace("-", "")
    raw_video_path = f"/tmp/{video_uuid}_raw.webm"
    enhanced_video_path = f"/tmp/{video_uuid}_enhanced.webm"
    

    # note the fps is hardcoded, this is a limitation from AnimateDiff I think?
    # (could we change this?)
    #
    # maybe to make things faster, we could *not* encode the video (as this uses files and external processes, which can be slow)
    # and instead return the unencoded frames to the frontend renderer?
    raw_video_path = export_to_video_file(output.frames[0], raw_video_path, fps=hardcoded_fps)

    final_video_path = raw_video_path
    
    # Optional frame interpolation
    if desired_duration > hardcoded_duration_sec or desired_duration < hardcoded_duration_sec or desired_fps > hardcoded_fps or desired_fps < hardcoded_fps:
        final_video_path = interpolate_video_frames(raw_video_path, enhanced_video_path, output_fps=desired_fps, desired_duration=desired_duration)

    # Read the content of the video file and encode it to base64
    with open(final_video_path, "rb") as video_file:
        video_base64 = base64.b64encode(video_file.read()).decode('utf-8')

    # clean-up (otherwise there is always a risk of "ghosting", eg. someone seeing the previous generated video,
    # of one of the steps go wrong - also we need to absolutely delete videos as we generate random files,
    # we can't afford to get a "tmp disk full" error)
    try:
        os.remove(raw_video_path)
        if final_video_path != raw_video_path:
            os.remove(final_video_path)
    except Exception as e:
        print("Failed to delete a video path:", e)
    
    # Prepend the appropriate data URI header with MIME type
    video_data_uri = 'data:video/webm;base64,' + video_base64

    return video_data_uri


# Gradio Interface
with gr.Blocks() as demo:
    gr.HTML("""
        <div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
        <div style="text-align: center; color: black;">
        <p style="color: black;">This space is a headless component of the cloud rendering engine used by AiTube.</p>
        <p style="color: black;">It is not available for public use, but you can use the <a href="https://huggingface.co/spaces/ByteDance/AnimateDiff-Lightning" target="_blank">original space</a>.</p>
        </div>
        </div>""")
    

    secret_token = gr.Text(label='Secret Token', max_lines=1)
    
    with gr.Group():
        with gr.Row():
            prompt = gr.Textbox(
                label='Prompt'
            )
        with gr.Row():
            select_base = gr.Dropdown(
                label='Base model',
                choices=[
                    "ToonYou", 
                    "epiCRealism",
                ],
                value=base_loaded
            )
            width = gr.Slider(
                label='Width',
                minimum=128,
                maximum=2048,
                step=32,
                value=512,
            )
            height = gr.Slider(
                label='Height',
                minimum=128,
                maximum=2048,
                step=32,
                value=288,
            )
            select_motion = gr.Dropdown(
                label='Motion',
                choices=[
                    ("Default", ""),
                    ("Zoom in", "guoyww/animatediff-motion-lora-zoom-in"),
                    ("Zoom out", "guoyww/animatediff-motion-lora-zoom-out"),
                    ("Tilt up", "guoyww/animatediff-motion-lora-tilt-up"),
                    ("Tilt down", "guoyww/animatediff-motion-lora-tilt-down"),
                    ("Pan left", "guoyww/animatediff-motion-lora-pan-left"),
                    ("Pan right", "guoyww/animatediff-motion-lora-pan-right"),
                    ("Roll left", "guoyww/animatediff-motion-lora-rolling-anticlockwise"),
                    ("Roll right", "guoyww/animatediff-motion-lora-rolling-clockwise"),
                ],
                value="",
            )
            select_step = gr.Dropdown(
                label='Inference steps',
                choices=[
                    ('1-Step', 1), 
                    ('2-Step', 2),
                    ('4-Step', 4),
                    ('8-Step', 8)],
                value=4,
            )
            duration_slider = gr.Slider(label="Desired Duration (seconds)", min_value=1, max_value=120, value=hardcoded_duration_sec, step=0.1)
            fps_slider = gr.Slider(label="Desired Frames Per Second", min_value=10, max_value=60, value=hardcoded_fps, step=1)
    
            submit = gr.Button()

    output_video_base64 = gr.Text()

    submit.click(
        fn=generate_image,
        inputs=[secret_token, prompt, select_base, width, height, select_motion, select_step, duration_slider, fps_slider],
        outputs=output_video_base64,
    )

demo.queue(max_size=12).launch(show_api=True)