ComfyUI-Advanced-ControlNet/adv_control/control_sparsectrl.py

#taken from: https://github.com/lllyasviel/ControlNet
#and modified
#and then taken from comfy/cldm/cldm.py and modified again

from abc import ABC, abstractmethod
import copy
import math
import numpy as np
from typing import Iterable, Union
import torch
import torch as th
import torch.nn as nn
from torch import Tensor
from einops import rearrange, repeat

from comfy.ldm.modules.diffusionmodules.util import (
    zero_module,
    timestep_embedding,
)

from comfy.cli_args import args
from comfy.cldm.cldm import ControlNet as ControlNetCLDM
from comfy.ldm.modules.attention import SpatialTransformer
from comfy.ldm.modules.attention import attention_basic, attention_pytorch, attention_split, attention_sub_quad, default
from comfy.ldm.modules.attention import FeedForward, SpatialTransformer
from comfy.ldm.modules.diffusionmodules.openaimodel import TimestepEmbedSequential
from comfy.model_patcher import ModelPatcher
import comfy.ops
import comfy.model_management
import comfy.utils

from .logger import logger
from .utils import (BIGMAX, AbstractPreprocWrapper, disable_weight_init_clean_groupnorm,
                    prepare_mask_batch, broadcast_image_to_extend, extend_to_batch_size)


# until xformers bug is fixed, do not use xformers for VersatileAttention! TODO: change this when fix is out
# logic for choosing optimized_attention method taken from comfy/ldm/modules/attention.py
# a fallback_attention_mm is selected to avoid CUDA configuration limitation with pytorch's scaled_dot_product
optimized_attention_mm = attention_basic
fallback_attention_mm = attention_basic
if comfy.model_management.xformers_enabled():
    pass
    #optimized_attention_mm = attention_xformers
if comfy.model_management.pytorch_attention_enabled():
    optimized_attention_mm = attention_pytorch
    if args.use_split_cross_attention:
        fallback_attention_mm = attention_split
    else:
        fallback_attention_mm = attention_sub_quad
else:
    if args.use_split_cross_attention:
        optimized_attention_mm = attention_split
    else:
        optimized_attention_mm = attention_sub_quad


class SparseConst:
    HINT_MULT = "sparse_hint_mult"
    NONHINT_MULT = "sparse_nonhint_mult"
    MASK_MULT = "sparse_mask_mult"


class SparseControlNet(ControlNetCLDM):
    def __init__(self, *args,**kwargs):
        super().__init__(*args, **kwargs)
        hint_channels = kwargs.get("hint_channels")
        operations: disable_weight_init_clean_groupnorm = kwargs.get("operations", disable_weight_init_clean_groupnorm)
        device = kwargs.get("device", None)
        self.use_simplified_conditioning_embedding = kwargs.get("use_simplified_conditioning_embedding", False)
        if self.use_simplified_conditioning_embedding:
            self.input_hint_block = TimestepEmbedSequential(
                zero_module(operations.conv_nd(self.dims, hint_channels, self.model_channels, 3, padding=1, dtype=self.dtype, device=device)),
            )
        self.motion_wrapper: SparseCtrlMotionWrapper = None
    
    def set_actual_length(self, actual_length: int, full_length: int):
        if self.motion_wrapper is not None:
            self.motion_wrapper.set_video_length(video_length=actual_length, full_length=full_length)

    def forward(self, x: Tensor, hint: Tensor, timesteps, context, y=None, **kwargs):
        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
        emb = self.time_embed(t_emb)

        # SparseCtrl sets noisy input to zeros
        x = torch.zeros_like(x)
        guided_hint = self.input_hint_block(hint, emb, context)

        out_output = []
        out_middle = []

        hs = []
        if self.num_classes is not None:
            assert y.shape[0] == x.shape[0]
            emb = emb + self.label_emb(y)

        h = x
        for module, zero_conv in zip(self.input_blocks, self.zero_convs):
            if guided_hint is not None:
                h = module(h, emb, context)
                h += guided_hint
                guided_hint = None
            else:
                h = module(h, emb, context)
            out_output.append(zero_conv(h, emb, context))

        h = self.middle_block(h, emb, context)
        out_middle.append(self.middle_block_out(h, emb, context))

        return {"middle": out_middle, "output": out_output}


