diff --git a/README.md b/README.md index 1eb1dab..89d29ca 100644 --- a/README.md +++ b/README.md @@ -18,3 +18,7 @@ Current `Visionatrix/UI` nodes list: - **VixUi-ListLogic** - to define the list of available mode options similar to `VixUi-CheckboxLogic` (up to 6 input options); - **VixUi-WorkflowMetadata** - mandatory node to fill the workflow metadata required for each Visionatrix flow for displaying in the UI list of workflows; - **VixUiAspectRatioSelector** - use it to display the desired image aspect ratio for your Flow. + +### Incorporated nodes + +- **StyleAlignedBatchAlign** - from the `style_aligned_comfy` repository diff --git a/__init__.py b/__init__.py index f775bbd..c7c34fe 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,7 @@ import json +from .style_aligned import StyleAlignedBatchAlign + class AnyType(str): def __ne__(self, __value: object) -> bool: @@ -427,7 +429,7 @@ NODE_CLASS_MAPPINGS = { "VixUiWorkflowMetadata": VixUiWorkflowMetadata, "VixDynamicLoraDefinition": VixDynamicLoraDefinition, "VixCheckboxLogic": VixCheckboxLogic, - + "StyleAlignedBatchAlign": StyleAlignedBatchAlign, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -443,4 +445,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VixUiWorkflowMetadata": "VixUI-WorkflowMetadata", "VixDynamicLoraDefinition": "Vix-DynamicLoraDefinition", "VixCheckboxLogic": "Vix-CheckboxLogic", + "StyleAlignedBatchAlign": "StyleAligned Batch Align", } diff --git a/pyproject.toml b/pyproject.toml index 32b5b36..d5e67d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-visionatrix" description = "The ComfyUI-Visionatrix nodes are designed for convenient ComfyUI to [a/Visionatrix](https://github.com/Visionatrix/Visionatrix) workflow support migration, in particular to extract prompt input params (input, textarea, checkbox, select, range, file) to be used in simplified Visionatrix UI." -version = "1.0.6" +version = "1.0.7" license = {file = "LICENSE"} urls.Repository = "https://github.com/Visionatrix/ComfyUI-Visionatrix" diff --git a/style_aligned.py b/style_aligned.py new file mode 100644 index 0000000..e1fa608 --- /dev/null +++ b/style_aligned.py @@ -0,0 +1,162 @@ +# MIT License +# +# Copyright (c) 2023 Brian Fitzgerald +# original repository: https://github.com/brianfitzgerald/style_aligned_comfy + +import torch.nn as nn +import torch + +from comfy.model_patcher import ModelPatcher + + +T = torch.Tensor + + +class StyleAlignedArgs: + def __init__(self, share_attn: str) -> None: + self.adain_keys = "k" in share_attn + self.adain_values = "v" in share_attn + self.adain_queries = "q" in share_attn + + share_attention: bool = True + adain_queries: bool = True + adain_keys: bool = True + adain_values: bool = True + + +def expand_first( + feat: T, + scale=1.0, +) -> T: + b = feat.shape[0] + feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1) + if scale == 1: + feat_style = feat_style.expand(2, b // 2, *feat.shape[1:]) + else: + feat_style = feat_style.repeat(1, b // 2, 1, 1, 1) + feat_style = torch.cat([feat_style[:, :1], scale * feat_style[:, 1:]], dim=1) + return feat_style.reshape(*feat.shape) + + +def concat_first(feat: T, dim=2, scale=1.0) -> T: + feat_style = expand_first(feat, scale=scale) + return torch.cat((feat, feat_style), dim=dim) + + +def calc_mean_std(feat, eps: float = 1e-5) -> "tuple[T, T]": + feat_std = (feat.var(dim=-2, keepdims=True) + eps).sqrt() + feat_mean = feat.mean(dim=-2, keepdims=True) + return feat_mean, feat_std + + +def adain(feat: T) -> T: + feat_mean, feat_std = calc_mean_std(feat) + feat_style_mean = expand_first(feat_mean) + feat_style_std = expand_first(feat_std) + feat = (feat - feat_mean) / feat_std + feat = feat * feat_style_std + feat_style_mean + return feat + + +class SharedAttentionProcessor: + def __init__(self, args: StyleAlignedArgs, scale: float): + self.args = args + self.scale = scale + + def __call__(self, q, k, v, extra_options): + if self.args.adain_queries: + q = adain(q) + if self.args.adain_keys: + k = adain(k) + if self.args.adain_values: + v = adain(v) + if self.args.share_attention: + k = concat_first(k, -2, scale=self.scale) + v = concat_first(v, -2) + + return q, k, v + + +def get_norm_layers( + layer: nn.Module, + norm_layers_: "dict[str, list[nn.GroupNorm | nn.LayerNorm]]", + share_layer_norm: bool, + share_group_norm: bool, +): + if isinstance(layer, nn.LayerNorm) and share_layer_norm: + norm_layers_["layer"].append(layer) + if isinstance(layer, nn.GroupNorm) and share_group_norm: + norm_layers_["group"].append(layer) + else: + for child_layer in layer.children(): + get_norm_layers( + child_layer, norm_layers_, share_layer_norm, share_group_norm + ) + + +def register_norm_forward( + norm_layer: nn.GroupNorm | nn.LayerNorm, +) -> nn.GroupNorm | nn.LayerNorm: + if not hasattr(norm_layer, "orig_forward"): + setattr(norm_layer, "orig_forward", norm_layer.forward) + orig_forward = norm_layer.orig_forward + + def forward_(hidden_states: T) -> T: + n = hidden_states.shape[-2] + hidden_states = concat_first(hidden_states, dim=-2) + hidden_states = orig_forward(hidden_states) # type: ignore + return hidden_states[..., :n, :] + + norm_layer.forward = forward_ # type: ignore + return norm_layer + + +def register_shared_norm( + model: ModelPatcher, + share_group_norm: bool = True, + share_layer_norm: bool = True, +): + norm_layers = {"group": [], "layer": []} + get_norm_layers(model.model, norm_layers, share_layer_norm, share_group_norm) + print( + f"Patching {len(norm_layers['group'])} group norms, {len(norm_layers['layer'])} layer norms." + ) + return [register_norm_forward(layer) for layer in norm_layers["group"]] + [ + register_norm_forward(layer) for layer in norm_layers["layer"] + ] + + +SHARE_NORM_OPTIONS = ["both", "group", "layer", "disabled"] +SHARE_ATTN_OPTIONS = ["q+k", "q+k+v", "disabled"] + + +class StyleAlignedBatchAlign: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "share_norm": (SHARE_NORM_OPTIONS,), + "share_attn": (SHARE_ATTN_OPTIONS,), + "scale": ("FLOAT", {"default": 1, "min": 0, "max": 1.0, "step": 0.1}), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "patch" + CATEGORY = "style_aligned" + + def patch( + self, + model: ModelPatcher, + share_norm: str, + share_attn: str, + scale: float, + ): + m = model.clone() + share_group_norm = share_norm in ["group", "both"] + share_layer_norm = share_norm in ["layer", "both"] + register_shared_norm(model, share_group_norm, share_layer_norm) + args = StyleAlignedArgs(share_attn) + m.set_model_attn1_patch(SharedAttentionProcessor(args, scale)) + return (m,)