test_latte_torch_compile.py

import argparse
import time

import torch

torch.set_float32_matmul_precision("high")

from diffusers import LattePipeline
from diffusers.utils import export_to_gif


def load_pipeline(args):
    id = "maxin-cn/Latte-1"
    pipe = LattePipeline.from_pretrained(
        id,
        torch_dtype=torch.float16
    ).to("cuda:1")
    pipe.set_progress_bar_config(disable=True)

    if args.use_compile:
        torch._inductor.config.conv_1x1_as_mm = True
        torch._inductor.config.coordinate_descent_tuning = True
        torch._inductor.config.epilogue_fusion = False
        torch._inductor.config.coordinate_descent_check_all_directions = True

        pipe.transformer.to(memory_format=torch.channels_last)
        pipe.vae.to(memory_format=torch.channels_last)
        pipe.transformer = torch.compile(pipe.transformer)
        pipe.vae.decode = torch.compile(pipe.vae.decode)
    return pipe


def run_benchmark(args):
    pipe = load_pipeline(args)

    prompt = "a dog wearing an astronaut suit floating in space"
    for _ in range(2):
        _ = pipe(
            prompt=prompt,
            height=512,
            width=512,
            video_length=16,
            num_inference_steps=25,
            generator=torch.manual_seed(42),
        )

    num_repeats = 10
    start = time.time()
    for _ in range(num_repeats):
        video = pipe(
            prompt=prompt,
            height=512,
            width=512,
            video_length=16,
            num_inference_steps=25,
            generator=torch.manual_seed(42),
        ).frames[0]
    end = time.time()
    avg_inference_time = (end - start) / num_repeats
    print(f"Average inference time: {avg_inference_time:.3f} seconds.")

    export_to_gif(video, f"latte_compile_test_{args.use_compile}.gif")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_compile", action="store_true", default=False)
    args = parser.parse_args()
    run_benchmark(args)