app.py

import gradio as gr
import torch
import os
import random
import time
import math
import spaces
from glob import glob
from pathlib import Path
from typing import Optional, List, Union

from diffusers import StableVideoDiffusionPipeline, StableVideoDragNUWAPipeline
from diffusers.utils import export_to_video, export_to_gif
from PIL import Image

if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

fps25Pipe = StableVideoDiffusionPipeline.from_pretrained(
    "vdo/stable-video-diffusion-img2vid-xt-1-1", torch_dtype=torch.float16, variant="fp16"
)
fps25Pipe.to(device)

fps14Pipe = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid", torch_dtype=torch.float16, variant="fp16"
)
fps14Pipe.to(device)

dragnuwaPipe = StableVideoDragNUWAPipeline.from_pretrained(
    "a-r-r-o-w/dragnuwa-svd", torch_dtype=torch.float16, variant="fp16", low_cpu_mem_usage=False, device_map=None
)
dragnuwaPipe.to(device)

max_64_bit_int = 2**63 - 1

def animate(image, seed=42, randomize_seed=True, motion_bucket_id=127, fps_id=25, noise_aug_strength=0.1, decoding_t=3, video_format="mp4", frame_format="webp", version="auto", width=1024, height=576, motion_control=False, num_inference_steps=25):
    start = time.time()

    if image is None:
        raise gr.Error("Please provide an image to animate.")

    output_folder = "outputs"
    image_data = resize_image(image, output_size=(width, height))
    if image_data.mode == "RGBA":
        image_data = image_data.convert("RGB")

    if motion_control:
        image_data = [image_data] * 3
        
    if randomize_seed:
        seed = random.randint(0, max_64_bit_int)
    
    if version == "auto":
        if 14 < fps_id:
            version = "svdxt"
        else:
            version = "svd"

    frames = animate_on_gpu(
        image_data,
        seed,
        motion_bucket_id,
        fps_id,
        noise_aug_strength,
        decoding_t,
        version,
        width,
        height,
        num_inference_steps
    )
    
    os.makedirs(output_folder, exist_ok=True)
    base_count = len(glob(os.path.join(output_folder, "*." + video_format)))
    result_path = os.path.join(output_folder, f"{base_count:06d}." + video_format)

    if video_format == "gif":
        video_path = None
        gif_path = result_path
        export_to_gif(image=frames, output_gif_path=gif_path, fps=fps_id)
    else:
        video_path = result_path
        gif_path = None
        export_to_video(frames, video_path, fps=fps_id)
    
    end = time.time()
    secondes = int(end - start)
    minutes = math.floor(secondes / 60)
    secondes = secondes - (minutes * 60)
    hours = math.floor(minutes / 60)
    minutes = minutes - (hours * 60)
    information = ("Start the process again if you want a different result. " if randomize_seed else "") + \
    "Wait 2 min before a new run to avoid quota penalty or use another computer. " + \
    "The video has been generated in " + \
    ((str(hours) + " h, ") if hours != 0 else "") + \
    ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + \
    str(secondes) + " sec."
    
    return [
        gr.update(value=video_path, visible=video_format != "gif"),
        gr.update(value=gif_path, visible=video_format == "gif"),
        gr.update(label="💾 Download animation in *." + video_format + " format", value=result_path, visible=True),
        gr.update(label="Generated frames in *." + frame_format + " format", format=frame_format, value=frames, visible=True),
        seed,
        gr.update(value=information, visible=True),
        gr.update(visible=True)
    ]

@torch.no_grad()
@spaces.GPU(queue=False)
def animate_on_gpu(image_data, seed=42, motion_bucket_id=127, fps_id=6, noise_aug_strength=0.1, decoding_t=3, version="svdxt", width=1024, height=576, num_inference_steps=25):
    generator = torch.manual_seed(seed)

    if version == "dragnuwa":
        return dragnuwaPipe(image_data, width=width, height=height, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25, num_inference_steps=num_inference_steps).frames[0]
    elif version == "svdxt":
        return fps25Pipe(image_data, width=width, height=height, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25, num_inference_steps=num_inference_steps).frames[0]
    else:
        return fps14Pipe(image_data, width=width, height=height, decode_chunk_size=decoding_t, generator=generator, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength, num_frames=25, num_inference_steps=num_inference_steps).frames[0]


def resize_image(image, output_size=(1024, 576)):
    if image.width == output_size[0] and image.height == output_size[1]:
        return image

    target_aspect = output_size[0] / output_size[1]
    image_aspect = image.width / image.height

    if image_aspect > target_aspect:
        new_height = output_size[1]
        new_width = int(new_height * image_aspect)
        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
        left = (new_width - output_size[0]) / 2
        top = 0
        right = (new_width + output_size[0]) / 2
        bottom = output_size[1]
    else:
        new_width = output_size[0]
        new_height = int(new_width / image_aspect)
        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
        left = 0
        top = (new_height - output_size[1]) / 2
        right = output_size[0]
        bottom = (new_height + output_size[1]) / 2

    return resized_image.crop((left, top, right, bottom))

def reset():
    return [
        None,
        random.randint(0, max_64_bit_int),
        True,
        127,
        6,
        0.1,
        3,
        "mp4",
        "webp",
        "auto",
        1024,
        576,
        False,
        25
    ]

