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