videox_fun/ui/controller.py

"""Modified from https://github.com/guoyww/AnimateDiff/blob/main/app.py
"""
import base64
import gc
import json
import os
import hashlib
import random
from datetime import datetime
from glob import glob

import cv2
import gradio as gr
import numpy as np
import pkg_resources
import requests
import torch
from diffusers import (CogVideoXDDIMScheduler, DDIMScheduler,
                       DPMSolverMultistepScheduler,
                       EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
                       FlowMatchEulerDiscreteScheduler, PNDMScheduler)
from omegaconf import OmegaConf
from PIL import Image
from safetensors import safe_open

from ..data.bucket_sampler import ASPECT_RATIO_512, get_closest_ratio
from ..utils.utils import save_videos_grid
from ..utils.fm_solvers import FlowDPMSolverMultistepScheduler
from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from ..dist import set_multi_gpus_devices

gradio_version = pkg_resources.get_distribution("gradio").version
gradio_version_is_above_4 = True if int(gradio_version.split('.')[0]) >= 4 else False

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

ddpm_scheduler_dict = {
    "Euler": EulerDiscreteScheduler,
    "Euler A": EulerAncestralDiscreteScheduler,
    "DPM++": DPMSolverMultistepScheduler, 
    "PNDM": PNDMScheduler,
    "DDIM": DDIMScheduler,
    "DDIM_Origin": DDIMScheduler,
    "DDIM_Cog": CogVideoXDDIMScheduler,
}
flow_scheduler_dict = {
    "Flow": FlowMatchEulerDiscreteScheduler,
    "Flow_Unipc": FlowUniPCMultistepScheduler,
    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
}
all_cheduler_dict = {**ddpm_scheduler_dict, **flow_scheduler_dict}

