File size: 12,741 Bytes
fa257b3
 
992a789
 
 
 
 
 
 
1298d15
39710bb
0ff2c60
 
 
 
 
 
 
992a789
 
 
 
 
 
 
 
0ff2c60
 
6204823
0ff2c60
 
992a789
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
beaef0c
 
 
 
 
 
0ff2c60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
992a789
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c9d5420
d4766e7
c9d5420
992a789
 
 
 
 
1298d15
c9d5420
992a789
6204823
 
 
 
28f726b
992a789
6204823
 
992a789
 
 
 
1298d15
 
 
 
beaef0c
c9d5420
8705203
ef4b87c
992a789
 
 
 
 
 
 
 
 
e0019a7
992a789
 
 
 
 
 
 
 
 
 
 
 
 
6204823
 
 
 
 
992a789
 
6204823
 
 
1298d15
 
 
c9d5420
ef4b87c
 
992a789
 
 
51d3bd9
 
 
 
 
 
 
992a789
8ab9fd5
1298d15
 
 
fb2478e
992a789
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6204823
 
 
fb2478e
 
 
992a789
 
 
6204823
fb2478e
992a789
fb2478e
992a789
 
 
 
6204823
fb2478e
992a789
 
 
 
 
28f726b
992a789
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1298d15
 
992a789
 
 
6204823
 
992a789
 
 
 
 
 
1298d15
992a789
 
 
6c2b3a4
 
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

import gradio as gr
# import gradio.helpers
import torch
import os
from glob import glob
from pathlib import Path
from typing import Optional

import base64
from io import BytesIO
import tempfile
import numpy as np
import cv2
import subprocess

from DeepCache import DeepCacheSDHelper

from PIL import Image
from diffusers.utils import load_image, export_to_video
from pipeline import StableVideoDiffusionPipeline

import random
from safetensors import safe_open
from lcm_scheduler import AnimateLCMSVDStochasticIterativeScheduler

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

hardcoded_fps = 8
hardcoded_duration_sec = 3

def get_safetensors_files():
    models_dir = "./safetensors"
    safetensors_files = [
        f for f in os.listdir(models_dir) if f.endswith(".safetensors")
    ]
    return safetensors_files


def model_select(selected_file):
    print("load model weights", selected_file)
    pipe.unet.cpu()
    file_path = os.path.join("./safetensors", selected_file)
    state_dict = {}
    with safe_open(file_path, framework="pt", device="cpu") as f:
        for key in f.keys():
            state_dict[key] = f.get_tensor(key)
    missing, unexpected = pipe.unet.load_state_dict(state_dict, strict=True)
    pipe.unet.cuda()
    del state_dict
    return

def decode_data_uri_to_image(data_uri):
    # parse the data uri
    header, encoded = data_uri.split(",", 1)
    data = base64.b64decode(encoded)
    img = Image.open(BytesIO(data))
    return img

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

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
        
# ----------------------------------- 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

noise_scheduler = AnimateLCMSVDStochasticIterativeScheduler(
    num_train_timesteps=40,
    sigma_min=0.002,
    sigma_max=700.0,
    sigma_data=1.0,
    s_noise=1.0,
    rho=7,
    clip_denoised=False,
)
pipe = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid-xt",
    scheduler=noise_scheduler,
    torch_dtype=torch.float16,
    variant="fp16",
)
pipe.to("cuda")
pipe.enable_model_cpu_offload()  # for smaller cost
model_select("AnimateLCM-SVD-xt-1.1.safetensors")
# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) # for faster inference


max_64_bit_int = 2**63 - 1

