app.py

import os

import imageio
import numpy as np
import torch
import random

import spaces

import gradio as gr

import torchvision
import torchvision.transforms as T
from einops import rearrange
from huggingface_hub import hf_hub_download
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights
from torchvision.utils import flow_to_image

from diffusers import AutoencoderKL, MotionAdapter, UNet2DConditionModel
from diffusers import DDIMScheduler
from transformers import CLIPTextModel, CLIPTokenizer

from onlyflow.models.flow_adaptor import FlowEncoder, FlowAdaptor
from onlyflow.models.unet import UNetMotionModel
from onlyflow.pipelines.pipeline_animation_long import FlowCtrlPipeline
from tools.optical_flow import get_optical_flow


def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
    videos = rearrange(videos, "b c t h w -> t b c h w")
    outputs = []
    for x in videos:
        x = torchvision.utils.make_grid(x, nrow=n_rows)
        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
        if rescale:
            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
        x = (x * 255).numpy().astype(np.uint8)
        outputs.append(x)

    os.makedirs(os.path.dirname(path), exist_ok=True)
    imageio.mimsave(path, outputs, fps=fps)

css = """
.toolbutton {
    margin-buttom: 0em 0em 0em 0em;
    max-width: 2.5em;
    min-width: 2.5em !important;
    height: 2.5em;
}
"""


class AnimateController:
    def __init__(self):

        # config dirs
        self.basedir                = os.getcwd()
        self.stable_diffusion_dir   = os.path.join(self.basedir, "models", "StableDiffusion")
        self.motion_module_dir      = os.path.join(self.basedir, "models", "Motion_Module")
        self.personalized_model_dir = os.path.join(self.basedir, "models", "DreamBooth_LoRA")
        self.savedir                = os.path.join(self.basedir, "samples")
        os.makedirs(self.savedir, exist_ok=True)


        ckpt_path = hf_hub_download('obvious-research/onlyflow', 'weights_fp16.ckpt')
        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
        self.flow_encoder_state_dict = ckpt['flow_encoder_state_dict']
        self.attention_processor_state_dict = ckpt['attention_processor_state_dict']

        self.tokenizer             = None
        self.text_encoder          = None
        self.vae                   = None
        self.unet                  = None
        self.motion_adapter        = None

    def update_base_model(self, base_model_id, progress=gr.Progress()):

        progress(0, desc="Starting...")

        self.tokenizer = CLIPTokenizer.from_pretrained(base_model_id, subfolder="tokenizer")
        self.text_encoder = CLIPTextModel.from_pretrained(base_model_id, subfolder="text_encoder")
        self.vae = AutoencoderKL.from_pretrained(base_model_id, subfolder="vae")
        self.unet = UNet2DConditionModel.from_pretrained(base_model_id, subfolder="unet")

        return base_model_id

    def update_motion_module(self, motion_module_id, progress=gr.Progress()):
        self.motion_adapter = MotionAdapter.from_pretrained(motion_module_id)

    def animate(
            self,
            id_base_model,
            id_motion_module,
            prompt_textbox_positive,
            prompt_textbox_negative,
            seed_textbox,
            input_video,
            height,
            width,
            flow_scale,
            cfg,
            diffusion_steps,
            temporal_ds,
            ctx_stride
    ):
        #if any([x is None for x in [self.tokenizer, self.text_encoder, self.vae, self.unet, self.motion_adapter]]) or isinstance(self.unet, str):
        self.update_base_model(id_base_model)
        self.update_motion_module(id_motion_module)

        self.unet = UNetMotionModel.from_unet2d(
            self.unet,
            motion_adapter=self.motion_adapter
        )

        self.raft = raft_large(weights=Raft_Large_Weights.DEFAULT, progress=False).eval()

        self.flow_encoder = FlowEncoder(
            downscale_factor=8,
            channels=[320, 640, 1280, 1280],
            nums_rb=2,
            ksize=1,
            sk=True,
            use_conv=False,
            compression_factor=1,
            temporal_attention_nhead=8,
            positional_embeddings="sinusoidal",
            num_positional_embeddings=16,
            checkpointing=False
        ).eval()

        self.vae.requires_grad_(False)
        self.text_encoder.requires_grad_(False)
        self.unet.requires_grad_(False)
        self.raft.requires_grad_(False)
        self.flow_encoder.requires_grad_(False)

        self.unet.set_all_attn(
            flow_channels=[320, 640, 1280, 1280],
            add_spatial=False,
            add_temporal=True,
            encoder_only=False,
            query_condition=True,
            key_value_condition=True,
            flow_scale=1.0,
        )

        self.flow_adaptor = FlowAdaptor(self.unet, self.flow_encoder).eval()

        # load the flow encoder weights
        pose_enc_m, pose_enc_u = self.flow_adaptor.flow_encoder.load_state_dict(
            self.flow_encoder_state_dict,
            strict=False
        )
        assert len(pose_enc_m) == 0 and len(pose_enc_u) == 0

        # load the attention processor weights
        _, attention_processor_u = self.flow_adaptor.unet.load_state_dict(
            self.attention_processor_state_dict,
            strict=False
        )
        assert len(attention_processor_u) == 0

        pipeline = FlowCtrlPipeline(
            vae=self.vae,
            text_encoder=self.text_encoder,
            tokenizer=self.tokenizer,
            unet=self.unet,
            motion_adapter=self.motion_adapter,
            flow_encoder=self.flow_encoder,
            scheduler=DDIMScheduler.from_pretrained(id_base_model, subfolder="scheduler"),
        )

        if int(seed_textbox) > 0:
            seed = int(seed_textbox)
        else:
            seed = random.randint(1, int(1e16))

        return animate_diffusion(seed, pipeline, self.raft, input_video, prompt_textbox_positive, prompt_textbox_negative, width, height, flow_scale, cfg, diffusion_steps, temporal_ds, ctx_stride)

