From be8a2deeec320a9254888328bc5ebc15defe2c4a Mon Sep 17 00:00:00 2001 From: brianf Date: Wed, 13 Dec 2023 22:29:38 +0000 Subject: [PATCH 1/4] Bugfix: register shared norm with the copied model --- __init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index ff7c0b5..796a4d0 100644 --- a/__init__.py +++ b/__init__.py @@ -214,13 +214,13 @@ class StyleAlignedReferenceSampler: x0_output = {} callback = latent_preview.prepare_callback( - model, sigmas.shape[-1] - 1, x0_output + m, sigmas.shape[-1] - 1, x0_output ) # Register shared norms 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) + register_shared_norm(m, share_group_norm, share_layer_norm) # Patch cross attn m.set_model_attn1_patch(SharedAttentionProcessor(self.args, scale)) From 915c562d5ea2930cf262cb6c8e5464bceb73a1ea Mon Sep 17 00:00:00 2001 From: brianf Date: Thu, 14 Dec 2023 01:41:20 +0000 Subject: [PATCH 2/4] add qkv options --- __init__.py | 37 ++++++++++++++++++++++--------------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/__init__.py b/__init__.py index 796a4d0..7bd3d76 100644 --- a/__init__.py +++ b/__init__.py @@ -22,20 +22,25 @@ def default(val, d): return d -@dataclass(frozen=True) 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 = False - shared_score_shift: float = 0.0 - only_self_level: float = 0.0 + adain_values: bool = True def expand_first( feat: T, scale=1.0, ) -> T: + """ + Expand the first element so it has the same shape as the rest of the batch. + """ b = feat.shape[0] feat_style = torch.stack((feat[0], feat[b // 2])).unsqueeze(1) if scale == 1: @@ -47,6 +52,9 @@ def expand_first( def concat_first(feat: T, dim=2, scale=1.0) -> T: + """ + concat the the feature and the style feature expanded above + """ feat_style = expand_first(feat, scale=scale) return torch.cat((feat, feat_style), dim=dim) @@ -142,6 +150,7 @@ def register_shared_norm( SHARE_NORM_OPTIONS = ["both", "group", "layer", "disabled"] +SHARE_ATTN_OPTIONS = ["q+k", "q+k+v", "disabled"] class StyleAlignedReferenceSampler: @@ -151,6 +160,7 @@ class StyleAlignedReferenceSampler: "required": { "model": ("MODEL",), "share_norm": (SHARE_NORM_OPTIONS,), + "share_attn": (SHARE_ATTN_OPTIONS,), "scale": ("FLOAT", {"default": 1, "min": 0, "max": 2.0, "step": 0.1}), "batch_size": ("INT", {"default": 2, "min": 1, "max": 8, "step": 1}), "noise_seed": ( @@ -180,13 +190,11 @@ class StyleAlignedReferenceSampler: FUNCTION = "patch" CATEGORY = "style_aligned" - def __init__(self) -> None: - self.args = StyleAlignedArgs() - def patch( self, model: ModelPatcher, share_norm: str, + share_attn: str, scale: float, batch_size: int, noise_seed: int, @@ -198,6 +206,7 @@ class StyleAlignedReferenceSampler: ref_latent: "dict[str, T]", ) -> "tuple[dict, dict]": m = model.clone() + args = StyleAlignedArgs(share_attn) # Concat batch with style latent style_latent_tensor = ref_latent["samples"] @@ -213,9 +222,7 @@ class StyleAlignedReferenceSampler: noise = torch.cat((ref_noise, noise), dim=0) x0_output = {} - callback = latent_preview.prepare_callback( - m, sigmas.shape[-1] - 1, x0_output - ) + callback = latent_preview.prepare_callback(m, sigmas.shape[-1] - 1, x0_output) # Register shared norms share_group_norm = share_norm in ["group", "both"] @@ -223,7 +230,7 @@ class StyleAlignedReferenceSampler: register_shared_norm(m, share_group_norm, share_layer_norm) # Patch cross attn - m.set_model_attn1_patch(SharedAttentionProcessor(self.args, scale)) + m.set_model_attn1_patch(SharedAttentionProcessor(args, scale)) disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED samples = comfy.sample.sample_custom( @@ -261,6 +268,7 @@ class StyleAlignedBatchAlign: "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}), } } @@ -269,20 +277,19 @@ class StyleAlignedBatchAlign: FUNCTION = "patch" CATEGORY = "style_aligned" - def __init__(self) -> None: - self.args = StyleAlignedArgs() - 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) - m.set_model_attn1_patch(SharedAttentionProcessor(self.args, scale)) + args = StyleAlignedArgs(share_attn) + m.set_model_attn1_patch(SharedAttentionProcessor(args, scale)) return (m,) From d92d756da8f920876c6e1d268d8c49ff52af615e Mon Sep 17 00:00:00 2001 From: brianf Date: Thu, 14 Dec 2023 01:43:00 +0000 Subject: [PATCH 3/4] update readme --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 2f74ea5..5fba9ca 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,7 @@ _A batch of generations with the same parameters and the Batch Align node applie ##### Parameters - `model`: The base model to patch. +- `share_attn`: Which components of self-attention are normalized. Defaults to `q+k`. Set to `q+k+v` for more extreme sharing, at the cost of quality in some cases. - `share_norm`: Whether to share normalization across the batch. Defaults to `both`. Set to `group` or `layer` to only share group or layer normalization, respectively. - `scale`: The scale at which to apply the style-alignment effect. Defaults to `1`. - `batch_size`, `noise_seed`, `control_after_generate`, `cfg`: Identical to the standard `KSampler` parameters. From 166afd5ad2f465fcdfee72bc410ce20b2b8d1291 Mon Sep 17 00:00:00 2001 From: brianf Date: Thu, 14 Dec 2023 01:43:37 +0000 Subject: [PATCH 4/4] update readme --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 5fba9ca..a3d94e8 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ _Above, a reference image, and a batch of images generated using the prompt 'a r ##### Parameters - `model`: The base model to patch. +- `share_attn`: Which components of self-attention are normalized. Defaults to `q+k`. Set to `q+k+v` for more extreme sharing, at the cost of quality in some cases. - `share_norm`: Whether to share normalization across the batch. Defaults to `both`. Set to `group` or `layer` to only share group or layer normalization, respectively. - `scale`: The scale at which to apply the style-alignment effect. Defaults to `1`.