def sample(
    secret_token: str,
    input_image_base64: str,
    seed: Optional[int] = 42,
    randomize_seed: bool = True,
    motion_bucket_id: int = 33,
    desired_duration: int = hardcoded_duration_sec,
    desired_fps: int = hardcoded_fps,
    max_guidance_scale: float = 1.2,
    min_guidance_scale: float = 1,
    width: int = 832,
    height: int = 448,
    num_inference_steps: int = 4,
    decoding_t: int = 4,  # Number of frames decoded at a time! This eats most VRAM. Reduce if necessary.
    output_folder: str = "outputs_gradio",
):
    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.')

    image = decode_data_uri_to_image(input_image_base64)

    print(f"seed={seed}\nrandomize_seed={randomize_seed}\nmotion_bucket_id={motion_bucket_id}\ndesired_duration={desired_duration}\ndesired_fps={desired_fps}\nmax_guidance_scale={max_guidance_scale}\nmin_guidance_scale={min_guidance_scale}\nwidth={width}\nheight={height}\nnum_inference_steps={num_inference_steps}\ndecoding_t={decoding_t}")
    
    if image.mode == "RGBA":
        image = image.convert("RGB")

    if randomize_seed:
        seed = random.randint(0, max_64_bit_int)
    generator = torch.manual_seed(seed)

    os.makedirs(output_folder, exist_ok=True)
    base_count = len(glob(os.path.join(output_folder, "*.mp4")))
    video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")

    with torch.autocast("cuda"):
        frames = pipe(
            image,
            decode_chunk_size=decoding_t,
            generator=generator,
            motion_bucket_id=motion_bucket_id,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            min_guidance_scale=min_guidance_scale,
            max_guidance_scale=max_guidance_scale,
        ).frames[0]

    # we leave default values here
    # alternatively we have implemented our own here: export_to_video_file(...)
    export_to_video(frames, video_path, fps=hardcoded_fps)
    
    torch.manual_seed(seed)

    final_video_path = interpolate_video_frames(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(video_path, "rb") as video_file:
        video_base64 = base64.b64encode(video_file.read()).decode('utf-8')

    # Prepend the appropriate data URI header with MIME type
    return 'data:video/mp4;base64,' + video_base64


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/doevent/AnimateLCM-SVD" target="_blank">original space</a>.</p>
          </div>
          </div>""")
    with gr.Row():
        secret_token = gr.Textbox()
        image_input_base64 = gr.Textbox()
        generate_btn = gr.Button("Generate")
        video_output_base64 = gr.Textbox()

        seed = gr.Slider(
            label="Seed",
            value=42,
            randomize=False,
            minimum=0,
            maximum=max_64_bit_int,
            step=1,
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
        motion_bucket_id = gr.Slider(
            label="Motion bucket id",
            info="Controls how much motion to add/remove from the image",
            value=80,
            minimum=1,
            maximum=255,
        )
        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=5, max_value=60, value=hardcoded_fps, step=1)
    
        # note: we want something that is close to 16:9 (1.7777)
        # 576 / 320 = 1.8
        #  448 / 256 = 1.75
        width = gr.Slider(
            label="Width of input image",
            info="It should be divisible by 64",
            value=832, # 576, # 256, 320, 384, 448
            minimum=256,
            maximum=2048,
            step=64,
        )
        height = gr.Slider(
            label="Height of input image",
            info="It should be divisible by 64",
            value=448, # 320, # 256, 320, 384, 448
            minimum=256,
            maximum=1152,
        )
        max_guidance_scale = gr.Slider(
            label="Max guidance scale",
            info="classifier-free guidance strength",
            value=1.2,
            minimum=1,
            maximum=2,
        )
        min_guidance_scale = gr.Slider(
            label="Min guidance scale",
            info="classifier-free guidance strength",
            value=1,
            minimum=1,
            maximum=1.5,
        )
        num_inference_steps = gr.Slider(
            label="Num inference steps",
            info="steps for inference",
            value=4,
            minimum=1,
            maximum=20,
            step=1,
        )

    generate_btn.click(
        fn=sample,
        inputs=[
            secret_token,
            image_input_base64,
            seed,
            randomize_seed,
            motion_bucket_id,
            duration_slider,
            fps_slider,
            max_guidance_scale,
            min_guidance_scale,
            width,
            height,
            num_inference_steps,
        ],
        outputs=video_output_base64,
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(show_error=True)