streaming_svd_inference.py

import numpy as np
from lib.farancia import IImage
from PIL import Image
from i2v_enhance import i2v_enhance_interface
from dataloader.dataset_factory import SingleImageDatasetFactory
from pytorch_lightning import Trainer, LightningDataModule, seed_everything
import math
from diffusion_trainer import streaming_svd as streaming_svd_model
import torch
from safetensors.torch import load_file as load_safetensors
from utils.loader import download_ckpt
from functools import partial
from dataloader.video_data_module import VideoDataModule
from pathlib import Path
from pytorch_lightning.cli import LightningCLI, LightningArgumentParser
from pytorch_lightning import LightningModule
import sys
import os
from copy import deepcopy
from utils.aux import ensure_annotation_class
from diffusers import FluxPipeline
from typing import Union


class CustomCLI(LightningCLI):

    def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
        parser.add_argument("--image", type=Path,
                            help="Path to the input image(s)")
        parser.add_argument("--output", type=Path,
                            help="Path to the output folder")
        parser.add_argument("--num_frames", type=int, default=100,
                            help="Number of frames to generate.")
        parser.add_argument("--out_fps", type=int, default=24,
                            help="Framerate of the generated video.")
        parser.add_argument("--chunk_size", type=int, default=38,
                            help="Chunk size used in randomized blending.")
        parser.add_argument("--overlap_size", type=int, default=12,
                            help="Overlap size used in randomized blending.")
        parser.add_argument("--use_randomized_blending", action="store_true",
                            help="Wether to use randomized blending.")
        parser.add_argument("--use_fp16", action="store_true",
                            help="Wether to use float16 quantization.")
        parser.add_argument("--prompt", type=str, default = "")

        return parser