class Fun_Controller:
    def __init__(
        self, GPU_memory_mode, scheduler_dict, model_name=None, model_type="Inpaint", 
        config_path=None, ulysses_degree=1, ring_degree=1, 
        fsdp_dit=False, fsdp_text_encoder=False, compile_dit=False, 
        weight_dtype=None, savedir_sample=None,
    ):
        # config dirs
        self.basedir                    = os.getcwd()
        self.config_dir                 = os.path.join(self.basedir, "config")
        self.diffusion_transformer_dir  = os.path.join(self.basedir, "models", "Diffusion_Transformer")
        self.motion_module_dir          = os.path.join(self.basedir, "models", "Motion_Module")
        self.personalized_model_dir     = os.path.join(self.basedir, "models", "Personalized_Model")
        if savedir_sample is None:
            self.savedir_sample         = os.path.join(self.basedir, "samples", datetime.now().strftime("Gradio-%Y-%m-%dT%H-%M-%S"))
        else:
            self.savedir_sample         = savedir_sample
        os.makedirs(self.savedir_sample, exist_ok=True)

        self.GPU_memory_mode                = GPU_memory_mode
        self.model_name                     = model_name
        self.diffusion_transformer_dropdown = model_name
        self.scheduler_dict             = scheduler_dict
        self.model_type                 = model_type
        if config_path is not None:
            self.config_path            = os.path.realpath(config_path)
            self.config                 = OmegaConf.load(config_path)
        else:
            self.config_path            = None
        self.ulysses_degree             = ulysses_degree
        self.ring_degree                = ring_degree
        self.fsdp_dit                   = fsdp_dit       
        self.fsdp_text_encoder          = fsdp_text_encoder
        self.compile_dit                = compile_dit
        self.weight_dtype               = weight_dtype
        self.device                     = set_multi_gpus_devices(self.ulysses_degree, self.ring_degree)

        self.diffusion_transformer_list = []
        self.motion_module_list         = []
        self.personalized_model_list    = []
        self.config_list                = []

        # config models
        self.tokenizer             = None
        self.text_encoder          = None
        self.vae                   = None
        self.transformer           = None
        self.transformer_2         = None
        self.pipeline              = None
        self.base_model_path       = "none"
        self.base_model_2_path     = "none"
        self.lora_model_path       = "none"
        self.lora_model_2_path     = "none"
        
        self.refresh_config()
        self.refresh_diffusion_transformer()
        self.refresh_personalized_model()
        if model_name != None:
            self.update_diffusion_transformer(model_name)

    def refresh_config(self):
        config_list = []
        for root, dirs, files in os.walk(self.config_dir):
            for file in files:
                if file.endswith(('.yaml', '.yml')):
                    full_path = os.path.join(root, file)
                    config_list.append(full_path)
        self.config_list = config_list

    def refresh_diffusion_transformer(self):
        self.diffusion_transformer_list = sorted(glob(os.path.join(self.diffusion_transformer_dir, "*/")))

    def refresh_personalized_model(self):
        personalized_model_list = sorted(glob(os.path.join(self.personalized_model_dir, "*.safetensors")))
        self.personalized_model_list = [os.path.basename(p) for p in personalized_model_list]

    def update_model_type(self, model_type):
        self.model_type = model_type

    def update_config(self, config_dropdown):
        self.config_path = config_dropdown
        self.config = OmegaConf.load(config_dropdown)
        print(f"Update config: {config_dropdown}")

    def update_diffusion_transformer(self, diffusion_transformer_dropdown):
        pass

    def update_base_model(self, base_model_dropdown, is_checkpoint_2=False):
        if not is_checkpoint_2:
            self.base_model_path = base_model_dropdown
        else:
            self.base_model_2_path = base_model_dropdown
        print(f"Update base model: {base_model_dropdown}")
        if base_model_dropdown == "none":
            return gr.update()
        if self.transformer is None and not is_checkpoint_2:
            gr.Info(f"Please select a pretrained model path.")
            print(f"Please select a pretrained model path.")
            return gr.update(value=None)
        elif self.transformer_2 is None and is_checkpoint_2:
            gr.Info(f"Please select a pretrained model path.")
            print(f"Please select a pretrained model path.")
            return gr.update(value=None)
        else:
            base_model_dropdown = os.path.join(self.personalized_model_dir, base_model_dropdown)
            base_model_state_dict = {}
            with safe_open(base_model_dropdown, framework="pt", device="cpu") as f:
                for key in f.keys():
                    base_model_state_dict[key] = f.get_tensor(key)
            if not is_checkpoint_2:
                self.transformer.load_state_dict(base_model_state_dict, strict=False)
            else:
                self.transformer_2.load_state_dict(base_model_state_dict, strict=False)
            print("Update base model done")
            return gr.update()

    def update_lora_model(self, lora_model_dropdown, is_checkpoint_2=False):
        print(f"Update lora model: {lora_model_dropdown}")
        if lora_model_dropdown == "none":
            self.lora_model_path = "none"
            return gr.update()
        lora_model_dropdown = os.path.join(self.personalized_model_dir, lora_model_dropdown)
        if not is_checkpoint_2:
            self.lora_model_path = lora_model_dropdown
        else:
            self.lora_model_2_path = lora_model_dropdown
        return gr.update()

    def clear_cache(self,):
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()

    def auto_model_clear_cache(self, model):
        origin_device = model.device
        model = model.to("cpu")
        gc.collect()
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        model = model.to(origin_device)
    
    def input_check(self,
        resize_method,
        generation_method, 
        start_image, 
        end_image, 
        validation_video,
        control_video,
        is_api = False,
    ):
        if self.transformer is None:
            if is_api:
                return "", f"Please select a pretrained model path."
            else:
                raise gr.Error(f"Please select a pretrained model path.")
        
        if control_video is not None and self.model_type == "Inpaint":
            if is_api:
                return "", f"If specifying the control video, please set the model_type == \"Control\". "
            else:
                raise gr.Error(f"If specifying the control video, please set the model_type == \"Control\". ")

        if control_video is None and self.model_type == "Control":
            if is_api:
                return "", f"If set the model_type == \"Control\", please specifying the control video. "
            else:
                raise gr.Error(f"If set the model_type == \"Control\", please specifying the control video. ")

        if resize_method == "Resize according to Reference":
            if start_image is None and validation_video is None and control_video is None:
                if is_api:
                    return "", f"Please upload an image when using \"Resize according to Reference\"."
                else:
                    raise gr.Error(f"Please upload an image when using \"Resize according to Reference\".")

        if self.transformer.config.in_channels == self.vae.config.latent_channels and start_image is not None:
            if is_api:
                return "", f"Please select an image to video pretrained model while using image to video."
            else:
                raise gr.Error(f"Please select an image to video pretrained model while using image to video.")

        if self.transformer.config.in_channels == self.vae.config.latent_channels and generation_method == "Long Video Generation":
            if is_api:
                return "", f"Please select an image to video pretrained model while using long video generation."
            else:
                raise gr.Error(f"Please select an image to video pretrained model while using long video generation.")
        
        if start_image is None and end_image is not None:
            if is_api:
                return "", f"If specifying the ending image of the video, please specify a starting image of the video."
            else:
                raise gr.Error(f"If specifying the ending image of the video, please specify a starting image of the video.")
        return "", "OK"

    def get_height_width_from_reference(
        self,
        base_resolution,
        start_image,
        validation_video,
        control_video,
    ):
        spatial_compression_ratio = self.vae.config.spatial_compression_ratio if hasattr(self.vae.config, "spatial_compression_ratio") else 8
        aspect_ratio_sample_size    = {key : [x / 512 * base_resolution for x in ASPECT_RATIO_512[key]] for key in ASPECT_RATIO_512.keys()}
        if self.model_type == "Inpaint":
            if validation_video is not None:
                original_width, original_height = Image.fromarray(cv2.VideoCapture(validation_video).read()[1]).size
            else:
                original_width, original_height = start_image[0].size if type(start_image) is list else Image.open(start_image).size
        else:
            original_width, original_height = Image.fromarray(cv2.VideoCapture(control_video).read()[1]).size
        closest_size, closest_ratio = get_closest_ratio(original_height, original_width, ratios=aspect_ratio_sample_size)
        height_slider, width_slider = [int(x / spatial_compression_ratio / 2) * spatial_compression_ratio * 2 for x in closest_size]
        return height_slider, width_slider

    def save_outputs(self, is_image, length_slider, sample, fps):
        def save_results():
            if not os.path.exists(self.savedir_sample):
                os.makedirs(self.savedir_sample, exist_ok=True)
            index = len([path for path in os.listdir(self.savedir_sample)]) + 1
            prefix = str(index).zfill(8)

            md5_hash = hashlib.md5(sample.cpu().numpy().tobytes()).hexdigest()

            if is_image or length_slider == 1:
                save_sample_path = os.path.join(self.savedir_sample, prefix + f"-{md5_hash}.png")
                print(f"Saving to {save_sample_path}")
                image = sample[0, :, 0]
                image = image.transpose(0, 1).transpose(1, 2)
                image = (image * 255).numpy().astype(np.uint8)
                image = Image.fromarray(image)
                image.save(save_sample_path)

            else:
                save_sample_path = os.path.join(self.savedir_sample, prefix + f"-{md5_hash}.mp4")
                print(f"Saving to {save_sample_path}")
                save_videos_grid(sample, save_sample_path, fps=fps)
            return save_sample_path

        if self.ulysses_degree * self.ring_degree > 1:
            import torch.distributed as dist
            if dist.get_rank() == 0:
                save_sample_path = save_results()
            else:
                save_sample_path = None
        else:
            save_sample_path = save_results()
        return save_sample_path

    def generate(
        self,
        diffusion_transformer_dropdown,
        base_model_dropdown,
        lora_model_dropdown, 
        lora_alpha_slider,
        prompt_textbox, 
        negative_prompt_textbox, 
        sampler_dropdown, 
        sample_step_slider, 
        resize_method,
        width_slider, 
        height_slider, 
        base_resolution, 
        generation_method, 
        length_slider, 
        overlap_video_length, 
        partial_video_length, 
        cfg_scale_slider, 
        start_image, 
        end_image, 
        validation_video,
        validation_video_mask,
        control_video,
        denoise_strength,
        seed_textbox,
        enable_teacache = None, 
        teacache_threshold = None, 
        num_skip_start_steps = None, 
        teacache_offload = None, 
        cfg_skip_ratio = None,
        enable_riflex = None, 
        riflex_k = None, 
        is_api = False,
    ):
        pass

def post_to_host(
    diffusion_transformer_dropdown,
    base_model_dropdown, lora_model_dropdown, lora_alpha_slider,
    prompt_textbox, negative_prompt_textbox, 
    sampler_dropdown, sample_step_slider, resize_method, width_slider, height_slider,
    base_resolution, generation_method, length_slider, cfg_scale_slider, 
    start_image, end_image, validation_video, validation_video_mask, denoise_strength, seed_textbox,
    ref_image = None, enable_teacache = None, teacache_threshold = None, num_skip_start_steps = None, 
    teacache_offload = None, cfg_skip_ratio = None,enable_riflex = None, riflex_k = None, 
):
    if start_image is not None:
        with open(start_image, 'rb') as file:
            file_content = file.read()
            start_image_encoded_content = base64.b64encode(file_content)
            start_image = start_image_encoded_content.decode('utf-8')

    if end_image is not None:
        with open(end_image, 'rb') as file:
            file_content = file.read()
            end_image_encoded_content = base64.b64encode(file_content)
            end_image = end_image_encoded_content.decode('utf-8')

    if validation_video is not None:
        with open(validation_video, 'rb') as file:
            file_content = file.read()
            validation_video_encoded_content = base64.b64encode(file_content)
            validation_video = validation_video_encoded_content.decode('utf-8')

    if validation_video_mask is not None:
        with open(validation_video_mask, 'rb') as file:
            file_content = file.read()
            validation_video_mask_encoded_content = base64.b64encode(file_content)
            validation_video_mask = validation_video_mask_encoded_content.decode('utf-8')

    if ref_image is not None:
        with open(ref_image, 'rb') as file:
            file_content = file.read()
            ref_image_encoded_content = base64.b64encode(file_content)
            ref_image = ref_image_encoded_content.decode('utf-8')

    datas = {
        "base_model_path": base_model_dropdown,
        "lora_model_path": lora_model_dropdown, 
        "lora_alpha_slider": lora_alpha_slider, 
        "prompt_textbox": prompt_textbox, 
        "negative_prompt_textbox": negative_prompt_textbox, 
        "sampler_dropdown": sampler_dropdown, 
        "sample_step_slider": sample_step_slider, 
        "resize_method": resize_method,
        "width_slider": width_slider, 
        "height_slider": height_slider, 
        "base_resolution": base_resolution,
        "generation_method": generation_method,
        "length_slider": length_slider,
        "cfg_scale_slider": cfg_scale_slider,
        "start_image": start_image,
        "end_image": end_image,
        "validation_video": validation_video,
        "validation_video_mask": validation_video_mask,
        "denoise_strength": denoise_strength,
        "seed_textbox": seed_textbox,

        "ref_image": ref_image,
        "enable_teacache": enable_teacache,
        "teacache_threshold": teacache_threshold,
        "num_skip_start_steps": num_skip_start_steps,
        "teacache_offload": teacache_offload,
        "cfg_skip_ratio": cfg_skip_ratio,
        "enable_riflex": enable_riflex,
        "riflex_k": riflex_k,
    }

    session = requests.session()
    session.headers.update({"Authorization": os.environ.get("EAS_TOKEN")})

    response = session.post(url=f'{os.environ.get("EAS_URL")}/videox_fun/infer_forward', json=datas, timeout=300)

    outputs = response.json()
    return outputs


class Fun_Controller_Client:
    def __init__(self, scheduler_dict, savedir_sample):
        self.basedir = os.getcwd()
        if savedir_sample is None:
            self.savedir_sample         = os.path.join(self.basedir, "samples", datetime.now().strftime("Gradio-%Y-%m-%dT%H-%M-%S"))
        else:
            self.savedir_sample         = savedir_sample
        os.makedirs(self.savedir_sample, exist_ok=True)
        
        self.scheduler_dict = scheduler_dict

    def generate(
        self,
        diffusion_transformer_dropdown,
        base_model_dropdown,
        lora_model_dropdown, 
        lora_alpha_slider,
        prompt_textbox, 
        negative_prompt_textbox, 
        sampler_dropdown, 
        sample_step_slider, 
        resize_method,
        width_slider, 
        height_slider, 
        base_resolution, 
        generation_method, 
        length_slider, 
        cfg_scale_slider, 
        start_image, 
        end_image, 
        validation_video, 
        validation_video_mask, 
        denoise_strength,
        seed_textbox,
        ref_image = None,
        enable_teacache = None, 
        teacache_threshold = None, 
        num_skip_start_steps = None, 
        teacache_offload = None, 
        cfg_skip_ratio = None,
        enable_riflex = None, 
        riflex_k = None, 
    ):
        is_image = True if generation_method == "Image Generation" else False

        outputs = post_to_host(
            diffusion_transformer_dropdown,
            base_model_dropdown, lora_model_dropdown, lora_alpha_slider,
            prompt_textbox, negative_prompt_textbox, 
            sampler_dropdown, sample_step_slider, resize_method, width_slider, height_slider,
            base_resolution, generation_method, length_slider, cfg_scale_slider, 
            start_image, end_image, validation_video, validation_video_mask, denoise_strength, 
            seed_textbox, ref_image = ref_image, enable_teacache = enable_teacache, teacache_threshold = teacache_threshold, 
            num_skip_start_steps = num_skip_start_steps, teacache_offload = teacache_offload, 
            cfg_skip_ratio = cfg_skip_ratio, enable_riflex = enable_riflex, riflex_k = riflex_k, 
        )

        try:
            base64_encoding = outputs["base64_encoding"]
        except:
            return gr.Image(visible=False, value=None), gr.Video(None, visible=True), outputs["message"]
            
        decoded_data = base64.b64decode(base64_encoding)

        if not os.path.exists(self.savedir_sample):
            os.makedirs(self.savedir_sample, exist_ok=True)
        md5_hash = hashlib.md5(decoded_data).hexdigest()

        index = len([path for path in os.listdir(self.savedir_sample)]) + 1
        prefix = str(index).zfill(8)
        
        if is_image or length_slider == 1:
            save_sample_path = os.path.join(self.savedir_sample, prefix + f"-{md5_hash}.png")
            print(f"Saving to {save_sample_path}")
            with open(save_sample_path, "wb") as file:
                file.write(decoded_data)
            if gradio_version_is_above_4:
                return gr.Image(value=save_sample_path, visible=True), gr.Video(value=None, visible=False), "Success"
            else:
                return gr.Image.update(value=save_sample_path, visible=True), gr.Video.update(value=None, visible=False), "Success"
        else:
            save_sample_path = os.path.join(self.savedir_sample, prefix + f"-{md5_hash}.mp4")
            print(f"Saving to {save_sample_path}")
            with open(save_sample_path, "wb") as file:
                file.write(decoded_data)
            if gradio_version_is_above_4:
                return gr.Image(visible=False, value=None), gr.Video(value=save_sample_path, visible=True), "Success"
            else:
                return gr.Image.update(visible=False, value=None), gr.Video.update(value=save_sample_path, visible=True), "Success"