app.py

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

# 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=10):
    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
# -----------------------------------------------------------------------------------

def interpolate_video_frames(input_file_path, output_file_path, output_fps=10, desired_duration=1.6, original_duration=1.6):
    scale_factor = original_duration / desired_duration
    interpolation_filter = f'minterpolate=fps={output_fps},setpts={scale_factor}*PTS'

    cmd = [
        'ffmpeg',
        '-i', input_file_path,
        '-filter:v', interpolation_filter,
        '-r', str(output_fps),
        output_file_path
    ]

    # Logging for debugging
    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=10)

    final_video_path = raw_video_path
    
    # Optional frame interpolation
    if desired_duration != 1 or desired_fps != 10:
        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=256,
            )
            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=1.6, step=0.1)
            fps_slider = gr.Slider(label="Desired Frames Per Second", min_value=10, max_value=60, value=10, 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)