diff --git a/__init__.py b/__init__.py index 424efd4..e33b1a7 100644 --- a/__init__.py +++ b/__init__.py @@ -1,12 +1,15 @@ from .nodes.apply_ref_unet_node import ApplyRefUNetNode -from .nodes.ref_apply_node import RefApplyNode +from .nodes.apply_ref_content_node import ApplyRefContentNode +from .nodes.apply_ref_style_node import ApplyRefStyleNode NODE_CLASS_MAPPINGS = { "ApplyRefUNetNode": ApplyRefUNetNode, - "RefApplyNode": RefApplyNode, + "ApplyRefContentNode": ApplyRefContentNode, + "ApplyRefStyleNode": ApplyRefStyleNode, } NODE_DISPLAY_NAME_MAPPINGS = { "ApplyRefUNetNode": "Apply Ref UNet", - "RefApplyNode": "Apply Reference", + "ApplyRefContentNode": "Apply Ref Content", + "ApplyRefStyleNode": "Apply Ref Style", } diff --git a/example_workflow/example_ref_workflow.json b/example_workflow/example_ref_workflow.json new file mode 100644 index 0000000..37fbf73 --- /dev/null +++ b/example_workflow/example_ref_workflow.json @@ -0,0 +1,964 @@ +{ + "last_node_id": 23, + "last_link_id": 38, + "nodes": [ + { + "id": 3, + "type": "CLIPTextEncode", + "pos": [ + 1, + 911 + ], + "size": { + "0": 362.39056396484375, + "1": 76 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 2 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 12 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 6, + "type": "EmptyLatentImage", + "pos": [ + 22, + 1055 + ], + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 13 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 512, + 512, + 1 + ] + }, + { + "id": 12, + "type": "VAEDecode", + "pos": [ + 786, + 563 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 15 + }, + { + "name": "vae", + "type": "VAE", + "link": 14 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 16 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 8, + "type": "ApplyRefUNetNode", + "pos": [ + -13, + 365 + ], + "size": { + "0": 210, + "1": 58 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 5 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 29 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ApplyRefUNetNode" + }, + "widgets_values": [ + "FULL" + ] + }, + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [ + -474, + 700 + ], + "size": { + "0": 315, + "1": 98 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 5 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 1, + 2, + 18, + 20 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 8, + 14, + 32 + ], + "shape": 3, + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "photon_v1.safetensors" + ] + }, + { + "id": 22, + "type": "ImageScale", + "pos": [ + 994, + -239 + ], + "size": { + "0": 210, + "1": 130 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 31 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 33 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "nearest-exact", + 512, + 512, + "disabled" + ] + }, + { + "id": 7, + "type": "ImageScale", + "pos": [ + 203, + -202 + ], + "size": { + "0": 210, + "1": 130 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 4 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 7 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ImageScale" + }, + "widgets_values": [ + "nearest-exact", + 512, + 512, + "disabled" + ] + }, + { + "id": 23, + "type": "VAEEncode", + "pos": [ + 723, + 59 + ], + "size": { + "0": 140, + "1": 46 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 33 + }, + { + "name": "vae", + "type": "VAE", + "link": 32 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 35 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 10, + "type": "VAEEncode", + "pos": [ + -27, + 97 + ], + "size": { + "0": 140, + "1": 46 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 28 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncode" + } + }, + { + "id": 15, + "type": "CLIPTextEncode", + "pos": [ + -46, + 617 + ], + "size": { + "0": 372.4996337890625, + "1": 85.3608627319336 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 18 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 37 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 5, + "type": "LoadImage", + "pos": [ + -109, + -305 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "gorilla.jpeg", + "image" + ] + }, + { + "id": 14, + "type": "SaveImage", + "pos": [ + 701, + 916 + ], + "size": { + "0": 315, + "1": 270 + }, + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 16 + } + ], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 18, + "type": "ApplyRefContentNode", + "pos": [ + 335, + 216 + ], + "size": { + "0": 210, + "1": 174 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 29 + }, + { + "name": "ref_latents", + "type": "LATENT", + "link": 28 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 30 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ApplyRefContentNode" + }, + "widgets_values": [ + true, + 0.8, + 1, + 0, + 0.8 + ] + }, + { + "id": 21, + "type": "LoadImage", + "pos": [ + 621, + -341 + ], + "size": { + "0": 315, + "1": 314 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 31 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "tigerstripes.jpeg", + "image" + ] + }, + { + "id": 16, + "type": "CLIPTextEncode", + "pos": [ + -37, + 473 + ], + "size": { + "0": 372.4996337890625, + "1": 85.3608627319336 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 20 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 36 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "orange and black stripes" + ] + }, + { + "id": 20, + "type": "ApplyRefStyleNode", + "pos": [ + 643, + 250 + ], + "size": { + "0": 210, + "1": 190 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 30 + }, + { + "name": "ref_latents", + "type": "LATENT", + "link": 35 + }, + { + "name": "style_positive", + "type": "CONDITIONING", + "link": 36 + }, + { + "name": "style_negative", + "type": "CONDITIONING", + "link": 37 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 38 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "ApplyRefStyleNode" + }, + "widgets_values": [ + true, + 2, + 0.2, + 0.7000000000000001 + ] + }, + { + "id": 11, + "type": "KSampler", + "pos": [ + 1106, + 666 + ], + "size": { + "0": 315, + "1": 262 + }, + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 38 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 11 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 12 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 13 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 15 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 1, + "fixed", + 20, + 7, + "dpmpp_2s_ancestral", + "normal", + 0.65 + ] + }, + { + "id": 2, + "type": "CLIPTextEncode", + "pos": [ + 0, + 773 + ], + "size": { + "0": 372.4996337890625, + "1": 85.3608627319336 + }, + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 1 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 11 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "orange gorilla stripes" + ] + } + ], + "links": [ + [ + 1, + 1, + 1, + 2, + 0, + "CLIP" + ], + [ + 2, + 1, + 1, + 3, + 0, + "CLIP" + ], + [ + 4, + 5, + 0, + 7, + 0, + "IMAGE" + ], + [ + 5, + 1, + 0, + 8, + 0, + "MODEL" + ], + [ + 7, + 7, + 0, + 10, + 0, + "IMAGE" + ], + [ + 8, + 1, + 2, + 10, + 1, + "VAE" + ], + [ + 11, + 2, + 0, + 11, + 1, + "CONDITIONING" + ], + [ + 12, + 3, + 0, + 11, + 2, + "CONDITIONING" + ], + [ + 13, + 6, + 0, + 11, + 3, + "LATENT" + ], + [ + 14, + 1, + 2, + 12, + 1, + "VAE" + ], + [ + 15, + 11, + 0, + 12, + 0, + "LATENT" + ], + [ + 16, + 12, + 0, + 14, + 0, + "IMAGE" + ], + [ + 18, + 1, + 1, + 15, + 0, + "CLIP" + ], + [ + 20, + 1, + 1, + 16, + 0, + "CLIP" + ], + [ + 28, + 10, + 0, + 18, + 1, + "LATENT" + ], + [ + 29, + 8, + 0, + 18, + 0, + "MODEL" + ], + [ + 30, + 18, + 0, + 20, + 0, + "MODEL" + ], + [ + 31, + 21, + 0, + 22, + 0, + "IMAGE" + ], + [ + 32, + 1, + 2, + 23, + 1, + "VAE" + ], + [ + 33, + 22, + 0, + 23, + 0, + "IMAGE" + ], + [ + 35, + 23, + 0, + 20, + 1, + "LATENT" + ], + [ + 36, + 16, + 0, + 20, + 2, + "CONDITIONING" + ], + [ + 37, + 15, + 0, + 20, + 3, + "CONDITIONING" + ], + [ + 38, + 20, + 0, + 11, + 0, + "MODEL" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file diff --git a/modules/attn_wrapper.py b/modules/attn_wrapper.py index 36046de..d89aee0 100644 --- a/modules/attn_wrapper.py +++ b/modules/attn_wrapper.py @@ -1,17 +1,51 @@ import torch +from comfy.ldm.modules.attention import optimized_attention, optimized_attention_masked +from comfy.ldm.util import default + from .ref_controller import RefMode + + def get_attn_wrapper(cls): class RefAttention(cls): ref_mode = RefMode.OFF - ref_bank = None + normal_mode = RefMode.OFF + kv_mode = RefMode.OFF + normal_bank = None + k_bank = None + v_bank = None is_cfg = True - style_fidelity = 0.5 + content_fidelity = 0.5 ref_attn_weight = 1.0 attention_bound = 1.0 hidden_states = None + block_type = None + block_idx = None + + def _attention_mechanism(self, q, k, v, mask): + if mask is None: + out = optimized_attention(q, k, v, self.heads) + else: + out = optimized_attention_masked(q, k, v, self.heads, mask) + return self.to_out(out) + + def _get_qkv(self, x, context=None, value=None): + q = self.to_q(x) + context = default(context, x) + k = self.to_k(context) + if value is not None: + v = self.to_v(value) + del value + else: + v = self.to_v(context) + return q,k,v + + def clean(self): + self.normal_bank = None + self.k_bank = None + self.v_bank = None def forward(self, *args, **kwargs): if self.ref_mode == RefMode.OFF: @@ -20,25 +54,45 @@ def get_attn_wrapper(cls): return self._ref_forward(*args, **kwargs) def _ref_forward(self, norm_hidden_states, context, **kwargs): - if self.ref_mode == RefMode.WRITE: - self.ref_bank = norm_hidden_states.detach().clone() - return super().forward(norm_hidden_states, context, **kwargs) - if self.ref_mode == RefMode.READ: - if self.attention_bound > self.ref_attn_weight: - bank = self.ref_bank + context = default(context, norm_hidden_states) + value = kwargs.get('value', None) + value = default(value, context) - self.ref_bank = None - bank_hidden_states = torch.cat([norm_hidden_states, bank], dim=1) - attn_output_uc = ( - super().forward( - norm_hidden_states, - bank_hidden_states, - **kwargs - ) - ) - attn_output = attn_output_uc - if self.is_cfg and self.style_fidelity > 0: - attn_output_c = attn_output_uc.clone() + if self.ref_mode == RefMode.WRITE: + if self.normal_mode == RefMode.WRITE and self.attention_bound > self.ref_attn_weight: + self.normal_bank = norm_hidden_states.detach().clone() + if self.kv_mode == RefMode.WRITE: + q,k,v = self._get_qkv(norm_hidden_states, context, **kwargs) + self.k_bank = k.detach().clone() + self.v_bank = v.detach().clone() + mask = kwargs['mask'] if 'mask' in kwargs else None + return self._attention_mechanism(q,k,v,mask=mask) + else: + return super().forward(norm_hidden_states, context, **kwargs) + if self.ref_mode == RefMode.READ: + normal_bank = None + if self.normal_mode == RefMode.READ and self.attention_bound > self.ref_attn_weight: + normal_bank = self.normal_bank + self.normal_bank = None + + if normal_bank is not None: + context = torch.cat([context, normal_bank], dim=1) + value = torch.cat([value, normal_bank], dim=1) + + + q, k, v = self._get_qkv(norm_hidden_states, context, value) + + if self.kv_mode == RefMode.READ: + k = self.k_bank + v = self.v_bank + self.k_bank = None + self.v_bank = None + + mask = kwargs['mask'] if 'mask' in kwargs else None + attn_output = self._attention_mechanism(q,k,v, mask=mask) + + if normal_bank is not None and self.is_cfg and self.content_fidelity > 0: + attn_output_c = attn_output.clone() uc_mask = torch.Tensor( [1] * (norm_hidden_states.shape[0] // 2) + [0] * (norm_hidden_states.shape[0] // 2) @@ -49,9 +103,7 @@ def get_attn_wrapper(cls): **kwargs ) attn_output_c[uc_mask] = attn_sub - attn_output = self.style_fidelity * attn_output_c + (1.0 - self.style_fidelity) * attn_output_uc - else: - attn_output = super().forward(norm_hidden_states, context, **kwargs) + return self.content_fidelity * attn_output_c + (1.0 - self.content_fidelity) * attn_output return attn_output return RefAttention diff --git a/modules/block_type.py b/modules/block_type.py new file mode 100644 index 0000000..791fca2 --- /dev/null +++ b/modules/block_type.py @@ -0,0 +1,4 @@ +import enum + +class BlockType(enum.Enum): + OUTPUT = 'OUTPUT' diff --git a/modules/ref_controller.py b/modules/ref_controller.py index 11d1dd1..b39451b 100644 --- a/modules/ref_controller.py +++ b/modules/ref_controller.py @@ -1,5 +1,7 @@ import enum +from .block_type import BlockType + class RefMode(enum.Enum): OFF = 1 @@ -17,43 +19,41 @@ class RefController: self.modules.append(module) def set_mode(self, mode: RefMode): - print('------- CTRL SETTING MODE TO', mode, 'FOR', len(self.modules)) for module in self.modules: module.ref_mode = mode + if mode == RefMode.OFF: + self.set_kv_mode(mode) + self.set_normal_mode(mode) + + def set_normal_mode(self, mode: RefMode): + for module in self.modules: + module.normal_mode = mode + + def set_kv_mode(self, mode: RefMode, count = None): + for module in self.modules: + module.kv_mode = RefMode.OFF + if mode == RefMode.OFF or count is None: + count = len(self.modules) + output_modules = list(filter(lambda m: m.block_type == BlockType.OUTPUT, self.modules)) + output_modules = sorted(output_modules, key=lambda m: -m.block_idx)[:count] + for module in output_modules: + module.kv_mode = mode def clear_modules(self): - print('------- CTRL CLEARING MODULES FOR', len(self.modules)) for module in self.modules: - module.ref_bank = None + module.clean() self.set_mode(RefMode.OFF) + self.set_kv_mode(RefMode.OFF, len(self.modules)) + self.set_normal_mode(RefMode.OFF) def set_cfg(self, cfg: bool): - print('------- CTRL SETTING CFG TO', cfg, 'FOR', len(self.modules)) for module in self.modules: module.is_cfg = cfg def set_attention_bound(self, attention_bound: float): - print('------- CTRL SETTING attention bound TO', attention_bound, 'FOR', len(self.modules)) for module in self.modules: module.attention_bound = attention_bound - def set_style_fidelity(self, style_fidelity: float): - print('------- CTRL SETTING style fidelity TO', style_fidelity, 'FOR', len(self.modules)) + def set_content_fidelity(self, content_fidelity: float): for module in self.modules: - module.style_fidelity = style_fidelity - - def add_temporal_transformer(self, module): - self.temp_modules.append(module) - - def get_default_temporal_values(self): - if len(self.temp_modules) > 0: - self.default_full_length = self.temp_modules[0].full_length - self.default_video_length = self.temp_modules[0].video_length - - def set_temporal_values_to(self, full_length, video_length): - for module in self.temp_modules: - module.full_length = full_length - module.video_length = video_length - - def set_temporal_values_to_default(self): - self.set_temporal_values_to(self.default_full_length, self.default_video_length) \ No newline at end of file + module.content_fidelity = content_fidelity diff --git a/modules/unet_wrapper.py b/modules/unet_wrapper.py index 58db8dd..51497d5 100644 --- a/modules/unet_wrapper.py +++ b/modules/unet_wrapper.py @@ -1,46 +1,97 @@ import torch from .ref_controller import RefController, RefMode -from ..utils.ref_config import RefConfig +from ..utils.ref_content_config import RefContentConfig +from ..utils.ref_style_config import RefStyleConfig from ..utils.noise_utils import add_noise def get_unet_wrapper(cls, ref_controller: RefController): class RefUNet(cls): is_ref = True - def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs): - if 'ref_config' not in transformer_options: - return super().forward(x, - timesteps=timesteps, - context=context, - y=y, - control=control, - transformer_options=transformer_options, - **kwargs) - - ref_config: RefConfig = transformer_options['ref_config'] + + def _ref_content_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs): + ref_config: RefContentConfig = transformer_options.get('ref_content_config', None) + if ref_config is None: + return RefMode.OFF ref_latent = ref_config.ref_latent ref_latent = torch.cat([ref_latent]*2) - sigma = ref_config.sampling.sigma(timesteps) + start_sigma = ref_config.sampling.percent_to_sigma(ref_config.content_start_percent) + end_sigma = ref_config.sampling.percent_to_sigma(ref_config.content_end_percent) + if not (start_sigma >= sigma[0] >= end_sigma): + return RefMode.OFF ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0]).to(x.device).to(x.dtype) - ref_controller.set_cfg(True) #TODO - ref_controller.set_attention_bound(ref_config.attention_bound) - ref_controller.set_style_fidelity(ref_config.style_fidelity) + ref_controller.set_cfg(len(transformer_options['cond_or_uncond']) > 1) # TODO + ref_controller.set_attention_bound(ref_config.content_attention_bound) + ref_controller.set_content_fidelity(ref_config.content_fideltiy) + ref_controller.set_normal_mode(RefMode.WRITE) + super().forward(ref_latent_noised, + timesteps=timesteps, + context=context, + y=y, + control=None, + transformer_options=transformer_options, + **kwargs) + ref_controller.set_normal_mode(RefMode.OFF) + + return RefMode.READ + + def _ref_style_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs): + ref_config: RefStyleConfig = transformer_options.get('ref_style_config', None) + if ref_config is None: + return RefMode.OFF, None + ref_latent = ref_config.style_latent + ref_latent = torch.cat([ref_latent]*2) + sigma = ref_config.sampling.sigma(timesteps) + start_sigma = ref_config.sampling.percent_to_sigma(ref_config.style_start_percent) + end_sigma = ref_config.sampling.percent_to_sigma(ref_config.style_end_percent) + if not (start_sigma >= sigma[0] >= end_sigma): + return RefMode.OFF, None + + ref_latent_noised = add_noise(ref_latent, torch.randn_like(ref_latent), sigma[0]).to(x.device).to(x.dtype) + + ref_controller.set_kv_mode(RefMode.WRITE, ref_config.style_count) + super().forward(ref_latent_noised, + timesteps=timesteps, + context=ref_config.style_prompt.to(ref_latent_noised.device).half(), + y=y, + control=None, + transformer_options=transformer_options, + **kwargs) + ref_controller.set_kv_mode(RefMode.OFF) + + return RefMode.READ, ref_config.style_count + + + def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs): + ref_controller.set_mode(RefMode.OFF) ref_controller.clear_modules() try: ref_controller.set_mode(RefMode.WRITE) - super().forward(ref_latent_noised, - timesteps=timesteps, - context=context, - y=y, - control=None, - transformer_options=transformer_options, - **kwargs) - + normal_read = self._ref_content_forward( + x, + timesteps=timesteps, + context=context, + y=y, + control=control, + transformer_options=transformer_options, + **kwargs + ) + kv_read, kv_count = self._ref_style_forward( + x, + timesteps=timesteps, + context=context, + y=y, + control=control, + transformer_options=transformer_options, + **kwargs + ) ref_controller.set_mode(RefMode.READ) + ref_controller.set_normal_mode(normal_read) + ref_controller.set_kv_mode(kv_read, kv_count) output = super().forward(x, timesteps=timesteps, context=context, @@ -50,9 +101,6 @@ def get_unet_wrapper(cls, ref_controller: RefController): **kwargs) ref_controller.set_mode(RefMode.OFF) return output - # except Exception as e: - # print('issue with ref') - # raise e finally: ref_controller.clear_modules() diff --git a/nodes/apply_ref_content_node.py b/nodes/apply_ref_content_node.py new file mode 100644 index 0000000..0816403 --- /dev/null +++ b/nodes/apply_ref_content_node.py @@ -0,0 +1,52 @@ +from ..utils.ref_sampling_utils import prepare_ref_latents +from ..utils.ref_content_config import RefContentConfig + + +class ApplyRefContentNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "ref_latents": ("LATENT",), + "enabled": ("BOOLEAN", {"default": True}), + "fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), + "attention_bound": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + + CATEGORY = "flow" + + def apply(self, + model, + ref_latents, + enabled, + fidelity, + attention_bound, + start_percent, + end_percent): + transformer_options = model.model_options.get( + 'transformer_options', {}) + model.model_options['transformer_options'] = transformer_options + if not enabled: + if 'ref_content_config' in transformer_options: + del transformer_options['ref_content_config'] + return (model, ) + + ref_latents = prepare_ref_latents(model, ref_latents) + sampling = model.model.model_sampling + + ref_config = RefContentConfig( + sampling, + ref_latents, + enabled, + fidelity, + attention_bound, + start_percent, + end_percent) + + transformer_options['ref_content_config'] = ref_config + + return (model, ) diff --git a/nodes/apply_ref_style_node.py b/nodes/apply_ref_style_node.py new file mode 100644 index 0000000..2d0cc5e --- /dev/null +++ b/nodes/apply_ref_style_node.py @@ -0,0 +1,57 @@ +import torch +from ..utils.ref_sampling_utils import prepare_ref_latents +from ..utils.ref_style_config import RefStyleConfig + + +class ApplyRefStyleNode: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL",), + "ref_latents": ("LATENT",), + "enabled": ("BOOLEAN", {"default": True}), + "style_positive": ("CONDITIONING",), + "style_negative": ("CONDITIONING",), + "attention_count": ("INT", {"default": 4, "min": 0, "max": 10, "step": 1}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), + "end_percent": ("FLOAT", {"default": 0.70, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), + }} + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + + CATEGORY = "flow" + + def apply(self, + model, + ref_latents, + enabled, + style_positive, + style_negative, + attention_count, + start_percent, + end_percent): + transformer_options = model.model_options.get( + 'transformer_options', {}) + model.model_options['transformer_options'] = transformer_options + if not enabled: + if 'ref_style_config' in transformer_options: + del transformer_options['ref_style_config'] + return (model, ) + + ref_latents = prepare_ref_latents(model, ref_latents) + sampling = model.model.model_sampling + + style_prompt = torch.cat([style_negative[0][0], style_positive[0][0]]) + + ref_config = RefStyleConfig( + sampling, + ref_latents, + enabled, + style_prompt, + attention_count, + start_percent, + end_percent) + + transformer_options['ref_style_config'] = ref_config + + return (model, ) diff --git a/nodes/ref_apply_node.py b/nodes/ref_apply_node.py deleted file mode 100644 index bdcfad3..0000000 --- a/nodes/ref_apply_node.py +++ /dev/null @@ -1,31 +0,0 @@ -from ..utils.ref_sampling_utils import prepare_ref_latents -from ..utils.ref_config import RefConfig - - -class RefApplyNode: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "model": ("MODEL",), - "ref_latents": ("LATENT",), - "style_fidelity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), - "attention_bound": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), - }} - RETURN_TYPES = ("MODEL",) - FUNCTION = "apply" - - CATEGORY = "flow" - - def apply(self, model, ref_latents, style_fidelity, attention_bound): - - ref_latents = prepare_ref_latents(model, ref_latents) - sampling = model.model.model_sampling - ref_config = RefConfig(sampling, ref_latents, - style_fidelity, attention_bound) - - transformer_options = model.model_options.get( - 'transformer_options', {}) - model.model_options['transformer_options'] = transformer_options - transformer_options['ref_config'] = ref_config - - return (model, ) diff --git a/utils/module_utils.py b/utils/module_utils.py index 16bcb2f..55d1bc3 100644 --- a/utils/module_utils.py +++ b/utils/module_utils.py @@ -3,6 +3,7 @@ import torch from ..modules.attn_wrapper import get_attn_wrapper from ..modules.unet_wrapper import get_unet_wrapper from ..modules.ref_controller import RefController +from ..modules.block_type import BlockType def isinstance_str(x: object, cls_name: str): @@ -53,7 +54,7 @@ def setup_ref_attn(model, attn_type): model = model.model.diffusion_model block_modules = list(filter(is_named_module_transformer_block, model.named_modules())) block_modules = list(map(lambda x: x[1], block_modules)) - block_modules = sorted(block_modules, key=sort_block) + ref_controller = RefController() @@ -65,6 +66,16 @@ def setup_ref_attn(model, attn_type): attn.ref_attn_weight = float(i) / float(len(block_modules)) ref_controller.add_module(attn) + block_modules = sorted(block_modules, key=sort_block) + for i, block_module in enumerate(block_modules): + attn.ref_attn_weight = float(i) / float(len(block_modules)) + + output_modules = list(filter(is_named_module_transformer_block, model.output_blocks.named_modules())) + output_modules = list(map(lambda x: x[1], output_modules)) + for i, block_module in enumerate(output_modules): + attn = block_module.attn1 + attn.block_type = BlockType.OUTPUT + attn.block_idx = i temp_transformers = list(filter(lambda x: isinstance_str(x, 'TemporalTransformer3DModel'), model.named_modules())) diff --git a/utils/ref_config.py b/utils/ref_config.py deleted file mode 100644 index fa4c50b..0000000 --- a/utils/ref_config.py +++ /dev/null @@ -1,9 +0,0 @@ - - - -class RefConfig: - def __init__(self, sampling, ref_latent, style_fidelity, attention_bound): - self.sampling = sampling - self.ref_latent = ref_latent - self.style_fidelity = style_fidelity - self.attention_bound = attention_bound \ No newline at end of file diff --git a/utils/ref_content_config.py b/utils/ref_content_config.py new file mode 100644 index 0000000..f588e77 --- /dev/null +++ b/utils/ref_content_config.py @@ -0,0 +1,18 @@ + +class RefContentConfig: + def __init__(self, + sampling, + ref_latent, + content_transfer, + content_fideltiy, + content_attention_bound, + content_start_percent, + content_end_percent, + ): + self.sampling = sampling + self.ref_latent = ref_latent + self.content_fideltiy = content_fideltiy + self.content_attention_bound = content_attention_bound + self.content_transfer = content_transfer + self.content_start_percent = content_start_percent + self.content_end_percent = content_end_percent \ No newline at end of file diff --git a/utils/ref_sampling_utils.py b/utils/ref_sampling_utils.py index be9f3da..6afb727 100644 --- a/utils/ref_sampling_utils.py +++ b/utils/ref_sampling_utils.py @@ -1,11 +1,5 @@ -import torch - import comfy.model_management -from .ref_config import RefConfig -from ..modules.ref_controller import RefController, RefMode -from ..utils.noise_utils import add_noise, add_noise_test - def get_base_model(model): base_model = model @@ -21,73 +15,3 @@ def prepare_ref_latents(model, ref_latent): device = comfy.model_management.get_torch_device() ref_latent = ref_latent.to(device) return ref_latent - - -def create_model_wrapper(model, sigmas, ref_latent, extra_args, ref_config:RefConfig): - # extra_args: ['cond', 'uncond', 'cond_scale', 'model_options', 'seed', 'denoise_mask'] - ref_latent = prepare_ref_latents(model, ref_latent) - - # Get ref_controller - model_options = extra_args.get('model_options', {}) - ref_controller: RefController = model_options.get('ref_controller', {}) - - ref_conds = [] - for cond in extra_args['cond']: - cond = {**cond} - if 'control' in cond: - del cond['control'] - ref_conds.append(cond) - - ref_unconds = [] - for cond in extra_args['uncond']: - cond = {**cond} - if 'control' in cond: - del cond['control'] - ref_unconds.append(cond) - - # Setup ref - ref_controller.clear_modules() - is_cfg = extra_args.get('cond_scale', 7) > 1 - ref_controller.set_cfg(is_cfg) - ref_controller.set_style_fidelity(ref_config.style_fidelity) - ref_controller.set_attention_bound(ref_config.attention_bound) - - @torch.no_grad() - def model_sample(latents, sigma, **kwargs): - ref_controller.get_default_temporal_values() - # Do Ref Sampling - ref_controller.set_temporal_values_to(1, 1) - ref_controller.set_mode(RefMode.WRITE) - # ref_latent_noised = add_noise(ref_latent.clone(), torch.randn_like(ref_latent), sigma[0]) - ref_latent_noised = ref_latent.clone() - ref_sigma = torch.tensor([sigma[0]]*len(ref_latent)).to(sigmas.device) - model(ref_latent_noised, ref_sigma, **{**kwargs, 'cond': ref_conds, 'uncond': ref_unconds}) - # Do Latent Sampling - ref_controller.set_temporal_values_to_default() - ref_controller.set_mode(RefMode.READ) - latents = model(latents, sigma, **kwargs) - - # Clear refs - ref_controller.clear_modules() - ref_controller.set_mode(RefMode.OFF) - return latents - - - class RefModelWrapper: - inner_model = model.inner_model - def __call__(self, *args, **kwargs): - return model_sample(*args, **kwargs) - - return RefModelWrapper() - - -def create_sampler(sample_fn, inversion_latent, ref_config): - @torch.no_grad() - def sample(model, latents, sigmas, extra_args=None, callback=None, disable=None, **extra_options): - # src_latents = prepare_src_latents(model, src_latent_image) - model_wrapper = create_model_wrapper(model, sigmas, inversion_latent, extra_args, ref_config) - output = sample_fn(model_wrapper, latents, sigmas, extra_args=extra_args, callback=callback, disable=disable) - del model_wrapper - return output - - return sample \ No newline at end of file diff --git a/utils/ref_style_config.py b/utils/ref_style_config.py new file mode 100644 index 0000000..458d057 --- /dev/null +++ b/utils/ref_style_config.py @@ -0,0 +1,18 @@ + +class RefStyleConfig: + def __init__(self, + sampling, + style_latent, + style_transfer, + style_prompt, + style_count, + style_start_percent, + style_end_percent + ): + self.sampling = sampling + self.style_latent = style_latent + self.style_transfer = style_transfer + self.style_count = style_count + self.style_prompt = style_prompt + self.style_start_percent = style_start_percent + self.style_end_percent = style_end_percent \ No newline at end of file