with gr.Blocks() as demo:
  gr.HTML("""
    <h1><center>Image-to-Video</center></h1>
    <big><center>Animate your image into 25 frames of 1024x576 pixels freely, without account, without watermark and download the video</center></big>
    <br/>
    
    <p>
    This demo is based on <i>Stable Video Diffusion</i> artificial intelligence.
    No prompt or camera control is handled here.
    To control motions, rather use <i><a href="https://huggingface.co/spaces/TencentARC/MotionCtrl_SVD">MotionCtrl SVD</a></i>.
    If you need 128 frames, rather use <i><a href="https://huggingface.co/spaces/modelscope/ExVideo-SVD-128f-v1">ExVideo</a></i>.
    </p>
    """)
  with gr.Row():
      with gr.Column():
          image = gr.Image(label="Upload your image", type="pil")
          with gr.Accordion("Advanced options", open=False):
              width = gr.Slider(label="Width", info="Width of the video", value=1024, minimum=256, maximum=1024, step=8)
              height = gr.Slider(label="Height", info="Height of the video", value=576, minimum=256, maximum=576, step=8)
              motion_control = gr.Checkbox(label="Motion control (experimental)", info="Fix the camera", value=False)
              video_format = gr.Radio([["*.mp4", "mp4"], ["*.avi", "avi"], ["*.wmv", "wmv"], ["*.mkv", "mkv"], ["*.mov", "mov"], ["*.gif", "gif"]], label="Video format for result", info="File extention", value="mp4", interactive=True)
              frame_format = gr.Radio([["*.webp", "webp"], ["*.png", "png"], ["*.jpeg", "jpeg"], ["*.gif (unanimated)", "gif"], ["*.bmp", "bmp"]], label="Image format for frames", info="File extention", value="webp", interactive=True)
              fps_id = gr.Slider(label="Frames per second", info="The length of your video in seconds will be 25/fps", value=25, minimum=5, maximum=30)
              motion_bucket_id = gr.Slider(label="Motion bucket id", info="Controls how much motion to add/remove from the image", value=127, minimum=1, maximum=255)
              noise_aug_strength = gr.Slider(label="Noise strength", info="The noise to add", value=0.1, minimum=0, maximum=1, step=0.1)
              num_inference_steps = gr.Slider(label="Number inference steps", info="More denoising steps usually lead to a higher quality video at the expense of slower inference", value=25, minimum=1, maximum=100, step=1)
              decoding_t = gr.Slider(label="Decoding", info="Number of frames decoded at a time; this eats more VRAM; reduce if necessary", value=3, minimum=1, maximum=5, step=1)
              version = gr.Radio([["Auto", "auto"], ["🏃🏻‍♀️ SVD (trained on 14 f/s)", "svd"], ["🏃🏻‍♀️💨 SVD-XT (trained on 25 f/s)", "svdxt"], ["DragNUWA (unstable)", "dragnuwa"]], label="Model", info="Trained model", value="auto", interactive=True)
              seed = gr.Slider(label="Seed", value=42, randomize=True, minimum=0, maximum=max_64_bit_int, step=1)
              randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

          generate_btn = gr.Button(value="🚀 Animate", variant="primary")
          reset_btn = gr.Button(value="🧹 Reinit page", variant="stop", elem_id="reset_button", visible=False)

      with gr.Column():
          video_output = gr.Video(label="Generated video", format="mp4", autoplay=True, show_download_button=False)
          gif_output = gr.Image(label="Generated video", format="gif", show_download_button=False, visible=False)
          download_button = gr.DownloadButton(label="💾 Download video", visible=False)
          information_msg = gr.HTML(visible=False)
          gallery = gr.Gallery(label="Generated frames", visible=False)
      
  generate_btn.click(fn=animate, inputs=[
      image,
      seed,
      randomize_seed,
      motion_bucket_id,
      fps_id,
      noise_aug_strength,
      decoding_t,
      video_format,
      frame_format,
      version,
      width,
      height,
      motion_control,
      num_inference_steps
  ], outputs=[
      video_output,
      gif_output,
      download_button,
      gallery,
      seed,
      information_msg,
      reset_btn
  ], api_name="video")

  reset_btn.click(fn=reset, inputs=[], outputs=[
      image,
      seed,
      randomize_seed,
      motion_bucket_id,
      fps_id,
      noise_aug_strength,
      decoding_t,
      video_format,
      frame_format,
      version,
      width,
      height,
      motion_control,
      num_inference_steps
  ], queue=False, show_progress=False)
    
  gr.Examples(
    examples=[
        ["Examples/Fire.webp", 42, True, 127, 25, 0.1, 3, "mp4", "png", "auto", 1024, 576, False, 25],
        ["Examples/Water.png", 42, True, 127, 25, 0.1, 3, "mp4", "png", "auto", 1024, 576, False, 25],
        ["Examples/Town.jpeg", 42, True, 127, 25, 0.1, 3, "mp4", "png", "auto", 1024, 576, False, 25]
    ],
    inputs=[image, seed, randomize_seed, motion_bucket_id, fps_id, noise_aug_strength, decoding_t, video_format, frame_format, version, width, height, motion_control, num_inference_steps],
    outputs=[video_output, gif_output, download_button, gallery, seed, information_msg, reset_btn],
    fn=animate,
    run_on_click=True,
    cache_examples=False,
  )

if __name__ == "__main__":
    demo.launch(share=True, show_api=False)