import gc import os import numpy as np import spaces import gradio as gr import torch from diffusers.training_utils import set_seed from depthcrafter.depth_crafter_ppl import DepthCrafterPipeline from depthcrafter.unet import DiffusersUNetSpatioTemporalConditionModelDepthCrafter import uuid import random from huggingface_hub import hf_hub_download from depthcrafter.utils import read_video_frames, vis_sequence_depth, save_video examples = [ ["examples/example_01.mp4", 10, 1.2, 1024, 60], ["examples/example_02.mp4", 10, 1.2, 1024, 60], ["examples/example_03.mp4", 10, 1.2, 1024, 60], ["examples/example_04.mp4", 10, 1.2, 1024, 60], ["examples/example_05.mp4", 10, 1.2, 1024, 60], ] unet = DiffusersUNetSpatioTemporalConditionModelDepthCrafter.from_pretrained( "tencent/DepthCrafter", low_cpu_mem_usage=True, torch_dtype=torch.float16, ) pipe = DepthCrafterPipeline.from_pretrained( "stabilityai/stable-video-diffusion-img2vid-xt", unet=unet, torch_dtype=torch.float16, variant="fp16", ) pipe.to("cuda") @spaces.GPU(duration=140) def infer_depth( video: str, num_denoising_steps: int, guidance_scale: float, max_res: int = 1024, process_length: int = 195, # save_folder: str = "./demo_output", window_size: int = 110, overlap: int = 25, target_fps: int = 15, seed: int = 42, track_time: bool = True, save_npz: bool = False, ): set_seed(seed) pipe.enable_xformers_memory_efficient_attention() frames, target_fps = read_video_frames(video, process_length, target_fps, max_res) print(f"==> video name: {video}, frames shape: {frames.shape}") # inference the depth map using the DepthCrafter pipeline with torch.inference_mode(): res = pipe( frames, height=frames.shape[1], width=frames.shape[2], output_type="np", guidance_scale=guidance_scale, num_inference_steps=num_denoising_steps, window_size=window_size, overlap=overlap, track_time=track_time, ).frames[0] # convert the three-channel output to a single channel depth map res = res.sum(-1) / res.shape[-1] # normalize the depth map to [0, 1] across the whole video res = (res - res.min()) / (res.max() - res.min()) # visualize the depth map and save the results vis = vis_sequence_depth(res) # save the depth map and visualization with the target FPS save_path = os.path.join(save_folder, os.path.splitext(os.path.basename(video))[0]) os.makedirs(os.path.dirname(save_path), exist_ok=True) if save_npz: np.savez_compressed(save_path + ".npz", depth=res) save_video(res, save_path + "_depth.mp4", fps=target_fps) save_video(vis, save_path + "_vis.mp4", fps=target_fps) save_video(frames, save_path + "_input.mp4", fps=target_fps) # clear the cache for the next video gc.collect() torch.cuda.empty_cache() return [ save_path + "_input.mp4", save_path + "_vis.mp4", # save_path + "_depth.mp4", ] def construct_demo(): with gr.Blocks(analytics_enabled=False) as depthcrafter_iface: gr.Markdown( """

DepthCrafter: Generating Consistent Long Depth Sequences for Open-world Videos

\

\ Wenbo Hu, \ Xiangjun Gao, \ Xiaoyu Li, \ Sijie Zhao, \ Xiaodong Cun, \ Yong Zhang, \ Long Quan, \ Ying Shan\

\ If you find DepthCrafter useful, please help ⭐ the \ [Github Repo]\ , which is important to Open-Source projects. Thanks!\ [ArXiv] \ [Project Page]
""" ) with gr.Row(equal_height=True): with gr.Column(scale=1): input_video = gr.Video(label="Input Video") # with gr.Tab(label="Output"): with gr.Column(scale=2): with gr.Row(equal_height=True): output_video_1 = gr.Video( label="Preprocessed video", interactive=False, autoplay=True, loop=True, show_share_button=True, scale=5, ) output_video_2 = gr.Video( label="Generated Depth Video", interactive=False, autoplay=True, loop=True, show_share_button=True, scale=5, ) with gr.Row(equal_height=True): with gr.Column(scale=1): with gr.Row(equal_height=False): with gr.Accordion("Advanced Settings", open=False): num_denoising_steps = gr.Slider( label="num denoising steps", minimum=1, maximum=25, value=10, step=1, ) guidance_scale = gr.Slider( label="cfg scale", minimum=1.0, maximum=1.2, value=1.2, step=0.1, ) max_res = gr.Slider( label="max resolution", minimum=512, maximum=2048, value=1024, step=64, ) process_length = gr.Slider( label="process length", minimum=1, maximum=280, value=60, step=1, ) generate_btn = gr.Button("Generate") with gr.Column(scale=2): pass gr.Examples( examples=examples, inputs=[ input_video, num_denoising_steps, guidance_scale, max_res, process_length, ], outputs=[output_video_1, output_video_2], fn=infer_depth, cache_examples="lazy", ) gr.Markdown( """ Note: For time quota consideration, we set the default parameters to be more efficient here, with a trade-off of shorter video length and slightly lower quality. You may adjust the parameters according to our [Github Repo] for better results if you have enough time quota. """ ) generate_btn.click( fn=infer_depth, inputs=[ input_video, num_denoising_steps, guidance_scale, max_res, process_length, ], outputs=[output_video_1, output_video_2], ) return depthcrafter_iface demo = construct_demo() if __name__ == "__main__": demo.queue() # demo.launch(server_name="0.0.0.0", server_port=80, debug=True) demo.launch(share=True)