class SparseModelPatcher(ModelPatcher):
    def __init__(self, *args, **kwargs):
        self.model: SparseControlNet
        super().__init__(*args, **kwargs)
    
    def load(self, device_to=None, lowvram_model_memory=0, *args, **kwargs):
        to_return = super().load(device_to=device_to, lowvram_model_memory=lowvram_model_memory, *args, **kwargs)
        if lowvram_model_memory > 0:
            self._patch_lowvram_extras(device_to=device_to)
        self._handle_float8_pe_tensors()
        return to_return

    def _patch_lowvram_extras(self, device_to=None):
        if self.model.motion_wrapper is not None:
            # figure out the tensors (likely pe's) that should be cast to device besides just the named_modules
            remaining_tensors = list(self.model.motion_wrapper.state_dict().keys())
            named_modules = []
            for n, _ in self.model.motion_wrapper.named_modules():
                named_modules.append(n)
                named_modules.append(f"{n}.weight")
                named_modules.append(f"{n}.bias")
            for name in named_modules:
                if name in remaining_tensors:
                    remaining_tensors.remove(name)

            for key in remaining_tensors:
                self.patch_weight_to_device(key, device_to)
                if device_to is not None:
                    comfy.utils.set_attr(self.model.motion_wrapper, key, comfy.utils.get_attr(self.model.motion_wrapper, key).to(device_to))

    def _handle_float8_pe_tensors(self):
        if self.model.motion_wrapper is not None:
            remaining_tensors = list(self.model.motion_wrapper.state_dict().keys())
            pe_tensors = [x for x in remaining_tensors if '.pe' in x]
            is_first = True
            for key in pe_tensors:
                if is_first:
                    is_first = False
                    if comfy.utils.get_attr(self.model.motion_wrapper, key).dtype not in [torch.float8_e5m2, torch.float8_e4m3fn]:
                        break
                comfy.utils.set_attr(self.model.motion_wrapper, key, comfy.utils.get_attr(self.model.motion_wrapper, key).half())

    # NOTE: no longer called by ComfyUI, but here for backwards compatibility
    def patch_model_lowvram(self, device_to=None, *args, **kwargs):
        patched_model = super().patch_model_lowvram(device_to, *args, **kwargs)
        self._patch_lowvram_extras(device_to=device_to)
        return patched_model

    def clone(self):
        # normal ModelPatcher clone actions
        n = SparseModelPatcher(self.model, self.load_device, self.offload_device, self.size, weight_inplace_update=self.weight_inplace_update)
        n.patches = {}
        for k in self.patches:
            n.patches[k] = self.patches[k][:]
        if hasattr(n, "patches_uuid"):
            self.patches_uuid = n.patches_uuid

        n.object_patches = self.object_patches.copy()
        n.model_options = copy.deepcopy(self.model_options)
        if hasattr(n, "model_keys"):
            n.model_keys = self.model_keys
        if hasattr(n, "backup"):
            self.backup = n.backup
        if hasattr(n, "object_patches_backup"):
            self.object_patches_backup = n.object_patches_backup


class PreprocSparseRGBWrapper(AbstractPreprocWrapper):
    error_msg = error_msg = "Invalid use of RGB SparseCtrl output. The output of RGB SparseCtrl preprocessor is NOT a usual image, but a latent pretending to be an image - you must connect the output directly to an Apply ControlNet node (advanced or otherwise). It cannot be used for anything else that accepts IMAGE input."
    def __init__(self, condhint: Tensor):
        super().__init__(condhint)


class SparseContextAware:
    NEAREST_HINT = "nearest_hint"
    OFF = "off"

    LIST = [NEAREST_HINT, OFF]


class SparseSettings:
    def __init__(self, sparse_method: 'SparseMethod', use_motion: bool=True, motion_strength=1.0, motion_scale=1.0, merged=False,
                 sparse_mask_mult=1.0, sparse_hint_mult=1.0, sparse_nonhint_mult=1.0, context_aware=SparseContextAware.NEAREST_HINT):
        # account for Steerable-Motion workflow incompatibility;
        # doing this to for my own peace of mind (not an issue with my code)
        if type(sparse_method) == str:
            logger.warn("Outdated Steerable-Motion workflow detected; attempting to auto-convert indexes input. If you experience an error here, consult Steerable-Motion github, NOT Advanced-ControlNet.")
            sparse_method = SparseIndexMethod(get_idx_list_from_str(sparse_method))
        self.sparse_method = sparse_method
        self.use_motion = use_motion
        self.motion_strength = motion_strength
        self.motion_scale = motion_scale
        self.merged = merged
        self.sparse_mask_mult = float(sparse_mask_mult)
        self.sparse_hint_mult = float(sparse_hint_mult)
        self.sparse_nonhint_mult = float(sparse_nonhint_mult)
        self.context_aware = context_aware
    
    def is_context_aware(self):
        return self.context_aware != SparseContextAware.OFF

    @classmethod
    def default(cls):
        return SparseSettings(sparse_method=SparseSpreadMethod(), use_motion=True)


