Update: remain injected option

advanced_sampler.py

@@ -7,10 +7,11 @@ def inject_ksamplerx0inpaint_call(original_ksampler_call_fn):
     def wrapper(self, x, sigma, denoise_mask, model_options={}, seed=None):
         
         # latent injection 옵션 확인
-        latent_inject_options = model_options.get("is_latent_inject")
+        latent_inject_options = model_options.get("is_latent_inject", None)
         if latent_inject_options is not None:
             start_sigma = latent_inject_options["start_sigma"]
             end_sigma = latent_inject_options["end_sigma"]
+            remain_injected = latent_inject_options["remain_injected"]
         
         if denoise_mask is not None:
             # 사용자 정의 마스크 함수 적용 (있는 경우)
@@ -38,8 +39,8 @@ def inject_ksamplerx0inpaint_call(original_ksampler_call_fn):
             else:
                 out = out * denoise_mask + self.latent_image * latent_mask
 
-        # 원래 함수로 복구
-        if self.sigmas[-2].item() >= sigma.item():
+        # remain injected 옵션에 따라 원래 함수로 복구
+        if remain_injected == False and self.sigmas[-2].item() >= sigma.item():
             KSamplerX0Inpaint.__call__ = original_ksampler_call_fn
             logging.debug("\033[94m[middlek latent injection] KSamplerX0Inpaint.__call__ return to original_ksampler_call_fn\033[0m")
 

node.py

@@ -38,12 +38,15 @@ class ApplyLatentInjection:
                              "inject_mask": ("MASK",),
                              "start_sigma": ("FLOAT", {"default": 15.0}),
                              "end_sigma": ("FLOAT", {"default": 0.0})
-                             }}
+                             },
+                "optional": {"remain_injected":([True, False], {"default":True})
+                           },
+                }
     RETURN_TYPES = ("MODEL", "LATENT",)
     FUNCTION = "apply_latent_injection"
     CATEGORY = "productfix"
 
-    def apply_latent_injection(self, model, latents, inject_image_embed, inject_mask, start_sigma, end_sigma):
+    def apply_latent_injection(self, model, latents, inject_image_embed, inject_mask, start_sigma, end_sigma, remain_injected=True):
         device = model_management.get_torch_device()
         dtype = model_management.VAE_DTYPES[0]
 
@@ -70,7 +73,7 @@ class ApplyLatentInjection:
 
         # 모델 옵션에 latent injection 파라미터 추가
         if hasattr(model, "model_options"):
-            model.model_options["is_latent_inject"] = {"start_sigma":start_sigma, "end_sigma":end_sigma}
+            model.model_options["is_latent_inject"] = {"start_sigma":start_sigma, "end_sigma":end_sigma, "remain_injected":remain_injected}
 
         return (model, latents, )