From aed4e1c63cc6394b5d4d093976fe95e85f5ea252 Mon Sep 17 00:00:00 2001
From: ExponentialML <59846140+ExponentialML@users.noreply.github.com>
Date: Thu, 14 Mar 2024 20:03:13 -0700
Subject: [PATCH] Add model output for consistency

---
 nodes.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/nodes.py b/nodes.py
index 9a8cecc..8d69840 100644
--- a/nodes.py
+++ b/nodes.py
@@ -19,8 +19,8 @@ class ApplyVisualStyle:
             } 
         }
 
-    RETURN_TYPES = ("CONDITIONING","CONDITIONING", "LATENT")
-    RETURN_NAMES = ("positive", "negative", "latents")
+    RETURN_TYPES = ("MODEL", "CONDITIONING","CONDITIONING", "LATENT")
+    RETURN_NAMES = ("model", "positive", "negative", "latents")
     
     FUNCTION = "apply_visual_style_prompt"
 
@@ -38,7 +38,7 @@ class ApplyVisualStyle:
         reference_latent = vae.encode(reference_image[:,:,:,:3])
         
         for n, m in model.model.diffusion_model.named_modules():
-            if m.__class__.__name__ == "CrossAttention":
+            if m.__class__.__name__  == "CrossAttention":
                 processor = VisualStyleProcessor(m, enabled=enabled)
                 setattr(m, 'forward', processor.visual_style_forward)
 
@@ -50,7 +50,7 @@ class ApplyVisualStyle:
         denoise_mask = torch.ones_like(latents)[:, :1, ...]
         denoise_mask[0] = 0.
 
-        return (conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask})    
+        return (model, conditioning_prompt, negative_prompt, {"samples": latents, "noise_mask": denoise_mask})    
 
 NODE_CLASS_MAPPINGS = {"ApplyVisualStyle": ApplyVisualStyle}