class SparseMethod(ABC):
    SPREAD = "spread"
    INDEX = "index"
    def __init__(self, method: str):
        self.method = method

    @abstractmethod
    def _get_indexes(self, hint_length: int, full_length: int) -> list[int]:
        pass

    def get_indexes(self, hint_length: int, full_length: int, sub_idxs: list[int]=None) -> tuple[list[int], list[int]]:
        returned_idxs = self._get_indexes(hint_length, full_length)
        if sub_idxs is None:
            return returned_idxs, None
        # need to map full indexes to condhint indexes
        index_mapping = {}
        for i, value in enumerate(returned_idxs):
            index_mapping[value] = i
        def get_mapped_idxs(idxs: list[int]):
            return [index_mapping[idx] for idx in idxs]
        # check if returned_idxs fit within subidxs
        fitting_idxs = []
        for sub_idx in sub_idxs:
            if sub_idx in returned_idxs:
                fitting_idxs.append(sub_idx)
        # if have any fitting_idxs, deal with it
        if len(fitting_idxs) > 0:
            return fitting_idxs, get_mapped_idxs(fitting_idxs)

        # since no returned_idxs fit in sub_idxs, need to get the next-closest hint images based on strategy
        def get_closest_idx(target_idx: int, idxs: list[int]):
            min_idx = -1
            min_dist = BIGMAX
            for idx in idxs:
                new_dist = abs(idx-target_idx)
                if new_dist < min_dist:
                    min_idx = idx
                    min_dist = new_dist
                    if min_dist == 1:
                        return min_idx, min_dist
            return min_idx, min_dist
        start_closest_idx, start_dist = get_closest_idx(sub_idxs[0], returned_idxs)
        end_closest_idx, end_dist = get_closest_idx(sub_idxs[-1], returned_idxs)
        # if only one cond hint exists, do special behavior 
        if hint_length == 1:
            # if same distance from start and end, 
            if start_dist == end_dist:
                # find center index of sub_idxs
                center_idx = sub_idxs[np.linspace(0, len(sub_idxs)-1, 3, endpoint=True, dtype=int)[1]]
                return [center_idx], get_mapped_idxs([start_closest_idx])
            # otherwise, return closest
            if start_dist < end_dist:
                return [sub_idxs[0]], get_mapped_idxs([start_closest_idx])
            return [sub_idxs[-1]], get_mapped_idxs([end_closest_idx])
        # otherwise, select up to two closest images, or just 1, whichever one applies best
        # if same distance from start and end, return two images to use 
        if start_dist == end_dist:
            return [sub_idxs[0], sub_idxs[-1]], get_mapped_idxs([start_closest_idx, end_closest_idx])
        # else, use just one
        if start_dist < end_dist:
            return [sub_idxs[0]], get_mapped_idxs([start_closest_idx])
        return [sub_idxs[-1]], get_mapped_idxs([end_closest_idx])


class SparseSpreadMethod(SparseMethod):
    UNIFORM = "uniform"
    STARTING = "starting"
    ENDING = "ending"
    CENTER = "center"

    LIST = [UNIFORM, STARTING, ENDING, CENTER]

    def __init__(self, spread=UNIFORM):
        super().__init__(self.SPREAD)
        self.spread = spread

    def _get_indexes(self, hint_length: int, full_length: int) -> list[int]:
        # if hint_length >= full_length, limit hints to full_length
        if hint_length >= full_length:
            return list(range(full_length))
        # handle special case of 1 hint image
        if hint_length == 1:
            if self.spread in [self.UNIFORM, self.STARTING]:
                return [0]
            elif self.spread == self.ENDING:
                return [full_length-1]
            elif self.spread == self.CENTER:
                # return second (of three) values as the center
                return [np.linspace(0, full_length-1, 3, endpoint=True, dtype=int)[1]]
            else:
                raise ValueError(f"Unrecognized spread: {self.spread}")
        # otherwise, handle other cases
        if self.spread == self.UNIFORM:
            return list(np.linspace(0, full_length-1, hint_length, endpoint=True, dtype=int))
        elif self.spread == self.STARTING:
            # make split 1 larger, remove last element
            return list(np.linspace(0, full_length-1, hint_length+1, endpoint=True, dtype=int))[:-1]
        elif self.spread == self.ENDING:
            # make split 1 larger, remove first element
            return list(np.linspace(0, full_length-1, hint_length+1, endpoint=True, dtype=int))[1:]
        elif self.spread == self.CENTER:
            # if hint length is not 3 greater than full length, do STARTING behavior
            if full_length-hint_length < 3:
                return list(np.linspace(0, full_length-1, hint_length+1, endpoint=True, dtype=int))[:-1]
            # otherwise, get linspace of 2 greater than needed, then cut off first and last
            return list(np.linspace(0, full_length-1, hint_length+2, endpoint=True, dtype=int))[1:-1]
        return ValueError(f"Unrecognized spread: {self.spread}")


class SparseIndexMethod(SparseMethod):
    def __init__(self, idxs: list[int]):
        super().__init__(self.INDEX)
        self.idxs = idxs

    def _get_indexes(self, hint_length: int, full_length: int) -> list[int]:
        orig_hint_length = hint_length
        if hint_length > full_length:
            hint_length = full_length
        # if idxs is less than hint_length, throw error
        if len(self.idxs) < hint_length:
            err_msg = f"There are not enough indexes ({len(self.idxs)}) provided to fit the usable {hint_length} input images."
            if orig_hint_length != hint_length:
                err_msg = f"{err_msg} (original input images: {orig_hint_length})"
            raise ValueError(err_msg)
        # cap idxs to hint_length
        idxs = self.idxs[:hint_length]
        new_idxs = []
        real_idxs = set()
        for idx in idxs:
            if idx < 0:
                real_idx = full_length+idx
                if real_idx in real_idxs:
                    raise ValueError(f"Index '{idx}' maps to '{real_idx}' and is duplicate - indexes in Sparse Index Method must be unique.")
            else:
                real_idx = idx
                if real_idx in real_idxs:
                    raise ValueError(f"Index '{idx}' is duplicate (or a negative index is equivalent) - indexes in Sparse Index Method must be unique.")
            real_idxs.add(real_idx)
            new_idxs.append(real_idx)
        return new_idxs  


