Spaces:
Running
on
Zero
Running
on
Zero
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 | |
import spaces | |
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(prompt, base, width, height, motion, step, desired_duration, desired_fps): | |
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 text-to-video API tool designed for Hugging Chat.</p> | |
<br/> | |
<p style="color: black;">✅ Uses ZeroGPU</p> | |
<p style="color: black;">✅ Simple, uncomplicated workflow</p> | |
<p style="color: black;">✅ Replies in less than 25 seconds</p> | |
<p style="color: black;">✅ Designed to be used as an API</p> | |
<br/> | |
<p style="color: black;">All credit due to the authors of the original space: <a href="https://huggingface.co/spaces/ByteDance/AnimateDiff-Lightning" target="_blank">ByteDance's AnimateDiff-Lightning</a>.</p> | |
</div> | |
</div>""") | |
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)", minimum=1, maximum=120, value=hardcoded_duration_sec, step=0.1) | |
fps_slider = gr.Slider(label="Desired Frames Per Second", minimum=10, maximum=60, value=hardcoded_fps, step=1) | |
submit = gr.Button() | |
output_video = gr.Video() | |
submit.click( | |
fn=generate_image, | |
inputs=[prompt, select_base, width, height, select_motion, select_step, duration_slider, fps_slider], | |
outputs=output_video, | |
) | |
demo.queue().launch(show_api=True) |