From 5deb00a64045912799d1c1b2af066dc22f484745 Mon Sep 17 00:00:00 2001 From: logtd Date: Mon, 26 Feb 2024 12:08:45 -0700 Subject: [PATCH] use model_function_wrapper --- nodes/flatten_ksampler_node.py | 9 +- nodes/flatten_unsampler_node.py | 161 ++----------------------------- nodes/load_flatten_model_node.py | 15 +++ 3 files changed, 31 insertions(+), 154 deletions(-) diff --git a/nodes/flatten_ksampler_node.py b/nodes/flatten_ksampler_node.py index 7b39130..7d5b185 100644 --- a/nodes/flatten_ksampler_node.py +++ b/nodes/flatten_ksampler_node.py @@ -39,8 +39,13 @@ class KSamplerFlattenNode: latent = latent_image latent_image = latent["samples"] - latent_image = rearrange(latent_image, "(b f) c h w -> b c f h w", b=1) + # latent_image = rearrange(latent_image, "(b f) c h w -> b c f h w", b=1) + transformer_options = {} + if 'transformer_options' in model.model_options: + transformer_options = model.model_options['transformer_options'] + transformer_options = { + **transformer_options, 'flatten': { 'trajs': trajectories, 'old_qk': old_qk @@ -107,7 +112,7 @@ class KSamplerFlattenNode: comfy.sample.cleanup_additional_models(models) out = latent.copy() - out["samples"] = rearrange(samples, 'b c f h w -> (b f) c h w') + out["samples"] = samples return (out, ) diff --git a/nodes/flatten_unsampler_node.py b/nodes/flatten_unsampler_node.py index 83515d5..7a8eadd 100644 --- a/nodes/flatten_unsampler_node.py +++ b/nodes/flatten_unsampler_node.py @@ -4,155 +4,6 @@ import torch class UnsamplerFlattenNode: - @classmethod - def INPUT_TYPES(s): - return {"required": - {"model": ("MODEL",), - "add_noise": (["enable", "disable"], ), - "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), - "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), - "old_qk": ("INT", {"default": 0, "min": 0, "max": 1}), - "trajectories": ("TRAJECTORY", ), - "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}), - "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), - "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), - "positive": ("CONDITIONING", ), - "negative": ("CONDITIONING", ), - "latent_image": ("LATENT", ), - "end_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}), - "return_with_leftover_noise": (["disable", "enable"], ), - } - } - - RETURN_TYPES = ("LATENT",) - FUNCTION = "sample" - - CATEGORY = "sampling" - - def sample(self, model, add_noise, noise_seed, steps, old_qk, trajectories, cfg, sampler_name, scheduler, positive, negative, latent_image, end_at_step, return_with_leftover_noise, denoise=1.0): - # PREPARTION - latent = latent_image - latent_image = latent["samples"] - latent_image = rearrange(latent_image, "(b f) c h w -> b c f h w", b=1) - transformer_options = { - 'flatten': { - 'trajs': trajectories, - 'old_qk': old_qk - } - } - model.model_options['transformer_options'] = transformer_options - # UNSAMPLER - device = comfy.model_management.get_torch_device() - - end_at_step = min(end_at_step, steps-1) - end_at_step = steps - end_at_step - - noise = torch.zeros(latent_image.size( - ), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu") - noise_mask = None - if "noise_mask" in latent: - noise_mask = comfy.sample.prepare_mask( - latent["noise_mask"], noise.shape, device) - - real_model = None - real_model = model.model - - noise = noise.to(device) - latent_image = latent_image.to(device) - - positive = comfy.sample.convert_cond(positive) - negative = comfy.sample.convert_cond(negative) - - models, inference_memory = comfy.sample.get_additional_models( - positive, negative, model.model_dtype()) - - comfy.model_management.load_models_gpu( - [model] + models, model.memory_required(noise.shape) + inference_memory) - - sampler = comfy.samplers.KSampler(real_model, steps=steps, device=device, sampler=sampler_name, - scheduler=scheduler, denoise=1.0, model_options=model.model_options) - - sigmas = sampler.sigmas.flip(0) + 0.0001 - - injection_dict = self._get_blank_injection_dict() - pbar = comfy.utils.ProgressBar(steps) - - def callback(step, x0, x, total_steps): - self._update_injections(model, injection_dict) - pbar.update_absolute(step + 1, total_steps) - - self._clear_injections(model) - latent_image = sampler.sample(noise, positive, negative, cfg=cfg, latent_image=latent_image, force_full_denoise=False, - denoise_mask=noise_mask, sigmas=sigmas, start_step=0, last_step=end_at_step, callback=callback) - self._clear_injections(model) - - del noise - latent_image = latent_image.cpu() - - comfy.sample.cleanup_additional_models(models) - - latent_image = rearrange(latent_image, 'b c f h w -> (b f) c h w') - - return ({'samples': latent_image, 'injections': injection_dict},) - - def _get_blank_injection_dict(self): - return { - 'features0': [], - 'features1': [], - 'features2': [], - 'q4': [], - 'k4': [], - 'q5': [], - 'k5': [], - 'q6': [], - 'k6': [], - 'q7': [], - 'k7': [], - 'q8': [], - 'k8': [], - 'q9': [], - 'k9': [] - } - - def _clear_injections(self, model): - model = model.model.diffusion_model - res_attn_dict = {1: [0, 1], 2: [0]} - for res in res_attn_dict: - for block in res_attn_dict[res]: - model.output_blocks[3*res+block][0].out_layers_features = None - attn_res_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0]} - for attn in attn_res_dict: - for block in attn_res_dict[attn]: - module = model.output_blocks[3*attn + - block][1].transformer_blocks[0].attn1 - module.q = None - module.k = None - module.inject_q = None - module.inject_k = None - - def _update_injections(self, model, injection): - model = model.model.diffusion_model - - res_dict = {1: [0, 1], 2: [0]} - res_idx = 0 - for res in res_dict: - for block in res_dict[res]: - injection[f'features{res_idx}'].append( - model.output_blocks[3*res+block][0].out_layers_features.cpu()) - res_idx += 1 - - attn_dict = {1: [1, 2], 2: [0, 1, 2], 3: [0]} - attn_idx = 4 - for attn in attn_dict: - for block in attn_dict[attn]: - module = model.output_blocks[3*attn + - block][1].transformer_blocks[0].attn1 - injection[f'q{attn_idx}'].append(module.q.cpu()) - injection[f'k{attn_idx}'].append(module.k.cpu()) - attn_idx += 1 - - -class Unsampler2: @classmethod def INPUT_TYPES(s): return {"required": @@ -167,6 +18,7 @@ class Unsampler2: "negative": ("CONDITIONING", ), "latent_image": ("LATENT", ), "trajectories": ("TRAJECTORY", ), + "old_qk": ("INT", {"default": 0, "min": 0, "max": 1}), }} RETURN_TYPES = ("LATENT",) @@ -174,15 +26,20 @@ class Unsampler2: CATEGORY = "sampling" - def unsampler(self, model, cfg, sampler_name, steps, end_at_step, scheduler, normalize, positive, negative, latent_image, trajectories): + def unsampler(self, model, cfg, sampler_name, steps, end_at_step, scheduler, normalize, positive, negative, latent_image, trajectories, old_qk): # PREPARTION latent = latent_image latent_image = latent["samples"] - latent_image = rearrange(latent_image, "(b f) c h w -> b c f h w", b=1) + # latent_image = rearrange(latent_image, "(b f) c h w -> b c f h w", b=1) + transformer_options = {} + if 'transformer_options' in model.model_options: + transformer_options = model.model_options['transformer_options'] + transformer_options = { + **transformer_options, 'flatten': { 'trajs': trajectories, - 'old_qk': 0 + 'old_qk': old_qk } } model.model_options['transformer_options'] = transformer_options diff --git a/nodes/load_flatten_model_node.py b/nodes/load_flatten_model_node.py index 9c83b5b..2bcec4b 100644 --- a/nodes/load_flatten_model_node.py +++ b/nodes/load_flatten_model_node.py @@ -1,3 +1,4 @@ +from einops import rearrange import comfy.sd import comfy.model_base import folder_paths @@ -27,4 +28,18 @@ class FlattenCheckpointLoaderNode: out = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings")) comfy.model_base.BaseModel = original_base + + def model_function_wrapper(apply_model_func, apply_params): + len_conds = len(apply_params['cond_or_uncond']) + input_x = apply_params['input'] + input_x = rearrange( + input_x, "(b f) c h w -> b c f h w", b=len_conds) + timestep_ = apply_params['timestep'] + model_out = apply_model_func(input_x, timestep_, **apply_params) + model_out = rearrange(model_out, 'b c f h w -> (b f) c h w') + return model_out + + model = out[0] + model.model_options['model_function_wrapper'] = model_function_wrapper + return out[:3]