def get_idx_list_from_str(indexes: str) -> list[int]:
    idxs = []
    unique_idxs = set()
    # get indeces from string
    str_idxs = [x.strip() for x in indexes.strip().split(",")]
    for str_idx in str_idxs:
        try:
            idx = int(str_idx)
            if idx in unique_idxs:
                raise ValueError(f"'{idx}' is duplicated; indexes must be unique.")
            idxs.append(idx)
            unique_idxs.add(idx)
        except ValueError:
            raise ValueError(f"'{str_idx}' is not a valid integer index.")
    if len(idxs) == 0:
        raise ValueError(f"No indexes were listed in Sparse Index Method.")
    return idxs


#########################################
# motion-related portion of controlnet
class BlockType:
    UP = "up"
    DOWN = "down"
    MID = "mid"

def get_down_block_max(mm_state_dict: dict[str, Tensor]) -> int:
    return get_block_max(mm_state_dict, "down_blocks")

def get_up_block_max(mm_state_dict: dict[str, Tensor]) -> int:
    return get_block_max(mm_state_dict, "up_blocks")

def get_block_max(mm_state_dict: dict[str, Tensor], block_name: str) -> int:
    # keep track of biggest down_block count in module
    biggest_block = -1
    for key in mm_state_dict.keys():
        if block_name in key:
            try:
                block_int = key.split(".")[1]
                block_num = int(block_int)
                if block_num > biggest_block:
                    biggest_block = block_num
            except ValueError:
                pass
    return biggest_block

def has_mid_block(mm_state_dict: dict[str, Tensor]):
    # check if keys contain mid_block
    for key in mm_state_dict.keys():
        if key.startswith("mid_block."):
            return True
    return False

def get_position_encoding_max_len(mm_state_dict: dict[str, Tensor], mm_name: str=None) -> int:
    # use pos_encoder.pe entries to determine max length - [1, {max_length}, {320|640|1280}]
    for key in mm_state_dict.keys():
        if key.endswith("pos_encoder.pe"):
            return mm_state_dict[key].size(1) # get middle dim
    raise ValueError(f"No pos_encoder.pe found in SparseCtrl state_dict - {mm_name} is not a valid SparseCtrl model!")


