improved logics

.gitattributes

@@ -1,2 +0,0 @@
-*.jpg filter=lfs diff=lfs merge=lfs -text
-*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -2,46 +2,31 @@
 This is an Extension for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), which is the joint research between me and <ins>TimothyAlexisVass</ins>.
 For more information, check out the original [Extension](https://github.com/Haoming02/sd-webui-diffusion-cg) for **Automatic1111**.
 
-## Nodes
-Some example workflows are included~
+## How to Use
+> Example workflows are included~
+- Attach the **Recenter** or **RecenterXL** node between `Empty Latent` and `KSampler` nodes
+    - Adjust the **strength** and **color** sliders as needed
+- Attach the **Normalization** or **NormalizationXL** node between `KSampler` and `VAE Decode` nodes
 
-- **Hook Recenter:** For **SD 1.5**. Hooks the callback to achieve the centering effect.
-  - Comes with `Effect Strength` slider and `CMYK` color settings
-- **Hook Recenter XL:** For **SDXL**. Hooks the callback to achieve the centering effect.
-  - Comes with `Effect Strength` slider and `Lab` color settings
-- **Unhook Recenter** (Optional)**:** Unhook the callback to disable the effects completely.
-  - If used, put near the end of the workflow
-- **Normalization:** For **SD 1.5**. Use before the `VAE Decode` node to achieve the normalization effect.
-- **NormalizationXL:** For **SDXL**. Use before the `VAE Decode` node to achieve the normalization effect.
-
-#### Important:
-- Toggle `custom_sampler` if you're using the **SamplerCustom** node.
-- In a single workflow, you only need to hook the callback once. The simplest way is to add it between the `Positive Prompt` and the `Sampler`.
-- Due to how `ComfyUI` works, if you also add **Unhook Recenter**, the effect may not work sometimes unless you also change the prompt.
-
-> ComfyUI doesn't go through a node unless it needs to be updated, so if you unhook the callback and the parameters didn't change *(**eg.** you're only iterating throguh seeds)*, then the callback will not be hooked again. Easiest way to solve this is just adding a space to the positive prompt, or just don't unhook the callback.
+### Important:
+- The **Recenter** is "global." If you want to disable it during later part of the workflow *(**eg.** during `Hires. Fix`)*, 
+you have to add another **Recenter** node and set its `strength` to `0.0`.
 
 ## Samples
 
 <p align="center">
 <b>SD 1.5</b><br>
-<img src="samples\1.5_off.jpg" width=384>
-<img src="workflows\1.5_on.png" width=384>
-<br><code>Off | On</code><br>
+<img src="samples\15_off.jpg" width=384>
+<img src="workflows\15_on.png" width=384>
+<br><code>Off | On</code>
 </p>
 
 <p align="center">
 <b>SDXL</b><br>
 <img src="samples\xl_off.jpg" width=384>
 <img src="workflows\xl_on.png" width=384>
-<br><code>Off | On</code><br>
+<br><code>Off | On</code>
 </p>
 
 ## Known Issue
-- Doesn't really work with `LCM` Sampler
-
-<hr>
-
-##### Checkpoints Used:
-- [UHD-23](https://civitai.com/models/22371/uhd-23)
-- [Juggernaut XL](https://civitai.com/models/133005/juggernaut-xl)
+- Doesn't work with certain Samplers

__init__.py

@@ -1,21 +1,16 @@
 from .normalization import Normalization, NormalizationXL
-from .recenter import HookCallback, HookCallbackXL, UnhookCallback
-# from .tensor_debug import Debug
+from .recenter import Recenter, RecenterXL
 
 NODE_CLASS_MAPPINGS = {
     "Normalization": Normalization,
     "NormalizationXL": NormalizationXL,
-    "Hook Recenter": HookCallback,
-    "Hook Recenter XL": HookCallbackXL,
-    "Unhook Recenter": UnhookCallback,
-  # "Tensor Debug": Debug,
+    "Recenter": Recenter,
+    "Recenter XL": RecenterXL
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "Normalization": "Normalization",
     "NormalizationXL": "NormalizationXL",
-    "Hook Recenter": "Hook Recenter",
-    "Hook Recenter XL": "Hook Recenter XL",
-    "Unhook Recenter": "Unhook Recenter",
-  # "Tensor Debug": "Tensor Debug",
+    "Recenter": "Recenter",
+    "Recenter XL": "RecenterXL"
 }

normalization.py

@@ -1,6 +1,7 @@
 DYNAMIC_RANGE = [18, 14, 14, 14]
 DYNAMIC_RANGE_XL = [20, 16, 16]
 
+
 def normalize_tensor(x, r):
     ratio = r / max(abs(float(x.min())), abs(float(x.max())))
     x *= max(ratio, 0.99)
@@ -12,6 +13,7 @@ def clone_latent(latent):
 
     return cloned_latent
 
+
 class Normalization:
     @classmethod
     def INPUT_TYPES(s):

recenter.py

@@ -1,20 +1,27 @@
 import comfy
 
+rc_strength = 0.0
+LUTs = []
+
 ORIGINAL_SAMPLE = comfy.sample.sample
 ORIGINAL_SAMPLE_CUSTOM = comfy.sample.sample_custom
 
-
-def hijack(SAMPLE, LUTs:list, strength:float):
+def hijack(SAMPLE):
 
     def sample_center(*args, **kwargs):
         original_callback = kwargs['callback']
 
         def hijack_callback(step, x0, x, total_steps):
+            global rc_strength
+            global LUTs
+
+            if rc_strength == 0 or len(LUTs) == 0:
+                return original_callback(step, x0, x, total_steps)
 
             batchSize = x.size(0)
             for b in range(batchSize):
                 for c in range(len(LUTs)):
-                    x[b][c] += (LUTs[c] - x[b][c].mean()) * strength
+                    x[b][c] += (LUTs[c] - x[b][c].mean()) * rc_strength
 
             return original_callback(step, x0, x, total_steps)
 
@@ -23,30 +30,16 @@ def hijack(SAMPLE, LUTs:list, strength:float):
 
     return sample_center
 
-
-class UnhookCallback:
-    @classmethod
-    def INPUT_TYPES(s):
-        return { "required": { "latent": ("LATENT", ) } }
-
-    RETURN_TYPES = ("LATENT", )
-    FUNCTION = "unhook"
-    CATEGORY = "Diffusion CG"
-
-    def unhook(self, latent):
-        comfy.sample.sample_custom = ORIGINAL_SAMPLE_CUSTOM
-        comfy.sample.sample = ORIGINAL_SAMPLE
-
-        return (latent,)
+comfy.sample.sample = hijack(ORIGINAL_SAMPLE)
+comfy.sample.sample_custom = hijack(ORIGINAL_SAMPLE_CUSTOM)
 
 
-class HookCallback:
+class Recenter:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
-                "prompt": ("CONDITIONING",),
-                "custom_sampler": ("BOOLEAN", {"default": False}),
+                "latent": ("LATENT",),
                 "strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0,
                                 "step": 0.1, "round": 0.1, "display": "slider"}),
                 "C": ("FLOAT", {"default": 0.01, "min": -1.00, "max": 1.00, "step": 0.01}),
@@ -56,26 +49,25 @@ class HookCallback:
             }
         }
 
