From 181641ca04ac524992f40e8717f5da01eb82ea79 Mon Sep 17 00:00:00 2001
From: asagi4 <130366179+asagi4@users.noreply.github.com>
Date: Sat, 3 May 2025 21:11:36 +0300
Subject: [PATCH] v0.4.0

---
 README.md                                       |  2 +-
 __init__.py                                     | 14 +++++++++++---
 {workflows => example_workflows}/AGExample.json |  0
 pyproject.toml                                  |  2 +-
 4 files changed, 13 insertions(+), 5 deletions(-)
 rename {workflows => example_workflows}/AGExample.json (100%)

diff --git a/README.md b/README.md
index 2684d77..aa6060a 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@ An implementation of adaptive guidance for ComfyUI
 
 See https://bcv-uniandes.github.io/adaptiveguidance-wp/
 
-Import [this workflow](workflows/AGExample.json?raw=1) into ComfyUI to compare Adaptive Guidance vs. normal CFG.
+Import [this workflow](example_workflows/AGExample.json?raw=1) into ComfyUI to compare Adaptive Guidance vs. normal CFG.
 
 ## What
 
diff --git a/__init__.py b/__init__.py
index 16e3807..3926b9e 100644
--- a/__init__.py
+++ b/__init__.py
@@ -215,8 +215,14 @@ class AdaptiveProjectedGuidanceFunction:
             cond = args["cond"]
             uncond = args["uncond"]
         cfg_scale = args["cond_scale"]
+        sigma = args["sigma"][0].item()
         step = args["model"].model_sampling.timestep(args["sigma"])[0].item()
-        x = args["input"]
+        x_orig = args["input"]
+        if self.mode == "vpred":
+            sigma = step
+            x = x_orig / (sigma * sigma + 1.0)
+            cond = ((x - (x_orig - cond)) * (sigma**2 + 1.0) ** 0.5) / (sigma)
+            uncond = ((x - (x_orig - uncond)) * (sigma**2 + 1.0) ** 0.5) / (sigma)
 
         if self.current_step < step:
             self.current_step = 999.0
@@ -249,7 +255,9 @@ class AdaptiveProjectedGuidanceFunction:
 
         pred = cond + (cfg_scale - 1) * (diff_orthogonal + self.eta * diff_parallel)
         if "denoised" == self.mode:
-            pred =  x - pred
+            pred = x_orig - pred
+        elif "vpred" == self.mode:
+            pred = x_orig - (x - pred * sigma / (sigma * sigma + 1.0) ** 0.5)
         return pred
 
 
@@ -262,7 +270,7 @@ class AdaptiveProjectedGuidance:
                 "momentum": ("FLOAT", {"default": 0.5, "min": -1.0, "max": 1.0, "step": 0.01}),
                 "eta": ("FLOAT", {"default": 1, "min": 0.0, "max": 1.0, "step": 0.01}),
                 "norm_threshold": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 50.0, "step": 0.1}),
-                "mode": (["normal", "denoised"],),
+                "mode": (["normal", "denoised", "vpred"],),
                 "adaptive_momentum": ("FLOAT", {"default": 0.18, "min": 0, "max": 1.0, "step": 0.01}),
             },
         }
diff --git a/workflows/AGExample.json b/example_workflows/AGExample.json
similarity index 100%
rename from workflows/AGExample.json
rename to example_workflows/AGExample.json
diff --git a/pyproject.toml b/pyproject.toml
index f4fe58f..ae65e54 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-adaptive-guidance"
 description = "An implementation of adaptive guidance for ComfyUI\nSee https://bcv-uniandes.github.io/adaptiveguidance-wp/"
-version = "0.3.1"
+version = "0.4.0"
 license = { text = "GNU General Public License v3.0" }
 
 [project.urls]