class SparseCtrlMotionWrapper(nn.Module):
    def __init__(self, mm_state_dict: dict[str, Tensor], ops=disable_weight_init_clean_groupnorm):
        super().__init__()
        self.down_blocks: Iterable[MotionModule] = None
        self.up_blocks: Iterable[MotionModule] = None
        self.mid_block: MotionModule = None
        self.encoding_max_len = get_position_encoding_max_len(mm_state_dict, "")
        layer_channels = (320, 640, 1280, 1280)
        if get_down_block_max(mm_state_dict) > -1:
            self.down_blocks = nn.ModuleList([])
            for c in layer_channels:
                self.down_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.DOWN, ops=ops))
        if get_up_block_max(mm_state_dict) > -1:
            self.up_blocks = nn.ModuleList([])
            for c in reversed(layer_channels):
                self.up_blocks.append(MotionModule(c, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.UP, ops=ops))
        if has_mid_block(mm_state_dict):
            self.mid_block = MotionModule(1280, temporal_position_encoding_max_len=self.encoding_max_len, block_type=BlockType.MID, ops=ops)

    def inject(self, unet: SparseControlNet):
        # inject input (down) blocks
        self._inject(unet.input_blocks, self.down_blocks)
        # inject mid block, if present
        if self.mid_block is not None:
            self._inject([unet.middle_block], [self.mid_block])
        unet.motion_wrapper = self

    def _inject(self, unet_blocks: nn.ModuleList, mm_blocks: nn.ModuleList):
        # Rules for injection:
        # For each component list in a unet block:
        #     if SpatialTransformer exists in list, place next block after last occurrence
        #     elif ResBlock exists in list, place next block after first occurrence
        #     else don't place block
        injection_count = 0
        unet_idx = 0
        # details about blocks passed in
        per_block = len(mm_blocks[0].motion_modules)
        injection_goal = len(mm_blocks) * per_block
        # only stop injecting when modules exhausted
        while injection_count < injection_goal:
            # figure out which VanillaTemporalModule from mm to inject
            mm_blk_idx, mm_vtm_idx = injection_count // per_block, injection_count % per_block
            # figure out layout of unet block components
            st_idx = -1 # SpatialTransformer index
            res_idx = -1 # first ResBlock index
            # first, figure out indeces of relevant blocks
            for idx, component in enumerate(unet_blocks[unet_idx]):
                if type(component) == SpatialTransformer:
                    st_idx = idx
                elif type(component).__name__ == "ResBlock" and res_idx < 0:
                    res_idx = idx
            # if SpatialTransformer exists, inject right after
            if st_idx >= 0:
                unet_blocks[unet_idx].insert(st_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
                injection_count += 1
            # otherwise, if only ResBlock exists, inject right after
            elif res_idx >= 0:
                unet_blocks[unet_idx].insert(res_idx+1, mm_blocks[mm_blk_idx].motion_modules[mm_vtm_idx])
                injection_count += 1
            # increment unet_idx
            unet_idx += 1

    def eject(self, unet: SparseControlNet):
        # remove from input blocks (downblocks)
        self._eject(unet.input_blocks)
        # remove from middle block (encapsulate in list to make compatible)
        self._eject([unet.middle_block])
        del unet.motion_wrapper
        unet.motion_wrapper = None

    def _eject(self, unet_blocks: nn.ModuleList):
        # eject all VanillaTemporalModule objects from all blocks
        for block in unet_blocks:
            idx_to_pop = []
            for idx, component in enumerate(block):
                if type(component) == VanillaTemporalModule:
                    idx_to_pop.append(idx)
            # pop in backwards order, as to not disturb what the indeces refer to
            for idx in sorted(idx_to_pop, reverse=True):
                block.pop(idx)

    def set_video_length(self, video_length: int, full_length: int):
        self.AD_video_length = video_length
        if self.down_blocks is not None:
            for block in self.down_blocks:
                block.set_video_length(video_length, full_length)
        if self.up_blocks is not None:
            for block in self.up_blocks:
                block.set_video_length(video_length, full_length)
        if self.mid_block is not None:
            self.mid_block.set_video_length(video_length, full_length)
    
    def set_scale_multiplier(self, multiplier: Union[float, None]):
        if self.down_blocks is not None:
            for block in self.down_blocks:
                block.set_scale_multiplier(multiplier)
        if self.up_blocks is not None:
            for block in self.up_blocks:
                block.set_scale_multiplier(multiplier)
        if self.mid_block is not None:
            self.mid_block.set_scale_multiplier(multiplier)

    def set_strength(self, strength: float):
        if self.down_blocks is not None:
            for block in self.down_blocks:
                block.set_strength(strength)
        if self.up_blocks is not None:
            for block in self.up_blocks:
                block.set_strength(strength)
        if self.mid_block is not None:
            self.mid_block.set_strength(strength)

    def reset_temp_vars(self):
        if self.down_blocks is not None:
            for block in self.down_blocks:
                block.reset_temp_vars()
        if self.up_blocks is not None:
            for block in self.up_blocks:
                block.reset_temp_vars()
        if self.mid_block is not None:
            self.mid_block.reset_temp_vars()

    def reset_scale_multiplier(self):
        self.set_scale_multiplier(None)

    def reset(self):
        self.reset_scale_multiplier()
        self.reset_temp_vars()


class MotionModule(nn.Module):
    def __init__(self, in_channels, temporal_position_encoding_max_len=24, block_type: str=BlockType.DOWN, ops=disable_weight_init_clean_groupnorm):
        super().__init__()
        if block_type == BlockType.MID:
            # mid blocks contain only a single VanillaTemporalModule
            self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList([get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops)])
        else:
            # down blocks contain two VanillaTemporalModules
            self.motion_modules: Iterable[VanillaTemporalModule] = nn.ModuleList(
                [
                    get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops),
                    get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops)
                ]
            )
            # up blocks contain one additional VanillaTemporalModule
            if block_type == BlockType.UP:
                self.motion_modules.append(get_motion_module(in_channels, temporal_position_encoding_max_len, ops=ops))
    
    def set_video_length(self, video_length: int, full_length: int):
        for motion_module in self.motion_modules:
            motion_module.set_video_length(video_length, full_length)
    
    def set_scale_multiplier(self, multiplier: Union[float, None]):
        for motion_module in self.motion_modules:
            motion_module.set_scale_multiplier(multiplier)
    
    def set_masks(self, masks: Tensor, min_val: float, max_val: float):
        for motion_module in self.motion_modules:
            motion_module.set_masks(masks, min_val, max_val)
    
    def set_sub_idxs(self, sub_idxs: list[int]):
        for motion_module in self.motion_modules:
            motion_module.set_sub_idxs(sub_idxs)

    def set_strength(self, strength: float):
        for motion_module in self.motion_modules:
            motion_module.set_strength(strength)

    def reset_temp_vars(self):
        for motion_module in self.motion_modules:
            motion_module.reset_temp_vars()


def get_motion_module(in_channels, temporal_position_encoding_max_len, ops=disable_weight_init_clean_groupnorm):
    # unlike normal AD, there is only one attention block expected in SparseCtrl models
    return VanillaTemporalModule(in_channels=in_channels, attention_block_types=("Temporal_Self",), temporal_position_encoding_max_len=temporal_position_encoding_max_len, ops=ops)