-    RETURN_TYPES = ("CONDITIONING",)
+    RETURN_TYPES = ("LATENT",)
     FUNCTION = "hook"
-    CATEGORY = "Diffusion CG"
+    CATEGORY = "latent"
 
-    def hook(self, prompt, custom_sampler, strength, C, M, Y, K):
-        if custom_sampler:
-            comfy.sample.sample_custom = hijack(ORIGINAL_SAMPLE_CUSTOM, [-K, -M, C, Y], strength)
-        else:
-            comfy.sample.sample = hijack(ORIGINAL_SAMPLE, [-K, -M, C, Y], strength)
+    def hook(self, latent, strength, C, M, Y, K):
+        global rc_strength
+        rc_strength = strength
+        global LUTs
+        LUTs = [-K, -M, C, Y]
 
-        return (prompt,)
+        return (latent,)
 
 
-class HookCallbackXL:
+class RecenterXL:
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
-                "prompt": ("CONDITIONING",),
-                "custom_sampler": ("BOOLEAN", {"default": False}),
+                "latent": ("LATENT",),
                 "strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0,
                                 "step": 0.1, "round": 0.1, "display": "slider"}),
                 "L": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.05}),
@@ -84,14 +76,14 @@ class HookCallbackXL:
             }
         }
 
-    RETURN_TYPES = ("CONDITIONING",)
+    RETURN_TYPES = ("LATENT",)
     FUNCTION = "hook"
-    CATEGORY = "Diffusion CG"
+    CATEGORY = "latent"
 
-    def hook(self, prompt, custom_sampler, strength, L, a, b):
-        if custom_sampler:
-            comfy.sample.sample_custom = hijack(ORIGINAL_SAMPLE_CUSTOM, [L, -a, b], strength)
-        else:
-            comfy.sample.sample = hijack(ORIGINAL_SAMPLE, [L, -a, b], strength)
+    def hook(self, latent, strength, L, a, b):
+        global rc_strength
+        rc_strength = strength
+        global LUTs
+        LUTs = [L, -a, b]
 
-        return (prompt,)
+        return (latent,)

samples/1.5_off.jpg

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:dbbda98d15ded20936bbfee25bb6ec9ea10de7e8762112b06d0478961cb3fa6a
-size 292905

tensor_debug.py

@@ -1,17 +0,0 @@
-class Debug:
-    @classmethod
-    def INPUT_TYPES(s):
-        return { "required": { "latent": ("LATENT",) } }
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "debug"
-    CATEGORY = "latent"
-
-    def debug(self, latent):
-
-        print('\n')
-        for c in range(4):
-            print(f'(min: {latent["samples"][0][c].min()}, max: {latent["samples"][0][c].max()}, mean: {latent["samples"][0][c].mean()})')
-        print('\n')
-
-        return (latent,)

workflows/1.5_on.png

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:2cefec2b4672d27c2f505e0615d769378d20048f1f58b2c1af33c3487e116979
-size 1551287