From be8a2deeec320a9254888328bc5ebc15defe2c4a Mon Sep 17 00:00:00 2001
From: brianf <brian@stability.ai>
Date: Wed, 13 Dec 2023 22:29:38 +0000
Subject: [PATCH] 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))