class VanillaTemporalModule(nn.Module):
    def __init__(
        self,
        in_channels,
        num_attention_heads=8,
        num_transformer_block=1,
        attention_block_types=("Temporal_Self", "Temporal_Self"),
        cross_frame_attention_mode=None,
        temporal_position_encoding=True,
        temporal_position_encoding_max_len=24,
        temporal_attention_dim_div=1,
        zero_initialize=True,
        ops=disable_weight_init_clean_groupnorm,
    ):
        super().__init__()
        self.strength = 1.0
        self.temporal_transformer = TemporalTransformer3DModel(
            in_channels=in_channels,
            num_attention_heads=num_attention_heads,
            attention_head_dim=in_channels
            // num_attention_heads
            // temporal_attention_dim_div,
            num_layers=num_transformer_block,
            attention_block_types=attention_block_types,
            cross_frame_attention_mode=cross_frame_attention_mode,
            temporal_position_encoding=temporal_position_encoding,
            temporal_position_encoding_max_len=temporal_position_encoding_max_len,
            ops=ops,
        )

        if zero_initialize:
            self.temporal_transformer.proj_out = zero_module(
                self.temporal_transformer.proj_out
            )

    def set_video_length(self, video_length: int, full_length: int):
        self.temporal_transformer.set_video_length(video_length, full_length)
    
    def set_scale_multiplier(self, multiplier: Union[float, None]):
        self.temporal_transformer.set_scale_multiplier(multiplier)

    def set_masks(self, masks: Tensor, min_val: float, max_val: float):
        self.temporal_transformer.set_masks(masks, min_val, max_val)
    
    def set_sub_idxs(self, sub_idxs: list[int]):
        self.temporal_transformer.set_sub_idxs(sub_idxs)

    def set_strength(self, strength: float):
        self.strength = strength

    def reset_temp_vars(self):
        self.set_strength(1.0)
        self.temporal_transformer.reset_temp_vars()

    def forward(self, input_tensor, encoder_hidden_states=None, attention_mask=None):
        if math.isclose(self.strength, 1.0):
            return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)
        elif math.isclose(self.strength, 0.0):
            return input_tensor
        # elif self.strength > 1.0:
        #     return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength
        else:
            return self.temporal_transformer(input_tensor, encoder_hidden_states, attention_mask)*self.strength + input_tensor*(1.0-self.strength)


