jbilcke-hf's picture
jbilcke-hf HF staff
Update app.py
6204823 verified
raw
history blame
12.7 kB
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}\nfps_id={fps_id}\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)