146 lines
3.4 KiB
Python
146 lines
3.4 KiB
Python
from typing import List, Callable, Any, Optional
|
|
import torch
|
|
import tqdm
|
|
from comfy.sd import CLIP, VAE
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
class CondForModels(torch.Tensor):
|
|
|
|
@staticmethod
|
|
def __new__(cls, x, ex, *args, **kwargs):
|
|
return super().__new__(cls, x, *args, **kwargs) # type: ignore
|
|
|
|
def __init__(self, x, ex: List[torch.Tensor], *args, **kwargs):
|
|
super().__init__()
|
|
self.ex = ex
|
|
|
|
ATTR_NAME = 'iter_fn'
|
|
|
|
def iterize_model(model: ModelPatcher) -> List[Callable[[],ModelPatcher]]:
|
|
if not hasattr(model, ATTR_NAME):
|
|
setattr(model, ATTR_NAME, [lambda: model])
|
|
return getattr(model, ATTR_NAME)
|
|
|
|
def iterize_clip(clip: CLIP) -> List[Callable[[],CLIP]]:
|
|
if hasattr(clip, ATTR_NAME):
|
|
return getattr(clip, ATTR_NAME)
|
|
|
|
setattr(clip, ATTR_NAME, [lambda: clip])
|
|
|
|
old_encode = CLIP.encode
|
|
|
|
def new_encode(*args, **kwargs):
|
|
xs = []
|
|
clips = getattr(clip, ATTR_NAME)
|
|
for fn in tqdm.tqdm(clips):
|
|
clip_: CLIP = fn()
|
|
if clip_ == clip:
|
|
x = old_encode(clip_, *args, **kwargs)
|
|
else:
|
|
x = clip_.encode(*args, **kwargs)
|
|
if x.dim() == 2:
|
|
x = x.unsqueeze(0)
|
|
xs.append(x)
|
|
return CondForModels(xs[0], xs)
|
|
|
|
clip.encode = new_encode
|
|
|
|
return getattr(clip, ATTR_NAME)
|
|
|
|
def iterize_vae(vae: VAE) -> List[Callable[[],VAE]]:
|
|
if hasattr(vae, ATTR_NAME):
|
|
return getattr(vae, ATTR_NAME)
|
|
|
|
setattr(vae, ATTR_NAME, [lambda: vae])
|
|
|
|
old_decode = VAE.decode
|
|
|
|
def new_decode(*args, **kwargs):
|
|
xs = []
|
|
vaes = getattr(vae, ATTR_NAME)
|
|
for fn in tqdm.tqdm(vaes):
|
|
vae_: VAE = fn()
|
|
if vae_ == vae:
|
|
x = old_decode(vae_, *args, **kwargs)
|
|
else:
|
|
x = vae_.decode(*args, **kwargs)
|
|
if x.dim() == 3:
|
|
x = x.unsqueeze(0)
|
|
xs.append(x)
|
|
return torch.cat(xs)
|
|
|
|
vae.decode = new_decode
|
|
|
|
return getattr(vae, ATTR_NAME)
|
|
|
|
def try_get_iter(obj) -> Optional[List[Callable[[],Any]]]:
|
|
return getattr(obj, ATTR_NAME, None)
|
|
|
|
class ModelIter:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'model1': ('MODEL', ),
|
|
'model2': ('MODEL', )
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('MODEL',)
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'model'
|
|
|
|
def execute(self, model1, model2):
|
|
fns = iterize_model(model1)
|
|
fns.append(lambda: model2)
|
|
return (model1,)
|
|
|
|
|
|
class CLIPIter:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'clip1': ('CLIP', ),
|
|
'clip2': ('CLIP', )
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('CLIP',)
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'model'
|
|
|
|
def execute(self, clip1, clip2):
|
|
fns = iterize_clip(clip1)
|
|
fns.append(lambda: clip2)
|
|
return (clip1,)
|
|
|
|
|
|
class VAEIter:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'vae1': ('VAE', ),
|
|
'vae2': ('VAE', )
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('VAE',)
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'model'
|
|
|
|
def execute(self, vae1, vae2):
|
|
fns = iterize_vae(vae1)
|
|
fns.append(lambda: vae2)
|
|
return (vae1,)
|