class TemporalTransformer3DModel(nn.Module):
    def __init__(
        self,
        in_channels,
        num_attention_heads,
        attention_head_dim,
        num_layers,
        attention_block_types=(
            "Temporal_Self",
            "Temporal_Self",
        ),
        dropout=0.0,
        norm_num_groups=32,
        cross_attention_dim=768,
        activation_fn="geglu",
        attention_bias=False,
        upcast_attention=False,
        cross_frame_attention_mode=None,
        temporal_position_encoding=False,
        temporal_position_encoding_max_len=24,
        ops=disable_weight_init_clean_groupnorm,
    ):
        super().__init__()
        self.video_length = 16
        self.full_length = 16
        self.scale_min = 1.0
        self.scale_max = 1.0
        self.raw_scale_mask: Union[Tensor, None] = None
        self.temp_scale_mask: Union[Tensor, None] = None
        self.sub_idxs: Union[list[int], None] = None
        self.prev_hidden_states_batch = 0


        inner_dim = num_attention_heads * attention_head_dim

        self.norm = ops.GroupNorm(
            num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
        )
        self.proj_in = ops.Linear(in_channels, inner_dim)

        self.transformer_blocks: Iterable[TemporalTransformerBlock] = nn.ModuleList(
            [
                TemporalTransformerBlock(
                    dim=inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    attention_block_types=attention_block_types,
                    dropout=dropout,
                    norm_num_groups=norm_num_groups,
                    cross_attention_dim=cross_attention_dim,
                    activation_fn=activation_fn,
                    attention_bias=attention_bias,
                    upcast_attention=upcast_attention,
                    cross_frame_attention_mode=cross_frame_attention_mode,
                    temporal_position_encoding=temporal_position_encoding,
                    temporal_position_encoding_max_len=temporal_position_encoding_max_len,
                    ops=ops,
                )
                for d in range(num_layers)
            ]
        )
        self.proj_out = ops.Linear(inner_dim, in_channels)

    def set_video_length(self, video_length: int, full_length: int):
        self.video_length = video_length
        self.full_length = full_length
    
    def set_scale_multiplier(self, multiplier: Union[float, None]):
        for block in self.transformer_blocks:
            block.set_scale_multiplier(multiplier)

    def set_masks(self, masks: Tensor, min_val: float, max_val: float):
        self.scale_min = min_val
        self.scale_max = max_val
        self.raw_scale_mask = masks

    def set_sub_idxs(self, sub_idxs: list[int]):
        self.sub_idxs = sub_idxs
        for block in self.transformer_blocks:
            block.set_sub_idxs(sub_idxs)

    def reset_temp_vars(self):
        del self.temp_scale_mask
        self.temp_scale_mask = None
        self.prev_hidden_states_batch = 0
        for block in self.transformer_blocks:
            block.reset_temp_vars()

    def get_scale_mask(self, hidden_states: Tensor) -> Union[Tensor, None]:
        # if no raw mask, return None
        if self.raw_scale_mask is None:
            return None
        shape = hidden_states.shape
        batch, channel, height, width = shape
        # if temp mask already calculated, return it
        if self.temp_scale_mask != None:
            # check if hidden_states batch matches
            if batch == self.prev_hidden_states_batch:
                if self.sub_idxs is not None:
                    return self.temp_scale_mask[:, self.sub_idxs, :]
                return self.temp_scale_mask
            # if does not match, reset cached temp_scale_mask and recalculate it
            del self.temp_scale_mask
            self.temp_scale_mask = None
        # otherwise, calculate temp mask
        self.prev_hidden_states_batch = batch
        mask = prepare_mask_batch(self.raw_scale_mask, shape=(self.full_length, 1, height, width))
        mask = extend_to_batch_size(mask, self.full_length)
        # if mask not the same amount length as full length, make it match
        if self.full_length != mask.shape[0]:
            mask = broadcast_image_to_extend(mask, self.full_length, 1)
        # reshape mask to attention K shape (h*w, latent_count, 1)
        batch, channel, height, width = mask.shape
        # first, perform same operations as on hidden_states,
        # turning (b, c, h, w) -> (b, h*w, c)
        mask = mask.permute(0, 2, 3, 1).reshape(batch, height*width, channel)
        # then, make it the same shape as attention's k, (h*w, b, c)
        mask = mask.permute(1, 0, 2)
        # make masks match the expected length of h*w
        batched_number = shape[0] // self.video_length
        if batched_number > 1:
            mask = torch.cat([mask] * batched_number, dim=0)
        # cache mask and set to proper device
        self.temp_scale_mask = mask
        # move temp_scale_mask to proper dtype + device
        self.temp_scale_mask = self.temp_scale_mask.to(dtype=hidden_states.dtype, device=hidden_states.device)
        # return subset of masks, if needed
        if self.sub_idxs is not None:
            return self.temp_scale_mask[:, self.sub_idxs, :]
        return self.temp_scale_mask

    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
        batch, channel, height, width = hidden_states.shape
        residual = hidden_states
        scale_mask = self.get_scale_mask(hidden_states)
        # add some casts for fp8 purposes - does not affect speed otherwise
        hidden_states = self.norm(hidden_states).to(hidden_states.dtype)
        inner_dim = hidden_states.shape[1]
        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
            batch, height * width, inner_dim
        )
        hidden_states = self.proj_in(hidden_states).to(hidden_states.dtype)

        # Transformer Blocks
        for block in self.transformer_blocks:
            hidden_states = block(
                hidden_states,
                encoder_hidden_states=encoder_hidden_states,
                attention_mask=attention_mask,
                video_length=self.video_length,
                scale_mask=scale_mask
            )

        # output
        hidden_states = self.proj_out(hidden_states)
        hidden_states = (
            hidden_states.reshape(batch, height, width, inner_dim)
            .permute(0, 3, 1, 2)
            .contiguous()
        )

        output = hidden_states + residual

        return output


