From 37b08f9b3ef2f767802b6c2a0fb0d98d55f7b716 Mon Sep 17 00:00:00 2001
From: hako-mikan <122196982+hako-mikan@users.noreply.github.com>
Date: Fri, 31 Jan 2025 19:14:53 +0900
Subject: [PATCH] #335

---
 scripts/latent.py |  5 +++++
 scripts/rp.py     | 19 +++++++++++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/scripts/latent.py b/scripts/latent.py
index 3e9d41e..90dc873 100644
--- a/scripts/latent.py
+++ b/scripts/latent.py
@@ -95,6 +95,11 @@ def denoiser_callback_s(self, params: CFGDenoiserParams):
     self.step = params.sampling_step
     self.total_step = params.total_sampling_steps
 
+    self.pn = self.pn_s
+
+    if self.only_r:
+        return
+
     if "Pro" in self.mode:  # in Prompt mode, make masks from sum of attension maps
         if self.x == None : cloneparams(params,self) # return to step 0 if mask is ready
         self.pfirst = True
diff --git a/scripts/rp.py b/scripts/rp.py
index c8b4524..faeebb2 100644
--- a/scripts/rp.py
+++ b/scripts/rp.py
@@ -229,6 +229,8 @@ class Script(modules.scripts.Script):
         self.count = 0
         self.eq = True
         self.pn = True
+        self.pn_s = True
+        self.only_r = False
         self.hr = False
         self.hr_scale = 0
         self.hr_w = 0
@@ -549,6 +551,7 @@ class Script(modules.scripts.Script):
                 else:                    
                     shared.batch_cond_uncond = orig_batch_cond_uncond
                 unloadlorafowards(p)
+                denoiserdealer(self, True)
             else:
                 hook_forwards(self, p, remove = True)
                 setuploras(self)
@@ -559,7 +562,7 @@ class Script(modules.scripts.Script):
             self.ex = "Ex" in self.mode
             if not usebase : bratios = "0"
             hook_forwards(self, p)
-            denoiserdealer(self, p)
+            denoiserdealer(self, False)
 
         if OPTCOUT in options: commentouter(p)
         neighbor(self,p)                                                    #detect other extention
@@ -631,7 +634,7 @@ class Script(modules.scripts.Script):
                 if self.lora_applied: # SBM Don't override orig twice on batch calls.
                     pass
                 else:
-                    denoiserdealer(self, p)
+                    denoiserdealer(self, False)
                     self.lora_applied = True
                 #escape reload loras in hires-fix
 
@@ -668,7 +671,14 @@ def unloader(self,p):
 
     unloadlorafowards(p)
 
-def denoiserdealer(self, p):
+def denoiserdealer(self, only_r):
+    if not hasattr(self,"dr_callbacks"):
+        self.dr_callbacks = on_cfg_denoiser(self.denoiser_callback)
+
+    if only_r:
+        self.only_r = True
+        return
+
     if self.calc =="Latent": # prompt mode use only denoiser callbacks
         if not hasattr(self,"dd_callbacks"):
             self.dd_callbacks = on_cfg_denoised(self.denoised_callback)
@@ -677,9 +687,6 @@ def denoiserdealer(self, p):
         else:                    
             shared.batch_cond_uncond = False
 
-    if not hasattr(self,"dr_callbacks"):
-        self.dr_callbacks = on_cfg_denoiser(self.denoiser_callback)
-
     if self.diff:
         if not hasattr(self,"dd_callbacks"):
             self.dd_callbacks = on_cfg_denoised(self.denoised_callback)