113 lines
4.0 KiB
Python
113 lines
4.0 KiB
Python
####################################################################################################
|
|
# DinkLink is my method of sharing classes/functions between my nodes.
|
|
#
|
|
# My DinkLink-compatible nodes will inject comfy.hooks with a __DINKLINK attr
|
|
# that stores a dictionary, where any of my node packs can store their stuff.
|
|
#
|
|
# It is not intended to be accessed by node packs that I don't develop, so things may change
|
|
# at any time.
|
|
#
|
|
# DinkLink also serves as a proof-of-concept for a future ComfyUI implementation of
|
|
# purposely exposing node pack classes/functions with other node packs.
|
|
####################################################################################################
|
|
from __future__ import annotations
|
|
from typing import Union
|
|
from torch import Tensor, nn
|
|
|
|
from comfy.model_patcher import ModelPatcher
|
|
import comfy.hooks
|
|
|
|
DINKLINK = "__DINKLINK"
|
|
|
|
|
|
def init_dinklink():
|
|
create_dinklink()
|
|
prepare_dinklink()
|
|
|
|
def create_dinklink():
|
|
if not hasattr(comfy.hooks, DINKLINK):
|
|
setattr(comfy.hooks, DINKLINK, {})
|
|
|
|
def get_dinklink() -> dict[str, dict[str]]:
|
|
create_dinklink()
|
|
return getattr(comfy.hooks, DINKLINK)
|
|
|
|
|
|
class DinkLinkConst:
|
|
VERSION = "version"
|
|
# ADE
|
|
ADE = "ADE"
|
|
ADE_ANIMATEDIFFMODEL = "AnimateDiffModel"
|
|
ADE_ANIMATEDIFFINFO = "AnimateDiffInfo"
|
|
ADE_CREATE_MOTIONMODELPATCHER = "create_MotionModelPatcher"
|
|
|
|
def prepare_dinklink():
|
|
pass
|
|
|
|
|
|
class InterfaceAnimateDiffInfo:
|
|
'''Class only used for IDE type hints; interface of ADE's AnimateDiffInfo'''
|
|
def __init__(self, sd_type: str, mm_format: str, mm_version: str, mm_name: str):
|
|
self.sd_type = sd_type
|
|
self.mm_format = mm_format
|
|
self.mm_version = mm_version
|
|
self.mm_name = mm_name
|
|
|
|
|
|
class InterfaceAnimateDiffModel(nn.Module):
|
|
'''Class only used for IDE type hints; interface of ADE's AnimateDiffModel'''
|
|
def __init__(self, mm_state_dict: dict[str, Tensor], mm_info: InterfaceAnimateDiffInfo, init_kwargs: dict[str]={}):
|
|
pass
|
|
|
|
def set_video_length(self, video_length: int, full_length: int) -> None:
|
|
raise NotImplemented()
|
|
|
|
def set_scale(self, scale: Union[float, Tensor, None], per_block_list: Union[list, None]=None) -> None:
|
|
raise NotImplemented()
|
|
|
|
def set_effect(self, multival: Union[float, Tensor, None], per_block_list: Union[list, None]=None) -> None:
|
|
raise NotImplemented()
|
|
|
|
def cleanup(self):
|
|
raise NotImplemented()
|
|
|
|
def inject(self, model: ModelPatcher):
|
|
pass
|
|
|
|
def eject(self, model: ModelPatcher):
|
|
pass
|
|
|
|
|
|
def get_CreateMotionModelPatcher(throw_exception=True):
|
|
d = get_dinklink()
|
|
try:
|
|
link_ade = d[DinkLinkConst.ADE]
|
|
return link_ade[DinkLinkConst.ADE_CREATE_MOTIONMODELPATCHER]
|
|
except KeyError:
|
|
if throw_exception:
|
|
raise Exception("Could not get create_MotionModelPatcher function. AnimateDiff-Evolved nodes need to be installed to use SparseCtrl; " + \
|
|
"they are either not installed or are of an insufficient version.")
|
|
return None
|
|
|
|
def get_AnimateDiffModel(throw_exception=True):
|
|
d = get_dinklink()
|
|
try:
|
|
link_ade = d[DinkLinkConst.ADE]
|
|
return link_ade[DinkLinkConst.ADE_ANIMATEDIFFMODEL]
|
|
except KeyError:
|
|
if throw_exception:
|
|
raise Exception("Could not get AnimateDiffModel class. AnimateDiff-Evolved nodes need to be installed to use SparseCtrl; " + \
|
|
"they are either not installed or are of an insufficient version.")
|
|
return None
|
|
|
|
def get_AnimateDiffInfo(throw_exception=True) -> InterfaceAnimateDiffInfo:
|
|
d = get_dinklink()
|
|
try:
|
|
link_ade = d[DinkLinkConst.ADE]
|
|
return link_ade[DinkLinkConst.ADE_ANIMATEDIFFINFO]
|
|
except KeyError:
|
|
if throw_exception:
|
|
raise Exception("Could not get AnimateDiffInfo class - AnimateDiff-Evolved nodes need to be installed to use SparseCtrl; " + \
|
|
"they are either not installed or are of an insufficient version.")
|
|
return None
|