diff --git a/README.md b/README.md index f6f9f67..7424dec 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,9 @@ |model|Dict2Model|`DICT`, (config_file)|`MODEL`|instantiate a model from given state_dict| |model|StateDictMerger|`DICT`, `DICT`, `FLOAT`|`MODEL`, `CLIP`, `VAE`|merge two or three models| |model|StateDictMergerBlockWeighted|`DICT`, `DICT`|`DICT`|merge two models with per-block weights| +|model|ModelIter|`MODEL`, `MODEL`|`MODEL`|iterate models| +|model|CLIPlIter|`CLIP`, `CLIP`|`CLIP`|iterate CLIPs| +|model|VAElIter|`VAE`, `VAE`|`VAE`|iterate VAEs| ## Output nodes diff --git a/__init__.py b/__init__.py index e541e2d..8bc0b8d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,9 @@ from .randomlatent import RandomLatentImage from .vae import VAEDecodeBatched, VAEEncodeBatched from .sample import KSamplerSetting, KSamplerOverrided, KSamplerXYZ -from .model import StateDictLoader, Dict2Model -from .model_merge import StateDictMerger, StateDictMergerBlockWeighted +from .model.loader import StateDictLoader, Dict2Model +from .model.iter import ModelIter, CLIPIter, VAEIter +from .model.merge import StateDictMerger, StateDictMergerBlockWeighted from .image import GridImage NODE_CLASS_MAPPINGS = { @@ -35,6 +36,15 @@ NODE_CLASS_MAPPINGS = { ## creates model from state_dict loaded by `StateDictLoader` 'Dict2Model': Dict2Model, + ## iterate two models for KSamplerXYZ + 'ModelIter': ModelIter, + + ## iterate two CLIPs for KSamplerXYZ + 'CLIPIter': CLIPIter, + + ## iterate two VAEs for KSamplerXYZ + 'VAEIter': VAEIter, + ## merge two (weighted sum) or three (add difference) state_dict 'StateDictMerger': StateDictMerger, diff --git a/model/iter.py b/model/iter.py new file mode 100644 index 0000000..1d0933c --- /dev/null +++ b/model/iter.py @@ -0,0 +1,143 @@ +from typing import List, Callable +import torch +import tqdm +from comfy.sd import ModelPatcher, CLIP, VAE + +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 + +def iterize_model(model: ModelPatcher) -> List[Callable[[],ModelPatcher]]: + ATTR_NAME = 'iter_fn' + 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]]: + ATTR_NAME = 'iter_fn' + 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]]: + ATTR_NAME = 'iter_fn' + 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) + + +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,) diff --git a/model.py b/model/loader.py similarity index 100% rename from model.py rename to model/loader.py diff --git a/model_merge.py b/model/merge.py similarity index 100% rename from model_merge.py rename to model/merge.py diff --git a/sample.py b/sample.py index 76d568c..16ef757 100644 --- a/sample.py +++ b/sample.py @@ -1,11 +1,12 @@ import re from itertools import product -from typing import Callable, List, Dict, Any, Union, Iterable +from typing import Callable, List, Dict, Any, Union, Tuple, cast import torch import model_management # type: ignore import comfy.samplers from nodes import common_ksampler from comfy.sd import ModelPatcher +from .model.iter import iterize_model, CondForModels re_int = re.compile(r"\s*([+-]?\s*\d+)\s*") re_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*") @@ -46,8 +47,83 @@ def get_cfg(noises: torch.Tensor, latent_image: torch.Tensor, cfgs: List[float]) cf = torch.FloatTensor(cfgs * noises.shape[0]) return torch.cat(ns), torch.cat(lat), cf[...,None,None,None] +def process_cond( + conds: List[List[Union[torch.Tensor,dict]]], + control_nets: Union[list,None], + noise_shape0: int, + device +): + conds_copy = [] + for p in conds: + t: torch.Tensor = p[0] # type: ignore + if t.shape[0] < noise_shape0: + t = torch.cat([t] * noise_shape0) + t = t.to(device) + if control_nets is not None and 'control' in p[1]: + control_nets += [p[1]['control']] # type: ignore + conds_copy += [[t] + p[1:]] + return conds_copy + +def process_cond_for_models( + conds: List[List[Union[torch.Tensor,CondForModels,dict]]], + control_nets: list, + noise_shape0: int, + device +): + assert ( + all(isinstance(p[0], CondForModels) for p in conds) + or not any(isinstance(p[0], CondForModels) for p in conds) + ) + + if isinstance(conds[0][0], CondForModels): + sizes = set( len(cast(CondForModels, p[0]).ex) for p in conds ) + assert len(sizes) == 1, f'number of conditions: {sizes}' + size = sizes.pop() + + # + # conds + # + [ CondForModels, dictA ] + # | .ex + condA for model1 + # | + condA for model2 + # | ... + # | L condA for model{size} + # + [ CondForModels, dictB ] + # | .ex + condB for model1 + # | + condB for model2 + # | ... + # | L condB for model{size} + # ... + # + # vvv + # + # conds + # + [ [ condA_for_model1, dictA ], [ condB_for_model1, dictB ], ... ] + # + [ [ condA_for_model2, dictA ], [ condB_for_model2, dictB ], ... ] + # ... + # + + result = [] + for model_index in range(size): + cs = [] + for cond_for_models in conds: + c: CondForModels = cond_for_models[0] # type: ignore + rest = cond_for_models[1:] + cond = c.ex[model_index] + cs.append([cond, *rest]) + result.append(process_cond( + cs, + control_nets if model_index == 0 else None, + noise_shape0, + device + )) + return result + + else: + return [ process_cond(conds, control_nets, noise_shape0, device) ] + + def common_ksampler_xyz( - model: Union[ModelPatcher,Iterable[ModelPatcher]], + model: ModelPatcher, seed: Union[int,List[int]], steps: Union[int,List[int]], cfg: Union[float,List[float]], @@ -66,9 +142,6 @@ def common_ksampler_xyz( noise_mask = None device = model_management.get_torch_device() - if not isinstance(model, Iterable): - model = (model,) - if not isinstance(seed, list): seed = [seed] @@ -99,41 +172,24 @@ def common_ksampler_xyz( latent_image = latent_image.to(device) cfg_ = cfg_.to(device) - positive_copy = [] - negative_copy = [] - control_nets = [] - for p in positive: - t = p[0] - if t.shape[0] < noise.shape[0]: - t = torch.cat([t] * noise.shape[0]) - t = t.to(device) - if 'control' in p[1]: - control_nets += [p[1]['control']] - positive_copy += [[t] + p[1:]] - for n in negative: - t = n[0] - if t.shape[0] < noise.shape[0]: - t = torch.cat([t] * noise.shape[0]) - t = t.to(device) - if 'control' in n[1]: - control_nets += [n[1]['control']] - negative_copy += [[t] + n[1:]] + positive_copies = process_cond_for_models(positive, control_nets, noise.shape[0], device) + negative_copies = process_cond_for_models(negative, control_nets, noise.shape[0], device) control_net_models = [] for x in control_nets: control_net_models += x.get_control_models() model_management.load_controlnet_gpu(control_net_models) - #samplers: List[comfy.samplers.KSampler] = [] samplers: List[Dict[str,Any]] = [] - for model_, sampler_name_, scheduler_, steps_ in product(model, sampler_name, scheduler, steps): + for (model_index, model_fn), sampler_name_, scheduler_, steps_ in product(enumerate(iterize_model(model)), sampler_name, scheduler, steps): if sampler_name_ not in comfy.samplers.KSampler.SAMPLERS: raise ValueError(f'unknown sampler name: {sampler_name_}') if scheduler_ not in comfy.samplers.KSampler.SCHEDULERS: raise ValueError(f'unknown scheduler name: {scheduler_}') samplers.append(dict( - model=model_, + model_index=model_index, + model=model_fn, steps=steps_, device=device, sampler=sampler_name_, @@ -143,12 +199,16 @@ def common_ksampler_xyz( all_samples: List[torch.Tensor] = [] for sampler_args in samplers: - model_ = sampler_args['model'] + model_ = sampler_args['model']() model_management.load_model_gpu(model_) sampler_args['model'] = model_.model + model_index = sampler_args.pop('model_index') + positive_copy = positive_copies[model_index] + negative_copy = negative_copies[model_index] + sampler = comfy.samplers.KSampler(**sampler_args) - print(f'XYZ sampler={sampler.sampler}/{sampler.scheduler} {sampler.steps}steps') + print(f'XYZ sampler=model@{model_index}/{sampler.sampler}/{sampler.scheduler} {sampler.steps}steps') samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg_, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask) samples = samples.cpu()