@spaces.GPU(duration=150)
def animate_diffusion(seed, pipeline, raft_model, base_video, prompt_textbox, negative_prompt_textbox, width_slider, height_slider, flow_scale, cfg, diffusion_steps, temporal_ds, context_stride):
    savedir = './samples'
    device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)

    raft_model = raft_model.to(device)
    pipeline = pipeline.to(device)

    pixel_values = torchvision.io.read_video(base_video, output_format="TCHW", pts_unit='sec')[0][::temporal_ds]
    print("Video loaded, shape:", pixel_values.shape)
    if width_slider/height_slider > pixel_values.shape[3]/pixel_values.shape[2]:
        print("Resizing video to fit width cause input video is not wide enough")
        temp_height = int(width_slider * pixel_values.shape[2]/pixel_values.shape[3])
        temp_width = width_slider
    else:
        print("Resizing video to fit height cause input video is not tall enough")
        temp_height = height_slider
        temp_width = int(height_slider * pixel_values.shape[3]/pixel_values.shape[2])
    print("Resizing video to:", temp_height, temp_width)
    pixel_values = T.Resize((temp_height, temp_width))(pixel_values)
    pixel_values = T.CenterCrop((height_slider, width_slider))(pixel_values)
    pixel_values = T.ConvertImageDtype(torch.float32)(pixel_values)[None, ...].contiguous().to(device)

    save_sample_path_input = os.path.join(savedir, f"input.mp4")
    pixel_values_save = pixel_values[0] * 255
    pixel_values_save = pixel_values_save.cpu()
    pixel_values_save = torch.permute(pixel_values_save, (0, 2, 3, 1))
    torchvision.io.write_video(save_sample_path_input, pixel_values_save, fps=8)
    del pixel_values_save

    print("Video loaded, shape:", pixel_values.shape)
    flow = get_optical_flow(
        raft_model,
        (pixel_values * 2) - 1,
        pixel_values.shape[1] - 1,
        encode_chunk_size=16,
    ).to('cpu')

    sample_flow = (flow_to_image(rearrange(flow[0], "c f h w -> f c h w")))  # N, 3, H, W
    save_sample_path_flow = os.path.join(savedir, f"flow.mp4")
    sample_flow = (sample_flow).cpu().to(torch.uint8).permute(0, 2, 3, 1)
    torchvision.io.write_video(save_sample_path_flow, sample_flow, fps=8)
    del sample_flow

    original_flow_shape = flow.shape
    print("Optical flow computed, shape:", flow.shape)
    if flow.shape[2] < 16:
        print("Video is too short, padding to 16 frames")
        video_length = 16
        n = 16 - flow.shape[2]
        # create a tensor containing the last frame optical flow repeated n times
        to_add = flow[:, :, -1].unsqueeze(2).expand(-1, -1, n, -1, -1)
        flow = torch.cat([flow, to_add], dim=2).to(device)
    elif flow.shape[2] > 16:
        print("Video is too long, enabling windowing")
        print("Enabling model CPU offload")
        pipeline.enable_model_cpu_offload()
        print("Enabling VAE slicing")
        pipeline.enable_vae_slicing()
        print("Enabling VAE tiling")
        pipeline.enable_vae_tiling()

        print("Enabling free noise")
        pipeline.enable_free_noise(
            context_length=16,
            context_stride=context_stride,
        )

        import math

        def find_divisors(n: int):
            """
            Return sorted list of all positive divisors of n.
            Uses a sqrt(n) approach for efficiency.
            """
            divs = set()
            limit = int(math.isqrt(n))
            for i in range(1, limit + 1):
                if n % i == 0:
                    divs.add(i)
                    divs.add(n // i)
            return sorted(divs)

        def multiples_in_range(k: int, min_val: int, max_val: int):
            """
            Return all multiples of k within [min_val, max_val].
            """
            if k == 0:
                return []

            # First multiple of k >= min_val
            start = ((min_val + k - 1) // k) * k
            # Last multiple of k <= max_val
            end = (max_val // k) * k

            return list(range(start, end + 1, k)) if start <= end else []

        def adjust_video_length(original_length: int,
                                context_stride: int,
                                chunk_size: int,
                                temporal_split_size: int) -> int:
            """
            Find the minimal video_length >= original_length satisfying:
              1) (video_length - 16) is divisible by context_stride.
              2) EITHER (2*video_length) is divisible by temporal_split_size
                 OR (2*video_length) is divisible by chunk_size
                 (when 2*video_length is not multiple of temporal_split_size).
            """

            # We start at least at 16 (though in practice original_length likely > 16)
            candidate = max(original_length, 16)

            # We want (candidate - 16) % context_stride == 0
            # so let n be the multiple to step.
            # n is how many times we add `context_stride` beyond 16.
            # This ensures (candidate - 16) is a multiple of context_stride.
            # Then we check the second condition, else keep stepping.

            # If candidate < 16, bump it to 16
            if candidate < 16:
                candidate = 16

            # Make sure we jump to the correct "starting multiple" of context_stride
            offset = (candidate - 16) % context_stride
            if offset != 0:
                candidate += (context_stride - offset)  # jump to the next multiple

            while True:
                # Condition: (candidate - 16) is multiple of context_stride (already enforced by stepping)
                # Check second part:
                # - if (2*candidate) % temporal_split_size == 0, we are good
                # - else we require (2*candidate) % chunk_size == 0
                twoL = 2 * candidate
                if (twoL % temporal_split_size == 0) or (twoL % chunk_size == 0):
                    return candidate

                # Go to next valid candidate
                candidate += context_stride

        def find_valid_configs(original_video_length: int,
                               width: int,
                               height: int,
                               context_stride: int):
            """
            Generate all valid tuples (chunk_size, spatial_split_size, temporal_split_size, video_length)
            subject to the constraints:
              1) chunk_size divides temporal_split_size
              2) chunk_size divides spatial_split_size
              3) chunk_size divides (2 * (width//64) * (height//64))
              4) if (2*video_length) % temporal_split_size != 0, then chunk_size divides (2*video_length)
              5) context_stride divides (video_length - 16)
              6) 128 <= spatial_split_size <= 512
              7) 1 <= temporal_split_size <= 32
              8) 1 <= chunk_size <= 16

            We allow increasing original_video_length minimally if needed to satisfy constraints #4 and #5.
            """

            factor = 2 * (width // 64) * (height // 64)

            # 1) find all possible chunk_size as divisors of factor, in [1..16]
            possible_chunks = [d for d in find_divisors(factor) if 1 <= d <= 32]

            # For storing results
            valid_tuples = []

            for chunk_size in possible_chunks:
                # 2) generate all spatial_split_size in [128..512] that are multiples of chunk_size
                spatial_splits = multiples_in_range(chunk_size, 480, 512)

                # 3) generate all temporal_split_size in [1..32] that are multiples of chunk_size
                temporal_splits = multiples_in_range(chunk_size, 1, 32)

                for ssp in spatial_splits:
                    for tsp in temporal_splits:
                        # 4) & 5) Adjust video_length minimally to satisfy constraints
                        final_length = adjust_video_length(original_video_length,
                                                           context_stride,
                                                           chunk_size,
                                                           tsp)
                        # Now we have a valid (chunk_size, ssp, tsp, final_length)
                        valid_tuples.append((chunk_size, ssp, tsp, final_length))

            return valid_tuples

        def find_pareto_optimal(configs):
            """
            Given a list of tuples (chunk_size, spatial_split_size, temporal_split_size, video_length),
            return the Pareto-optimal subset under the criteria:
              - chunk_size: larger is better
              - spatial_split_size: larger is better
              - temporal_split_size: larger is better
              - video_length: smaller is better
            """

            def dominates(A, B):
                cA, sA, tA, lA = A
                cB, sB, tB, lB = B

                # A dominates B if:
                #   cA >= cB, sA >= sB, tA >= tB, and lA <= lB
                #   AND at least one of these is a strict inequality.

                better_or_equal = (cA >= cB) and (tA >= tB) and (lA <= lB)
                strictly_better = (cA > cB) or (tA > tB) or (lA < lB)

                return better_or_equal and strictly_better

            pareto = []
            for i, cfg_i in enumerate(configs):
                # Check if cfg_i is dominated by any cfg_j
                is_dominated = False
                for j, cfg_j in enumerate(configs):
                    if i == j:
                        continue
                    if dominates(cfg_j, cfg_i):
                        is_dominated = True
                        break
                if not is_dominated:
                    pareto.append(cfg_i)

            return pareto

        print("Finding valid configurations...")
        valid_configs = find_valid_configs(
            original_video_length=flow.shape[2],
            width=width_slider,
            height=height_slider,
            context_stride=context_stride
        )

        print("Found", len(valid_configs), "valid configurations")
        print("Finding Pareto-optimal configurations...")
        pareto_optimal = find_pareto_optimal(valid_configs)

        print("Found", pareto_optimal)

        criteria = lambda cs, sss, tss, vl: cs + tss - 3 * int(abs(flow.shape[2] - vl) / 10)
        pareto_optimal.sort(key=lambda x: criteria(*x), reverse=True)

        print("Found sorted", pareto_optimal)

        solution = pareto_optimal[0]
        chunk_size, spatial_split_size, temporal_split_size, video_length = solution

        n = video_length - original_flow_shape[2]
        to_add = flow[:, :, -1].unsqueeze(2).expand(-1, -1, n, -1, -1)
        flow = torch.cat([flow, to_add], dim=2)

        pipeline.enable_free_noise_split_inference(
            temporal_split_size=temporal_split_size,
            spatial_split_size=spatial_split_size
        )
        pipeline.unet.enable_forward_chunking(chunk_size)

        print("Chunking enabled with chunk size:", chunk_size)
        print("Temporal split size:", temporal_split_size)
        print("Spatial split size:", spatial_split_size)
        print("Context stride:", context_stride)
        print("Temporal downscale:", temporal_ds)
        print("Video length:", video_length)
        print("Flow shape:", flow.shape)
    else:
        print("Video is just right, no padding or windowing needed")
        flow = flow.to(device)
        video_length = flow.shape[2]

    sample_vid = pipeline(
        prompt_textbox,
        negative_prompt=negative_prompt_textbox,
        optical_flow=flow,
        num_inference_steps=diffusion_steps,
        guidance_scale=cfg,
        width=width_slider,
        height=height_slider,
        num_frames=video_length,
        val_scale_factor_temporal=flow_scale,
        generator=generator,
    ).frames[0]

    del flow
    if device == "cuda":
        torch.cuda.synchronize()
        torch.cuda.empty_cache()

    save_sample_path_video = os.path.join(savedir, f"sample.mp4")
    sample_vid = sample_vid[:original_flow_shape[2]] * 255.
    sample_vid = sample_vid.cpu().numpy()
    sample_vid = np.transpose(sample_vid, axes=(0, 2, 3, 1))
    torchvision.io.write_video(save_sample_path_video, sample_vid, fps=8)

    return gr.Video(value=save_sample_path_flow), gr.Video(value=save_sample_path_video)

controller = AnimateController()


def find_closest_ratio(target_ratio):
    width_list = list(reversed(range(256, 1025, 64)))
    height_list = list(reversed(range(256, 1025, 64)))
    ratio_list = [(h, w, w/h) for h in height_list for w in width_list]
    ratio_list.sort(key=lambda x: abs(x[2] - target_ratio))
    ratio_list = list(filter(lambda x: x[2] == ratio_list[0][2], ratio_list))
    ratio_list.sort(key=lambda x: abs(x[0]*x[1] - 512*512))
    return ratio_list[0][:2]


def find_dimension(video):
    import av
    container = av.open(open(video, 'rb'))
    height, width = container.streams.video[0].height, container.streams.video[0].width
    target_ratio = width / height
    return find_closest_ratio(target_ratio)


def ui():
    with gr.Blocks(css=css) as demo:
        gr.Markdown(
            """
            # <p style="text-align:center;">OnlyFlow: Optical Flow based Motion Conditioning for Video Diffusion Models</p>
            Mathis Koroglu, Hugo Caselles-Dupré, Guillaume Jeanneret Sanmiguel, Matthieu Cord<br>
            [Arxiv Report](https://arxiv.org/abs/2411.10501) | [Project Page](https://obvious-research.github.io/onlyflow/) | [Github](https://github.com/obvious-research/onlyflow/)
            """
        )
        gr.Markdown(
            """
            ### Quick Start:

            1. Select desired `Base Model`.
            2. Select `Motion Module`. We recommend trying guoyww/animatediff-motion-adapter-v1-5-3 for the best results.
            3. Provide `Positive Prompt` and `Negative Prompt`. You are encouraged to refer to each model's webpage on HuggingFace Hub or CivitAI to learn how to write prompts for them.
            4. Upload a video to extract optical flow from.
            5. Select a 'Flow Scale' to modulate the input video optical flow conditioning.
            6. Select a 'CFG' and 'Diffusion Steps' to control the quality of the generated video and prompt adherence.
            7. Select a 'Temporal Downsample' to reduce the number of frames in the input video.
            8. If you want to use a custom dimension, check the `Custom Dimension` box and adjust the `Width` and `Height` sliders.
            9. If the video is too long, you can adjust the generation window offset with the `Context Stride` slider.
            10. Click `Generate`, wait for ~1/3 min, and enjoy the result!

            If you have any error concerning GPU limits, please try again later when your ZeroGPU quota is reset, or try with a shorter video.
            Otherwise, you can also duplicate this space and select a custom GPU plan.
            """
        )
        with gr.Row():
            with gr.Column():

                gr.Markdown("# INPUTS")

                with gr.Row(equal_height=True, show_progress=True):
                    base_model = gr.Dropdown(
                        label="Select or type a base model id",
                        choices=[
                        "stable-diffusion-v1-5/stable-diffusion-v1-5",
                        "digiplay/Photon_v1",
                    ],
                        interactive=True,
                        scale=4,
                        allow_custom_value=True,
                        show_label=True
                    )
                    base_model_btn = gr.Button(value="Update", scale=1, size='lg')
                with gr.Row(equal_height=True, show_progress=True):
                    motion_module  = gr.Dropdown(
                        label="Select or type a motion module id",
                        choices=[
                            "guoyww/animatediff-motion-adapter-v1-5-3",
                            "guoyww/animatediff-motion-adapter-v1-5-2"
                        ],
                        interactive=True,
                        scale=4
                    )
                    motion_module_btn = gr.Button(value="Update", scale=1, size='lg')

                base_model_btn.click(fn=controller.update_base_model, inputs=[base_model])
                motion_module_btn.click(fn=controller.update_motion_module, inputs=[motion_module])

                prompt_textbox_positive = gr.Textbox(label="Positive Prompt", lines=3)
                prompt_textbox_negative = gr.Textbox(label="Negative Prompt", lines=2, value="worst quality, low quality, nsfw, logo")

                flow_scale = gr.Slider(label="Flow Scale", value=1.0, minimum=0, maximum=2, step=0.025)
                diffusion_steps = gr.Slider(label="Diffusion Steps", value=25, minimum=0, maximum=100, step=1)
                cfg = gr.Slider(label="CFG", value=7.5, minimum=0, maximum=30, step=0.1)

                temporal_ds = gr.Slider(label="Temporal Downsample", value=1, minimum=1, maximum=30, step=1)

                input_video = gr.Video(label="Input Video", interactive=True)
                ctx_stride = gr.State(12)

                with gr.Accordion("Advanced", open=False):
                    use_custom_dim = gr.Checkbox(label="Custom Dimension", value=False)

                    with gr.Row(equal_height=True):

                        height, width = gr.State(512), gr.State(512)

                        @gr.render(inputs=[use_custom_dim, input_video])
                        def render_custom_dim(use_custom_dim, input_video):
                            if input_video is not None:
                                loc_height, loc_width = find_dimension(input_video)
                            else:
                                loc_height, loc_width = 512, 512
                            slider_width = gr.Slider(label="Width", value=loc_width, minimum=256, maximum=1024,
                                                     step=64, visible=use_custom_dim)
                            slider_height = gr.Slider(label="Height", value=loc_height, minimum=256, maximum=1024,
                                                      step=64, visible=use_custom_dim)

                            slider_width.change(lambda x: x, inputs=[slider_width], outputs=[width])
                            slider_height.change(lambda x: x, inputs=[slider_height], outputs=[height])


                    with gr.Row():
                        @gr.render(inputs=input_video)
                        def render_ctx_stride(input_video):
                            if input_video is not None:
                                video  = open(input_video, 'rb')
                                import av
                                container = av.open(video)
                                num_frames = container.streams.video[0].frames
                                if num_frames > 17:
                                    stride_slider = gr.Slider(label="Context Stride", value=12, minimum=1, maximum=16, step=1)
                                    stride_slider.input(lambda x: x, inputs=[stride_slider], outputs=[ctx_stride])
                                if num_frames > 32:
                                    gr.Warning(f"Video is long ({num_frames} frames), consider using a shorter video, increasing the context stride, or selecting a custom GPU plan.")
                                elif num_frames > 64:
                                    raise gr.Error(f"Video is too long ({num_frames} frames), please use a shorter video, increase the context stride, or select a custom GPU plan. The current parameters won't allow generation on ZeroGPU.")

                    with gr.Row(equal_height=True):
                        seed_textbox = gr.Textbox(label="Seed",  value='-1')

                        seed_button  = gr.Button(value="\U0001F3B2", elem_classes="toolbutton")
                        seed_button.click(
                            fn=lambda: random.randint(1, int(1e16)),
                            inputs=[],
                            outputs=[seed_textbox]
                        )

                with gr.Row():
                    clear_btn = gr.ClearButton(value="Clear & Reset", size='lg', variant='secondary', scale=1)
                    generate_button = gr.Button(value="Generate", variant='primary', scale=2, size='lg')

                    clear_btn.add([base_model, motion_module, input_video, prompt_textbox_positive, prompt_textbox_negative, seed_textbox, use_custom_dim, ctx_stride])

            with gr.Column():

                gr.Markdown("# OUTPUTS")

                result_optical_flow = gr.Video(label="Optical Flow", interactive=False)
                result_video = gr.Video(label="Generated Animation", interactive=False)

            inputs  = [base_model, motion_module, prompt_textbox_positive, prompt_textbox_negative, seed_textbox, input_video, height, width, flow_scale, cfg, diffusion_steps, temporal_ds, ctx_stride]
            outputs = [result_optical_flow, result_video]

            generate_button.click(fn=controller.animate, inputs=inputs, outputs=outputs)

    return demo


if __name__ == "__main__":
    demo = ui()
    demo.queue(max_size=20)
    demo.launch()