From c3361e04e7203b504b06a24212c2565eb1b4c3f2 Mon Sep 17 00:00:00 2001
From: nolbert82 <nolanofrance@gmail.com>
Date: Wed, 22 Oct 2025 21:33:06 +0200
Subject: [PATCH] Fixed guidance end

---
 modules/processing_callbacks.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py
index 3aec83fca..e36bd0e75 100644
--- a/modules/processing_callbacks.py
+++ b/modules/processing_callbacks.py
@@ -93,15 +93,28 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {}
     if step != getattr(pipe, 'num_timesteps', 0):
         kwargs = processing_correction.correction_callback(p, timestep, kwargs, initial=step == 0)
     kwargs = prompt_callback(step, kwargs)  # monkey patch for diffusers callback issues
-    if step == int(getattr(pipe, 'num_timesteps', 100) * p.cfg_end) and 'prompt_embeds' in kwargs and 'negative_prompt_embeds' in kwargs:
+
+    if step == 0:
+        setattr(pipe, "_cfg_end_applied", False)
+
+    cfg_end = getattr(p, "cfg_end", 1.0) or 1.0
+    total_steps = getattr(pipe, "num_timesteps", 0)
+    target_step = int(total_steps * cfg_end) if total_steps else 0
+    if (
+        cfg_end < 1.0
+        and not getattr(pipe, "_cfg_end_applied", False)
+        and step >= target_step
+    ):
+        setattr(pipe, "_cfg_end_applied", True)
         if "PAG" in shared.sd_model.__class__.__name__:
             pipe._guidance_scale = 1.001 if pipe._guidance_scale > 1 else pipe._guidance_scale  # pylint: disable=protected-access
             pipe._pag_scale = 0.001  # pylint: disable=protected-access
         else:
             pipe._guidance_scale = 0.0  # pylint: disable=protected-access
-            for key in {"prompt_embeds", "negative_prompt_embeds", "add_text_embeds", "add_time_ids"} & set(kwargs):
-                if kwargs[key] is not None:
-                    kwargs[key] = kwargs[key].chunk(2)[-1]
+            for key in {"prompt_embeds", "negative_prompt_embeds", "add_text_embeds", "add_time_ids"}:
+                tensor = kwargs.get(key, None)
+                if tensor is not None and hasattr(tensor, "chunk") and tensor.shape[0] % 2 == 0:
+                    kwargs[key] = tensor.chunk(2)[-1]
     try:
         current_noise_pred = kwargs.get("noise_pred", None)
         if current_noise_pred is None: