custom_nodes/ComfyUI-AnimateDiff-Evolved/animatediff/utils_model.py

import hashlib
from pathlib import Path
from typing import Callable, Union
from collections.abc import Iterable
from time import time
import copy

from torch import Tensor
import torch
import numpy as np

import folder_paths
from comfy.model_base import SD21UNCLIP, SDXL, BaseModel, SDXLRefiner, SVD_img2vid, model_sampling, ModelType
from comfy.model_management import xformers_enabled
from comfy.model_patcher import ModelPatcher
from comfy.sd import VAE
from comfy.utils import ProgressBar

import comfy.model_sampling
import comfy_extras.nodes_model_advanced

from .logger import logger

BIGMIN = -(2**53-1)
BIGMAX = (2**53-1)

MAX_RESOLUTION = 16384  # mirrors ComfyUI's nodes.py MAX_RESOLUTION

class MachineState:
    READ = "read"
    WRITE = "write"
    READ_WRITE = "read_write"
    OFF = "off"


def vae_encode_raw_dynamic_batched(vae: VAE, pixels: Tensor, max_batch=16, min_batch=1, max_size=512*512, show_pbar=False):
    b, h, w, c = pixels.shape
    actual_size = h*w
    actual_batch_size = int(max(min_batch, min(max_batch, max_batch // max((actual_size / max_size), 1.0))))
    return vae_encode_raw_batched(vae=vae, pixels=pixels, per_batch=actual_batch_size, show_pbar=show_pbar)


def vae_decode_raw_dynamic_batched(vae: VAE, latents: Tensor, max_batch=16, min_batch=1, max_size=512*512, show_pbar=False):
    b, c, h, w = latents.shape
    actual_size = (h*vae.downscale_ratio)*(w*vae.downscale_ratio)
    actual_batch_size = int(max(min_batch, min(max_batch, max_batch // max((actual_size / max_size), 1.0))))
    return vae_decode_raw_batched(vae=vae, latents=latents, per_batch=actual_batch_size, show_pbar=show_pbar)


def vae_encode_raw_batched(vae: VAE, pixels: Tensor, per_batch=16, show_pbar=False):
    encoded = []
    pbar = None
    if show_pbar:
        pbar = ProgressBar(pixels.shape[0])
    for start_idx in range(0, pixels.shape[0], per_batch):
        sub_encoded = vae.encode(pixels[start_idx:start_idx+per_batch][:,:,:,:3])
        encoded.append(sub_encoded)
        if pbar is not None:
            pbar.update(sub_encoded.shape[0])
    return torch.cat(encoded, dim=0)


def vae_decode_raw_batched(vae: VAE, latents: Tensor, per_batch=16, show_pbar=False):
    decoded = []
    pbar = None
    if show_pbar:
        pbar = ProgressBar(latents.shape[0])
    for start_idx in range(0, latents.shape[0], per_batch):
        sub_decoded = vae.decode(latents[start_idx:start_idx+per_batch])
        decoded.append(sub_decoded)
        if pbar is not None:
            pbar.update(sub_decoded.shape[0])
    return torch.cat(decoded, dim=0)


class ModelSamplingConfig:
    def __init__(self, beta_schedule: str, linear_start: float=None, linear_end: float=None, given_betas: Tensor=None, timesteps: int=None):
        self.sampling_settings = {"beta_schedule": beta_schedule}
        if linear_start is not None:
            self.sampling_settings["linear_start"] = linear_start
        if linear_end is not None:
            self.sampling_settings["linear_end"] = linear_end
        if given_betas is not None:
            self.sampling_settings["given_betas"] = given_betas
        if timesteps is not None:
            self.sampling_settings["timesteps"] = timesteps


class ModelSamplingType:
    EPS = "eps"
    V_PREDICTION = "v_prediction"
    LCM = "lcm"

    _NON_LCM_LIST = [EPS, V_PREDICTION]
    _FULL_LIST = [EPS, V_PREDICTION, LCM]

    MAP = {
        EPS: ModelType.EPS,
        V_PREDICTION: ModelType.V_PREDICTION,
        LCM: comfy_extras.nodes_model_advanced.LCM,
    }

    @classmethod
    def from_alias(cls, alias: str):
        return cls.MAP[alias]


def factory_model_sampling_discrete_distilled(original_timesteps=50):
    class ModelSamplingDiscreteDistilledEvolved(comfy_extras.nodes_model_advanced.ModelSamplingDiscreteDistilled):
        def __init__(self, *args, **kwargs):
            self.original_timesteps = original_timesteps  # normal LCM has 50
            super().__init__(*args, **kwargs)
    return ModelSamplingDiscreteDistilledEvolved


# based on code in comfy_extras/nodes_model_advanced.py
def evolved_model_sampling(model_config: ModelSamplingConfig, model_type: ModelType, alias: str, original_timesteps: Union[int, None]=None):
    # if LCM, need to handle manually
    if BetaSchedules.is_lcm(alias) or original_timesteps is not None:
        sampling_type = comfy_extras.nodes_model_advanced.LCM
        if original_timesteps is not None:
            sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=original_timesteps)
        elif alias == BetaSchedules.LCM_100:
            sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=100)
        elif alias == BetaSchedules.LCM_25:
            sampling_base = factory_model_sampling_discrete_distilled(original_timesteps=25)
        else:
            sampling_base = comfy_extras.nodes_model_advanced.ModelSamplingDiscreteDistilled
        class ModelSamplingAdvancedEvolved(sampling_base, sampling_type):
            pass
        # NOTE: if I want to support zsnr, this is where I would add that code
        return ModelSamplingAdvancedEvolved(model_config)
    # otherwise, use vanilla model_sampling function
    ms = model_sampling(model_config, model_type)
    if "given_betas" in model_config.sampling_settings:
        beta_schedule = model_config.sampling_settings.get("beta_schedule", "linear")
        linear_start = model_config.sampling_settings.get("linear_start", 0.00085)
        linear_end = model_config.sampling_settings.get("linear_end", 0.012)
        timesteps = model_config.sampling_settings.get("timesteps", 1000)
        given_betas = model_config.sampling_settings.get("given_betas", None)
        ms._register_schedule(given_betas=given_betas, beta_schedule=beta_schedule,
                              timesteps=timesteps, linear_start=linear_start, linear_end=linear_end)
    return ms


class BetaSchedules:
    AUTOSELECT = "autoselect"
    SQRT_LINEAR = "sqrt_linear (AnimateDiff)"
    LINEAR_ADXL = "linear (AnimateDiff-SDXL)"
    LINEAR = "linear (HotshotXL/default)"
    AVG_LINEAR_SQRT_LINEAR = "avg(sqrt_linear,linear)"
    LCM_AVG_LINEAR_SQRT_LINEAR = "lcm avg(sqrt_linear,linear)"
    LCM = "lcm"
    LCM_100 = "lcm[100_ots]"
    LCM_25 = "lcm[25_ots]"
    LCM_SQRT_LINEAR = "lcm >> sqrt_linear"
    USE_EXISTING = "use existing"
    SQRT = "sqrt"
    COSINE = "cosine"
    SQUAREDCOS_CAP_V2 = "squaredcos_cap_v2"
    RAW_LINEAR = "linear"
    RAW_SQRT_LINEAR = "sqrt_linear"

    RAW_BETA_SCHEDULE_LIST = [RAW_LINEAR, RAW_SQRT_LINEAR, SQRT, COSINE, SQUAREDCOS_CAP_V2]

    ALIAS_LCM_LIST = [LCM, LCM_100, LCM_25, LCM_SQRT_LINEAR]

    ALIAS_ACTIVE_LIST = [SQRT_LINEAR, LINEAR_ADXL, LINEAR, AVG_LINEAR_SQRT_LINEAR, LCM_AVG_LINEAR_SQRT_LINEAR, LCM, LCM_100, LCM_SQRT_LINEAR, # LCM_25 is purposely omitted
                  SQRT, COSINE, SQUAREDCOS_CAP_V2]

    ALIAS_LIST = [AUTOSELECT, USE_EXISTING] + ALIAS_ACTIVE_LIST

    

    ALIAS_MAP = {
        SQRT_LINEAR: "sqrt_linear",
        LINEAR_ADXL: "linear", # also linear, but has different linear_end (0.020)
        LINEAR: "linear",
        LCM_100: "linear",  # distilled, 100 original timesteps
        LCM_25: "linear",  # distilled, 25 original timesteps
        LCM: "linear",  # distilled
        LCM_SQRT_LINEAR: "sqrt_linear", # distilled, sqrt_linear
        SQRT: "sqrt",
        COSINE: "cosine",
        SQUAREDCOS_CAP_V2: "squaredcos_cap_v2",
        RAW_LINEAR: "linear",
        RAW_SQRT_LINEAR: "sqrt_linear"
    }

    @classmethod
    def is_lcm(cls, alias: str):
        return alias in cls.ALIAS_LCM_LIST

    @classmethod
    def to_name(cls, alias: str):
        return cls.ALIAS_MAP[alias]
    
    @classmethod
    def to_config(cls, alias: str) -> ModelSamplingConfig:
        linear_start = None
        linear_end = None
        if alias == cls.LINEAR_ADXL:
            # uses linear_end=0.020
            linear_end = 0.020
        return ModelSamplingConfig(cls.to_name(alias), linear_start=linear_start, linear_end=linear_end)
    
    @classmethod
    def _to_model_sampling(cls, alias: str, model_type: ModelType, config_override: Union[ModelSamplingConfig,None]=None, original_timesteps: Union[int,None]=None):
        if alias == cls.USE_EXISTING:
            return None
        elif config_override != None:
            return evolved_model_sampling(config_override, model_type=model_type, alias=alias, original_timesteps=original_timesteps)
        elif alias == cls.AVG_LINEAR_SQRT_LINEAR:
            ms_linear = evolved_model_sampling(cls.to_config(cls.LINEAR), model_type=model_type, alias=cls.LINEAR)
            ms_sqrt_linear = evolved_model_sampling(cls.to_config(cls.SQRT_LINEAR), model_type=model_type, alias=cls.SQRT_LINEAR)
            avg_sigmas = (ms_linear.sigmas + ms_sqrt_linear.sigmas) / 2
            ms_linear.set_sigmas(avg_sigmas)
            return ms_linear
        elif alias == cls.LCM_AVG_LINEAR_SQRT_LINEAR:
            ms_linear = evolved_model_sampling(cls.to_config(cls.LCM), model_type=model_type, alias=cls.LCM)
            ms_sqrt_linear = evolved_model_sampling(cls.to_config(cls.LCM_SQRT_LINEAR), model_type=model_type, alias=cls.LCM_SQRT_LINEAR)
            avg_sigmas = (ms_linear.sigmas + ms_sqrt_linear.sigmas) / 2
            ms_linear.set_sigmas(avg_sigmas)
            return ms_linear
            # average out the sigmas
        ms_obj = evolved_model_sampling(cls.to_config(alias), model_type=model_type, alias=alias, original_timesteps=original_timesteps)
        return ms_obj

    @classmethod
    def to_model_sampling(cls, alias: str, model: ModelPatcher):
        return cls._to_model_sampling(alias=alias, model_type=model.model.model_type)

    @staticmethod
    def get_alias_list_with_first_element(first_element: str):
        new_list = BetaSchedules.ALIAS_LIST.copy()
        element_index = new_list.index(first_element)
        new_list[0], new_list[element_index] = new_list[element_index], new_list[0]
        return new_list


class SigmaSchedule:
    def __init__(self, model_sampling: comfy.model_sampling.ModelSamplingDiscrete, model_type: ModelType):
        self.model_sampling = model_sampling
        #self.config = config
        self.model_type = model_type
        self.original_timesteps = getattr(self.model_sampling, "original_timesteps", None)
    
    def is_lcm(self):
        return self.original_timesteps is not None

    def total_sigmas(self):
        return len(self.model_sampling.sigmas)
    
    def clone(self) -> 'SigmaSchedule':
        new_model_sampling = copy.deepcopy(self.model_sampling)
        #new_config = copy.deepcopy(self.config)
        return SigmaSchedule(model_sampling=new_model_sampling, model_type=self.model_type)

    # def clone(self):
    #     pass

    @staticmethod
    def apply_zsnr(new_model_sampling: comfy.model_sampling.ModelSamplingDiscrete):
        new_model_sampling.set_sigmas(comfy_extras.nodes_model_advanced.rescale_zero_terminal_snr_sigmas(new_model_sampling.sigmas))

    # def get_lcmified(self, original_timesteps=50, zsnr=False) -> 'SigmaSchedule':
    #     new_model_sampling = evolved_model_sampling(model_config=self.config, model_type=self.model_type, alias=None, original_timesteps=original_timesteps)
    #     if zsnr:
    #         new_model_sampling.set_sigmas(comfy_extras.nodes_model_advanced.rescale_zero_terminal_snr_sigmas(new_model_sampling.sigmas))
    #     return SigmaSchedule(model_sampling=new_model_sampling, config=self.config, model_type=self.model_type, is_lcm=True)
        

class InterpolationMethod:
    LINEAR = "linear"
    EASE_IN = "ease_in"
    EASE_OUT = "ease_out"
    EASE_IN_OUT = "ease_in_out"

    _LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]

    @classmethod
    def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
        diff = num_to - num_from
        if method == cls.LINEAR:
            weights = torch.linspace(num_from, num_to, length)
        elif method == cls.EASE_IN:
            index = torch.linspace(0, 1, length)
            weights = diff * np.power(index, 2) + num_from
        elif method == cls.EASE_OUT:
            index = torch.linspace(0, 1, length)
            weights = diff * (1 - np.power(1 - index, 2)) + num_from
        elif method == cls.EASE_IN_OUT:
            index = torch.linspace(0, 1, length)
            weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
        else:
            raise ValueError(f"Unrecognized interpolation method '{method}'.")
        if reverse:
            weights = weights.flip(dims=(0,))
        return weights


class ScaleMethods:
    NEAREST_EXACT = "nearest-exact"
    BILINEAR = "bilinear"
    AREA = "area"
    BICUBIC = "bicubic"
    LANCZOS = "lanczos"

    _LIST_IMAGE = [NEAREST_EXACT, BILINEAR, AREA, BICUBIC, LANCZOS]


class CropMethods:
    DISABLED = "disabled"
    CENTER = "center"

    _LIST = [DISABLED, CENTER]


class Folders:
    ANIMATEDIFF_MODELS = "animatediff_models"
    MOTION_LORA = "animatediff_motion_lora"
    VIDEO_FORMATS = "animatediff_video_formats"


def add_extension_to_folder_path(folder_name: str, extensions: Union[str, list[str]]):
    if folder_name in folder_paths.folder_names_and_paths:
        if isinstance(extensions, str):
            folder_paths.folder_names_and_paths[folder_name][1].add(extensions)
        elif isinstance(extensions, Iterable):
            for ext in extensions:
                folder_paths.folder_names_and_paths[folder_name][1].add(ext) 


def try_mkdir(full_path: str):
    try:
        Path(full_path).mkdir()
    except Exception:
        pass


# register motion models folder(s)
folder_paths.add_model_folder_path(Folders.ANIMATEDIFF_MODELS, str(Path(__file__).parent.parent / "models"))
folder_paths.add_model_folder_path(Folders.ANIMATEDIFF_MODELS, str(Path(folder_paths.models_dir) / Folders.ANIMATEDIFF_MODELS))
add_extension_to_folder_path(Folders.ANIMATEDIFF_MODELS, folder_paths.supported_pt_extensions)
try_mkdir(str(Path(folder_paths.models_dir) / Folders.ANIMATEDIFF_MODELS))

# register motion LoRA folder(s)
folder_paths.add_model_folder_path(Folders.MOTION_LORA, str(Path(__file__).parent.parent / "motion_lora"))
folder_paths.add_model_folder_path(Folders.MOTION_LORA, str(Path(folder_paths.models_dir) / Folders.MOTION_LORA))
add_extension_to_folder_path(Folders.MOTION_LORA, folder_paths.supported_pt_extensions)
try_mkdir(str(Path(folder_paths.models_dir) / Folders.MOTION_LORA))

# register video_formats folder
folder_paths.add_model_folder_path(Folders.VIDEO_FORMATS, str(Path(__file__).parent.parent / "video_formats"))
add_extension_to_folder_path(Folders.VIDEO_FORMATS, ".json")


def get_available_motion_models():
    return folder_paths.get_filename_list(Folders.ANIMATEDIFF_MODELS)


def get_motion_model_path(model_name: str):
    return folder_paths.get_full_path(Folders.ANIMATEDIFF_MODELS, model_name)


def get_available_motion_loras():
    return folder_paths.get_filename_list(Folders.MOTION_LORA)


def get_motion_lora_path(lora_name: str):
    return folder_paths.get_full_path(Folders.MOTION_LORA, lora_name)


# modified from https://stackoverflow.com/questions/22058048/hashing-a-file-in-python
def calculate_file_hash(filename: str, hash_every_n: int = 50):
    h = hashlib.sha256()
    b = bytearray(1024*1024)
    mv = memoryview(b)
    with open(filename, 'rb', buffering=0) as f:
        i = 0
        # don't hash entire file, only portions of it
        while n := f.readinto(mv):
            if i%hash_every_n == 0:
                h.update(mv[:n])
            i += 1
    return h.hexdigest()


def calculate_model_hash(model: ModelPatcher):
    unet = model.model.diff
    t = unet.input_blocks[1]
    m = hashlib.sha256()
    for buf in t.buffers():
        m.update(buf.cpu().numpy().view(np.uint8))
    return m.hexdigest()


def strip_path(path):
    # removes whitespace and single quotes from either end of string, if present
    path = path.strip()
    if path.startswith("\""):
        path = path[1:]
    if path.endswith("\""):
        path = path[:-1]
    return path


class ModelTypeSD:
    SD1_5 = "SD1.5"
    SD2_1 = "SD2.1"
    SDXL = "SDXL"
    SDXL_REFINER = "SDXL_Refiner"
    SVD = "SVD"

    _LIST = [SD1_5, SD2_1, SDXL, SDXL_REFINER, SVD]


def get_sd_model_type(model: ModelPatcher) -> str:
    if model is None:
        return None
    type_str = str(type(model.model).__name__)
    # instructpix2pix models should be allowed to work with AD
    if type(model.model) == BaseModel or type_str == "SD15_instructpix2pix":
        return ModelTypeSD.SD1_5
    elif type(model.model) == SDXL or type_str == "SDXL_instructpix2pix":
        return ModelTypeSD.SDXL
    elif type(model.model) == SD21UNCLIP:
        return ModelTypeSD.SD2_1
    elif type(model.model) == SDXLRefiner:
        return ModelTypeSD.SDXL_REFINER
    elif type(model.model) == SVD_img2vid:
        return ModelTypeSD.SVD
    else:
        return type_str

def is_checkpoint_sd1_5(model: ModelPatcher):
    return False if model is None else type(model.model) == BaseModel

def is_checkpoint_sdxl(model: ModelPatcher):
    return False if model is None else type(model.model) == SDXL


def raise_if_not_checkpoint_sd1_5(model: ModelPatcher):
    if not is_checkpoint_sd1_5(model):
        raise ValueError(f"For AnimateDiff, SD Checkpoint (model) is expected to be SD1.5-based (BaseModel), but was: {type(model.model).__name__}")


# TODO: remove this filth when xformers bug gets fixed in future xformers version
# NOTE: avoid using this for now to avoid false positives with pytorch or non-AD stuff like SVD
def wrap_function_to_inject_xformers_bug_info(function_to_wrap: Callable) -> Callable:
    if not xformers_enabled:
        return function_to_wrap
    else:
        def wrapped_function(*args, **kwargs):
            try:
                return function_to_wrap(*args, **kwargs)
            except RuntimeError as e:
                if str(e).startswith("CUDA error: invalid configuration argument"):
                    raise RuntimeError(f"An xformers bug was encountered in AnimateDiff - this is unexpected, \
                                       report this to Kosinkadink/ComfyUI-AnimateDiff-Evolved repo as an issue, \
                                       and a workaround for now is to run ComfyUI with the --disable-xformers argument.")
                raise
        return wrapped_function


class Timer(object):
    __slots__ = ("start_time", "end_time")

    def __init__(self) -> None:
        self.start_time = 0.0
        self.end_time = 0.0

    def start(self) -> None:
        self.start_time = time()

    def update(self) -> None:
        self.start()

    def stop(self) -> float:
        self.end_time = time()
        return self.get_time_diff()

    def get_time_diff(self) -> float:
        return self.end_time - self.start_time

    def get_time_current(self) -> float:
        return time() - self.start_time


# TODO: possibly add configuration file in future when needed?
# # Load config settings
# ADE_DIR = Path(__file__).parent.parent
# ADE_CONFIG_FILE = ADE_DIR / "ade_config.json"

# class ADE_Settings:
#     USE_XFORMERS_IN_VERSATILE_ATTENTION = "use_xformers_in_VersatileAttention"

# # Create ADE config if not present
# ABS_CONFIG = {
#     ADE_Settings.USE_XFORMERS_IN_VERSATILE_ATTENTION: True
# }
# if not ADE_CONFIG_FILE.exists():
#     with ADE_CONFIG_FILE.open("w") as f:
#         json.dumps(ABS_CONFIG, indent=4)
# # otherwise, load it and use values
# else:
#     loaded_values: dict = None
#     with ADE_CONFIG_FILE.open("r") as f:
#         loaded_values = json.load(f)
#     if loaded_values is not None:
#         for key, value in loaded_values.items():
#             if key in ABS_CONFIG:
#                 ABS_CONFIG[key] = value