class TemporalTransformerBlock(nn.Module):
    def __init__(
        self,
        dim,
        num_attention_heads,
        attention_head_dim,
        attention_block_types=(
            "Temporal_Self",
            "Temporal_Self",
        ),
        dropout=0.0,
        norm_num_groups=32,
        cross_attention_dim=768,
        activation_fn="geglu",
        attention_bias=False,
        upcast_attention=False,
        cross_frame_attention_mode=None,
        temporal_position_encoding=False,
        temporal_position_encoding_max_len=24,
        ops=disable_weight_init_clean_groupnorm,
    ):
        super().__init__()

        attention_blocks = []
        norms = []

        for block_name in attention_block_types:
            attention_blocks.append(
                VersatileAttention(
                    attention_mode=block_name.split("_")[0],
                    context_dim=cross_attention_dim # called context_dim for ComfyUI impl
                    if block_name.endswith("_Cross")
                    else None,
                    query_dim=dim,
                    heads=num_attention_heads,
                    dim_head=attention_head_dim,
                    dropout=dropout,
                    #bias=attention_bias, # remove for Comfy CrossAttention
                    #upcast_attention=upcast_attention, # remove for Comfy CrossAttention
                    cross_frame_attention_mode=cross_frame_attention_mode,
                    temporal_position_encoding=temporal_position_encoding,
                    temporal_position_encoding_max_len=temporal_position_encoding_max_len,
                    ops=ops,
                )
            )
            norms.append(ops.LayerNorm(dim))

        self.attention_blocks: Iterable[VersatileAttention] = nn.ModuleList(attention_blocks)
        self.norms = nn.ModuleList(norms)

        self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn == "geglu"), operations=ops)
        self.ff_norm = ops.LayerNorm(dim)

    def set_scale_multiplier(self, multiplier: Union[float, None]):
        for block in self.attention_blocks:
            block.set_scale_multiplier(multiplier)

    def set_sub_idxs(self, sub_idxs: list[int]):
        for block in self.attention_blocks:
            block.set_sub_idxs(sub_idxs)

    def reset_temp_vars(self):
        for block in self.attention_blocks:
            block.reset_temp_vars()

    def forward(
        self,
        hidden_states,
        encoder_hidden_states=None,
        attention_mask=None,
        video_length=None,
        scale_mask=None
    ):
        for attention_block, norm in zip(self.attention_blocks, self.norms):
            norm_hidden_states = norm(hidden_states).to(hidden_states.dtype)
            hidden_states = (
                attention_block(
                    norm_hidden_states,
                    encoder_hidden_states=encoder_hidden_states
                    if attention_block.is_cross_attention
                    else None,
                    attention_mask=attention_mask,
                    video_length=video_length,
                    scale_mask=scale_mask
                )
                + hidden_states
            )

        hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states

        output = hidden_states
        return output


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.0, max_len=24):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(1, max_len, d_model)
        pe[0, :, 0::2] = torch.sin(position * div_term)
        pe[0, :, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)
        self.sub_idxs = None

    def set_sub_idxs(self, sub_idxs: list[int]):
        self.sub_idxs = sub_idxs

    def forward(self, x):
        #if self.sub_idxs is not None:
        #    x = x + self.pe[:, self.sub_idxs]
        #else:
        x = x + self.pe[:, : x.size(1)]
        return self.dropout(x)


class CrossAttentionMMSparse(nn.Module):
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None,
                 operations=disable_weight_init_clean_groupnorm):
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)

        self.actual_attention = optimized_attention_mm
        self.heads = heads
        self.dim_head = dim_head
        self.scale = None

        self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
        self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
        self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)

        self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))

    def reset_attention_type(self):
        self.actual_attention = optimized_attention_mm

    def forward(self, x, context=None, value=None, mask=None, scale_mask=None):
        q = self.to_q(x)
        context = default(context, x)
        k: Tensor = self.to_k(context)
        if value is not None:
            v = self.to_v(value)
            del value
        else:
            v = self.to_v(context)

        # apply custom scale by multiplying k by scale factor
        if self.scale is not None:
            k *= self.scale
        
        # apply scale mask, if present
        if scale_mask is not None:
            k *= scale_mask

        try:
            out = self.actual_attention(q, k, v, self.heads, mask)
        except RuntimeError as e:
            if str(e).startswith("CUDA error: invalid configuration argument"):
                self.actual_attention = fallback_attention_mm
                out = self.actual_attention(q, k, v, self.heads, mask)
            else:
                raise
        return self.to_out(out)


class VersatileAttention(CrossAttentionMMSparse):
    def __init__(
        self,
        attention_mode=None,
        cross_frame_attention_mode=None,
        temporal_position_encoding=False,
        temporal_position_encoding_max_len=24,
        ops=disable_weight_init_clean_groupnorm,
        *args,
        **kwargs,
    ):
        super().__init__(operations=ops, *args, **kwargs)
        assert attention_mode == "Temporal"

        self.attention_mode = attention_mode
        self.is_cross_attention = kwargs["context_dim"] is not None

        self.pos_encoder = (
            PositionalEncoding(
                kwargs["query_dim"],
                dropout=0.0,
                max_len=temporal_position_encoding_max_len,
            )
            if (temporal_position_encoding and attention_mode == "Temporal")
            else None
        )

    def extra_repr(self):
        return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"

    def set_scale_multiplier(self, multiplier: Union[float, None]):
        if multiplier is None or math.isclose(multiplier, 1.0):
            self.scale = None
        else:
            self.scale = multiplier

    def set_sub_idxs(self, sub_idxs: list[int]):
        if self.pos_encoder != None:
            self.pos_encoder.set_sub_idxs(sub_idxs)

    def reset_temp_vars(self):
        self.reset_attention_type()

    def forward(
        self,
        hidden_states: Tensor,
        encoder_hidden_states=None,
        attention_mask=None,
        video_length=None,
        scale_mask=None,
    ):
        if self.attention_mode != "Temporal":
            raise NotImplementedError

        d = hidden_states.shape[1]
        hidden_states = rearrange(
            hidden_states, "(b f) d c -> (b d) f c", f=video_length
        )

        if self.pos_encoder is not None:
           hidden_states = self.pos_encoder(hidden_states).to(hidden_states.dtype)

        encoder_hidden_states = (
            repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)
            if encoder_hidden_states is not None
            else encoder_hidden_states
        )

        hidden_states = super().forward(
            hidden_states,
            encoder_hidden_states,
            value=None,
            mask=attention_mask,
            scale_mask=scale_mask,
        )

        hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)

        return hidden_states