class StreamingSVD():

    def __init__(self, load_argv = True) -> None:

        call_fol = Path(os.getcwd()).resolve()

        code_fol = Path(__file__).resolve().parent
        code_fol = os.path.relpath(code_fol, call_fol)
        argv_backup = deepcopy(sys.argv)

        if "--use_fp16" in sys.argv:
            os.environ["STREAMING_USE_FP16"] = "True"
        sys.argv = [__file__]
        sys.argv.extend(self.__config_call(argv_backup[1:] if load_argv else [], code_fol))
        cli = CustomCLI(LightningModule, run=False, subclass_mode_model=True, parser_kwargs={
            "parser_mode": "omegaconf"}, save_config_callback=None)
        self.__init_models(cli)
        self.__init_fields(cli)

        sys.argv = argv_backup

    def __init_models(self, cli):
        model = cli.model
        trainer = cli.trainer

        path = download_ckpt(
            local_path=model.diff_trainer_params.streamingsvd_ckpt.ckpt_path_local,
            global_path=model.diff_trainer_params.streamingsvd_ckpt.ckpt_path_global
        )
        if path.endswith(".safetensors"):
            ckpt = load_safetensors(path)
        else:
            ckpt = torch.load(path, map_location="cpu")["state_dict"]

        model.load_state_dict(ckpt)  # load trained model
        trainer = cli.trainer
        data_module_loader = partial(VideoDataModule, workers=0)
        vfi = i2v_enhance_interface.vfi_init(model.vfi)

        enhance_pipeline, enhance_generator = i2v_enhance_interface.i2v_enhance_init(
            model.i2v_enhance)
        enhance_pipeline.unet.enable_forward_chunking(chunk_size=1, dim=1)
        flux_pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
        flux_pipe.enable_model_cpu_offload()

        # store of objects
        model: streaming_svd_model
        data_module_loader: LightningDataModule
        trainer: Trainer

        self.model = model
        self.vfi = vfi
        self.data_module_loader = data_module_loader
        self.enhance_pipeline = enhance_pipeline
        self.enhance_generator = enhance_generator
        self.trainer = trainer
        self.flux_pipe = flux_pipe

    def __init_fields(self, cli):
        self.input_path = cli.config["image"]
        self.output_path = cli.config["output"]
        self.num_frames = cli.config["num_frames"]
        self.fps = cli.config["out_fps"]
        self.use_randomized_blending = cli.config["use_randomized_blending"]
        self.chunk_size = cli.config["chunk_size"]
        self.overlap_size = cli.config["overlap_size"]
        self.prompt = cli.config["prompt"]

    def __config_call(self, config_cmds, code_fol):
        cmds = [cmd for cmd in config_cmds if len(cmd) > 0]
        cmd_init = []
        cmd_init.append(f"--config")
        cmd_init.append(f"{code_fol}/config.yaml")
        if "--use_fp16" in config_cmds:
            cmd_init.append(f"--trainer.precision=16-true")
        cmd_init.extend(cmds)
        return cmd_init

    # interfaces

    def streaming_t2v(self, prompt, num_frames: int, use_randomized_blending: bool = False, chunk_size: int = 38, overlap_size: int = 12, seed=33):
        image = self.text_to_image(prompt=prompt)
        return self.streaming_i2v(image, num_frames=num_frames, use_randomized_blending=use_randomized_blending, chunk_size=chunk_size, overlap_size=overlap_size, seed=seed)

    def streaming_i2v(self, image, num_frames: int, use_randomized_blending: bool = False, chunk_size: int = 38, overlap_size: int = 12, seed=33) -> np.array:
        video, scaled_outpainted_image, expanded_size = self.image_to_video(
            image, num_frames=(num_frames+1)//2, seed=seed)
        max_memory_allocated = torch.cuda.max_memory_allocated()
        print(
            f"max_memory_allocated at image_to_video: {max_memory_allocated}")
        video = self.enhance_video(image=IImage(scaled_outpainted_image).numpy(), video=video, chunk_size=chunk_size, overlap_size=overlap_size,
                                   use_randomized_blending=use_randomized_blending, seed=seed)
        video = self.interpolate_video(video, dest_num_frames=num_frames)

        # scale/crop back to input size
        if image.shape[0] == 1:
            image = image[0]
        video = IImage(video, vmin=0, vmax=255).resize(expanded_size[::-1]).crop((0, 0, image.shape[1], image.shape[0])).numpy()

        print(
            f"max_memory_allocated at interpolate_video: {max_memory_allocated}")
        return video

    # StreamingSVD pipeline
    def streaming(self, image: np.ndarray):

        datamodule = self.data_module_loader(predict_dataset_factory=SingleImageDatasetFactory(
            file=image))
        self.trainer.predict(model=self.model, datamodule=datamodule)
        video = self.trainer.generated_video
        expanded_size = self.trainer.expanded_size
        scaled_outpainted_image = self.trainer.scaled_outpainted_image
    
        return video, scaled_outpainted_image, expanded_size

    def image_to_video(self, image: Union[np.ndarray, str], num_frames: int, seed=33) -> tuple[np.ndarray,Image,list[int]]:
        seed_everything(seed)
        if isinstance(image, str):
            image = IImage.open(image).numpy()

        if image.shape[0] == 1 and image.ndim == 4:
            image = image[0]

        assert image.shape[-1] == 3 and image.shape[0] > 1, "Wrong image format. Assuming shape [H W C], with C = 3."
        assert image.dtype == "uint8", "Wrong dtype for input image. Must be uint8."
        # compute necessary number of chunks
        n_cond_frames = self.model.inference_params.num_conditional_frames
        n_frames_per_gen = self.model.sampler.guider.num_frames
        n_autoregressive_generations = math.ceil(
            (num_frames - n_frames_per_gen) / (n_frames_per_gen - n_cond_frames))
        self.model.inference_params.n_autoregressive_generations = int(
            n_autoregressive_generations)

        print(" --- STREAMING ----- [START]")
        video, scaled_outpainted_image, expanded_size = self.streaming(
            image=image)
        print(f" --- STREAMING ----- [FINISHED]: {video.shape}")

        video = video[:num_frames]
        return video, scaled_outpainted_image, expanded_size

    def enhance_video(self, video: Union[np.ndarray, str], image: np.ndarray = None, chunk_size = 38, overlap_size=12, strength=0.97, use_randomized_blending=False, seed=33,num_frames = None):

        seed_everything(seed)
        if isinstance(video, str):
            video = IImage.open(video).numpy()
            if image is None:
                image = video[0]
            print("ATTENTION: We take first frame of previous stage as input frame for enhance. ")

        if num_frames is not None:
            video = video[:num_frames, ...]

        if not use_randomized_blending:
            chunk_size = video.shape[0]
            overlap_size = 0
        if image.ndim == 3:
            image = image[None]
        image = [Image.fromarray(
            IImage(image, vmin=0, vmax=255).resize((720, 1280)).numpy()[0])]

        video = np.split(video, video.shape[0])
        video = [Image.fromarray(frame[0]).resize((1280, 720))
                 for frame in video]

        print(
            f"---- ENHANCE  ---- [START]. Video length = {len(video)}. Randomized Blending = {use_randomized_blending}. Chunk size = {chunk_size}. Overlap size = {overlap_size}.")
        video_enhanced = i2v_enhance_interface.i2v_enhance_process(
            image=image, video=video, pipeline=self.enhance_pipeline, generator=self.enhance_generator,
            chunk_size=chunk_size, overlap_size=overlap_size, strength=strength, use_randomized_blending=use_randomized_blending)
        video_enhanced = np.stack([np.asarray(frame)
                                   for frame in video_enhanced], axis=0)
        print("---- ENHANCE  ---- [FINISHED].")
        return video_enhanced

    def interpolate_video(self, video: np.ndarray, dest_num_frames: int):
        video = np.split(video, len(video))
        video = [frame[0] for frame in video]

        print(" ---- VFI  ---- [START]")
        self.vfi.device()
        video_vfi = i2v_enhance_interface.vfi_process(
            video=video, vfi=self.vfi, video_len=dest_num_frames)
        video_vfi = np.stack([np.asarray(frame)
                              for frame in video_vfi], axis=0)
        self.vfi.unload()
        print(f"---- VFI  ---- [FINISHED]. Video length = {len(video_vfi)}")
        return video_vfi

    # T2I method

    def text_to_image(self, prompt, seed=33):
        # FLUX
        print("[FLUX] Generating image from text prompt")
        out = self.flux_pipe(
            prompt=prompt,
            guidance_scale=0,
            height=720,
            width=1280,
            num_inference_steps=4,
            max_sequence_length=256,
            generator=torch.Generator(
                device=self.model.device).manual_seed(seed),
        ).images[0]
        print("[FLUX] Finished")
        return np.array(out)


if __name__ == "__main__":

    @ensure_annotation_class
    def get_input_data(input_path: Path = None):
        if input_path.is_file():
            inputs = [input_path]
        else:
            suffixes = ["*.[jJ][pP][gG]", "*.[pP][nN][gG]",
                        "*.[jJ][pP][eE][gG]", "*.[bB][mM][pP]"]  # loading png, jpg and bmp images
            inputs = []
            for suffix in suffixes:
                inputs.extend(list(input_path.glob(suffix)))
        assert len(
            inputs) > 0, "No images found. Please make sure the input path is correct."

        img_as_np = [IImage.open(input).numpy() for input in inputs]
        return zip(img_as_np, inputs)

    streaming_svd = StreamingSVD()
    num_frames = streaming_svd.num_frames
    chunk_size = streaming_svd.chunk_size
    overlap_size = streaming_svd.overlap_size
    use_randomized_blending = streaming_svd.use_randomized_blending
    if not use_randomized_blending:
        chunk_size = (num_frames + 1)//2
        overlap_size = 0
    result_path = Path(streaming_svd.output_path)
    seed = 33

    assert result_path.exists() is False or result_path.is_dir(
    ), "Output path must be the path to a folder."
    prompt = streaming_svd.prompt
    if len(prompt) == 0:
        for img, img_path in get_input_data(streaming_svd.input_path):
            video = streaming_svd.streaming_i2v(
                image=img, num_frames=num_frames, use_randomized_blending=use_randomized_blending, chunk_size=chunk_size, overlap_size=overlap_size, seed=33)
            if not result_path.exists():
                result_path.mkdir(parents=True)
            result_file = result_path / (img_path.stem+".mp4")
            result_file = result_file.as_posix()
            IImage(video, vmin=0, vmax=255).setFps(
                streaming_svd.fps).save(result_file)
            print(f"Video created at: {result_file}")
    else:
        video = streaming_svd.streaming_t2v(
            prompt=prompt, num_frames=num_frames, use_randomized_blending=use_randomized_blending, chunk_size=chunk_size, overlap_size=overlap_size, seed=33)
        prompt_file = prompt.replace(" ", "_").replace(
            ".", "_").replace("/", "_").replace(":", "_")
        prompt_file = prompt_file[:15]
        if not result_path.exists():
                result_path.mkdir(parents=True)
        result_file = result_path / (prompt_file+".mp4")
        result_file = result_file.as_posix()
        IImage(video, vmin=0, vmax=255).setFps(
            streaming_svd.fps).save(result_file)
        print(f"Video created at: {result_file}")