Bugfix: register shared norm with the copied model

2023-12-13 22:29:38 +00:00
parent 28804b7116
commit be8a2deeec
1 changed files with 2 additions and 2 deletions
@@ -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))