nodes/ttulttul-ComfyUI-FlowMatching-Inverter
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Complete node pack

Browse Source
This commit is contained in:
Ken Simpson
2025-09-28 12:53:32 -07:00
parent 363cbc16a9
commit a13d650773
2 changed files with 497 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+108
View File
@@ -0,0 +1,108 @@
# ComfyUI Qwen Rectified Flow Tools
This custom node pack for ComfyUI provides a suite of advanced tools for performing "invert-and-resample" workflows with Rectified Flow models like Qwen-VL. It allows you to take an existing image, convert it to a noisy latent using a guided inversion process, and then use that latent as a starting point to generate creative, high-quality variations.
The pack provides both the individual building blocks for maximum flexibility and a final "convenience node" that encapsulates the optimal workflow for ease of use.
## Features
- **Guided Latent Inversion:** Go beyond simple noise addition. Use the model's own understanding of the image (its velocity field) to intelligently re-noise a latent.
- **Creative Controls:** Deterministically **amplify** the inversion process or stochastically **perturb** it with seeded noise to create unique variations.
- **Stability Tools:** Includes specialized nodes to solve common artifacts in advanced workflows, such as VAE "waffle patterns" and sampler "brown-out" collapse.
- **Modular and Combined Nodes:** Use the individual components to experiment, or use the final `Latent Hybrid Inverter` for a powerful, one-node solution.
## Nodes Included
### 1. Qwen Rectified Flow Inverter
This is the core node of the pack. It takes a clean latent and "re-noises" it by running the Rectified Flow ODE integration in reverse.
- **Key Inputs:**
- `inversion_strength`: How far to re-noise the image (0.0 = no change, 1.0 = full noise).
- `velocity_amplification`: Deterministically speeds up or slows down the inversion.
- `velocity_perturb_strength`: Adds seeded random chaos to the inversion path for creative variations.
- `normalize_output`: **Crucial for stability.** Fixes the statistics of the output latent to prevent sampler errors.
### 2. Latent Gaussian Blur
A utility node to fix high-frequency artifacts in latents. Its primary purpose is to eliminate the "waffle pattern" that can appear in img2img workflows.
- **Key Inputs:**
- `sigma`: The strength of the blur. A small value (0.5-1.5) is usually enough.
- `blur_mode`: Choose between `Spatial Only` (safe, recommended) or the experimental `Spatial and Channel` for artistic effects.
### 3. Add Latent Noise (Seeded)
A simple but powerful utility to add a configurable amount of seeded Gaussian noise to any latent. Useful for re-introducing "perfect" noise to a latent that has an unnatural distribution.
- **Key Inputs:**
- `strength`: How much noise to add, relative to the latent's existing signal strength.
- `seed`: Makes the random noise pattern repeatable.
### 4. Forward Diffusion (Add Scheduled Noise)
This node performs the "natural" noising process that diffusion models are trained on. It takes a clean latent and adds the mathematically "perfect" amount of noise for a given step in the schedule. This creates a stable "anchor" latent.
- **Key Inputs:**
- `noise_strength`: The target noise level (0.0 to 1.0). Should match your KSampler's `denoise` value.
### 5. Latent Hybrid Inverter (Qwen)
The final convenience node that combines the best practices discovered into a single, powerful tool. It internally performs the creative inversion, creates a stable anchor latent via forward diffusion, and blends them together to produce an ideal starting latent for resampling.
- **Key Inputs:**
- `strength`: The target noise level for both passes.
- `blend_factor`: Controls the mix between the creative inverter latent (1.0) and the stable anchor latent (0.0). A 50/50 blend (`0.5`) is recommended.
- Contains all the creative controls from the inverter and separate seeds for each internal process.
## Workflows
### Basic Invert and Resample (Img2Img)
This workflow demonstrates the basic principle but may be prone to artifacts at high denoise strengths.
1. **VAE Encode** your source image.
2. Connect the latent to the **Qwen Rectified Flow Inverter**.
- Set `inversion_strength` to `0.8`.
3. Connect the output to a **KSampler**.
- **Crucially, set `denoise` to match the `inversion_strength` (`0.8`).**
4. **VAE Decode** the result.
### Advanced: The "Anchor" Workflow (Recommended for High Quality)
This is the optimal workflow that solves both the "waffle pattern" and "brown-out" artifacts, allowing for high-strength, high-quality variations.
1. **VAE Encode** the source image.
2. **Create the "Creative" Latent:**
- Feed the clean latent into the **Qwen Rectified Flow Inverter**. Set your desired `inversion_strength`, `amplification`, `perturbation`, etc.
- (Optional but Recommended) Feed the result into a **Latent Gaussian Blur** with a low `sigma` (e.g., 1.0) to remove VAE artifacts.
3. **Create the "Anchor" Latent:**
- Feed the original clean latent from the VAE Encode into the **Forward Diffusion (Add Scheduled Noise)** node.
- Set its `noise_strength` to *exactly match* the `inversion_strength` from the inverter. Use a different seed.
4. **Blend the Latents:**
- Use a `Latent Blend` or `Latent Slerp` node to mix the "Creative" and "Anchor" latents. A blend factor of `0.5` is a great starting point.
5. **Sample from the Blend:**
- Feed the blended latent into a **KSampler** (or **KSampler Advanced** with `add_noise` disabled).
- Set the KSampler's `denoise` value to *exactly match* the `strength` used in the previous steps.
6. **VAE Decode** the final image.
![image](https://user-images.githubusercontent.com/12345/67890.png)*(Note: A diagram of the advanced workflow would be ideal here)*
### Ultimate Convenience: The Hybrid Inverter Workflow
This workflow uses the single convenience node to achieve the same result as the advanced workflow with fewer nodes.
1. **VAE Encode** the source image.
2. Connect the latent to the **Latent Hybrid Inverter**.
- Set the desired `strength` (e.g., 0.75).
- Set the `blend_factor` (e.g., 0.5).
- Configure seeds and creative parameters.
3. Connect the output to a **KSampler**.
- Set the `denoise` to **match the `strength`** from the hybrid inverter.
4. **VAE Decode** the result.
## Installation
1. Clone or download this repository.
2. Place the entire folder (e.g., `ComfyUI-Qwen-Inverter-Pack`) into the `ComfyUI/custom_nodes/` directory.
3. Restart ComfyUI. The nodes will be available under the `Qwen/Sampling` and `Latent` categories.
__init__.py
+389 -11
View File
@@ -1,7 +1,135 @@
import torch
import comfy.model_management
import comfy.utils
import torch
import torch.nn.functional as F
import torchvision.transforms.functional as TF
def slerp(val, low, high):
"""
Spherical Linear Interpolation for PyTorch tensors, adapted for latents.
This is a more 'natural' way to blend for high-dimensional spaces.
Handles 4D or 5D tensors by flattening them for the operation.
"""
# Get original shape and flatten the tensors
original_shape = low.shape
low = low.flatten()
high = high.flatten()
# Normalize the vectors to be unit vectors
low_norm = F.normalize(low, p=2, dim=0)
high_norm = F.normalize(high, p=2, dim=0)
# Calculate the angle between the vectors
omega = torch.acos((low_norm * high_norm).sum())
# Handle the case where vectors are very close
if omega.abs().item() < 1e-3:
return low.reshape(original_shape) * (1.0 - val) + high.reshape(original_shape) * val
so = torch.sin(omega)
# Perform the SLERP interpolation
res = (torch.sin((1.0 - val) * omega) / so) * low + (torch.sin(val * omega) / so) * high
# Reshape back to the original tensor shape
return res.reshape(original_shape)
class LatentHybridInverter:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"latent_image": ("LATENT",),
"positive": ("CONDITIONING", ),
"steps": ("INT", {"default": 20, "min": 1, "max": 1000}),
"strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Strength for both inversion and forward diffusion. The target noise level."}),
"blend_factor": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "0.0 = 100% Forward Diffusion (Stable), 1.0 = 100% Inverter (Creative)"}),
"inverter_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed for the creative, perturbed inversion process."}),
"forward_diffusion_seed": ("INT", {"default": 1, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed for the stable, 'perfect' noise anchor."}),
"velocity_amplification": ("FLOAT", {"default": 1.0, "min": -0.9, "max": 2.0, "step": 0.05}),
"velocity_perturb_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 5.0, "step": 0.01}),
"internal_precision": (["bfloat16", "float16", "force_float32"], {"default": "bfloat16"}),
"normalize_output": (["enable", "disable"], {"default": "enable", "tooltip": "Normalizing the inverter output is highly recommended."}),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "hybrid_invert"
CATEGORY = "Qwen/Sampling"
def hybrid_invert(self, model, latent_image, positive, steps, strength, blend_factor, inverter_seed, forward_diffusion_seed, velocity_amplification, velocity_perturb_strength, internal_precision, normalize_output):
if strength == 0.0:
return (latent_image,)
device = comfy.model_management.get_torch_device()
native_dtype = comfy.model_management.unet_dtype()
autocast_dtype = torch.bfloat16 if internal_precision == "bfloat16" and torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32 if internal_precision == "force_float32" else torch.float16
clean_latent = latent_image["samples"].clone().to(device, native_dtype)
context = positive[0][0].to(device, native_dtype)
qwen_transformer = model.model.diffusion_model
comfy.model_management.load_models_gpu([model])
# --- PASS 1: CREATIVE INVERSION ---
inversion_steps = max(1, int(steps * strength))
end_t = strength
timesteps = torch.linspace(0, end_t, inversion_steps + 1, device=device, dtype=native_dtype)
x_t_inverted = clean_latent.unsqueeze(2) if clean_latent.dim() == 4 else clean_latent
with torch.no_grad():
for i in range(inversion_steps):
t_current, t_next = timesteps[i], timesteps[i+1]
t_current_b = torch.full((x_t_inverted.shape[0],), t_current, device=device, dtype=native_dtype)
with torch.autocast(device_type="cuda", dtype=autocast_dtype):
predicted_velocity = qwen_transformer(x=x_t_inverted, timestep=t_current_b, context=context, attention_mask=None)
if predicted_velocity.dim() == 4: predicted_velocity = predicted_velocity.unsqueeze(2)
if velocity_amplification != 0.0: predicted_velocity *= (1.0 + velocity_amplification)
if velocity_perturb_strength > 0:
step_seed = inverter_seed + i
generator = torch.Generator(device=device).manual_seed(step_seed)
velocity_std = torch.std(predicted_velocity)
noise = torch.randn(predicted_velocity.shape, generator=generator, device=device, dtype=native_dtype)
predicted_velocity += noise * velocity_std * velocity_perturb_strength
dt = t_next - t_current
x_t_inverted += predicted_velocity * dt
if torch.isnan(x_t_inverted).any():
print(f"!!! NaN in Inverter Pass. Stopping. !!!"); x_t_inverted.fill_(0); break
if normalize_output == "enable":
for i in range(x_t_inverted.shape[0]):
item = x_t_inverted[i]
mean, std = torch.mean(item), torch.std(item)
if std > 1e-5: x_t_inverted[i] = (item - mean) / std
else: x_t_inverted[i] = item - mean
# --- PASS 2: STABLE FORWARD DIFFUSION ---
sampler = comfy.samplers.KSampler(model, steps=steps, device=device)
sigmas = sampler.sigmas
start_step = steps - int(steps * strength)
if start_step >= len(sigmas): start_step = len(sigmas) - 1
sigma = sigmas[start_step].to(device)
generator = torch.Generator(device=device).manual_seed(forward_diffusion_seed)
noise = torch.randn(clean_latent.shape, generator=generator, device=device, dtype=native_dtype)
# We need to unsqueeze noise to 5D to match the inverter output if needed
if x_t_inverted.dim() == 5 and noise.dim() == 4:
noise = noise.unsqueeze(2)
x_t_ideal = clean_latent + noise * sigma
# --- PASS 3: SLERP BLEND ---
# If blend_factor is 0, just use the stable latent; if 1, use the creative one.
if blend_factor == 0.0:
final_latent = x_t_ideal
elif blend_factor == 1.0:
final_latent = x_t_inverted
else:
final_latent = slerp(blend_factor, x_t_ideal, x_t_inverted)
out = latent_image.copy()
out["samples"] = final_latent.cpu()
return (out, )
class QwenRectifiedFlowInverter:
@classmethod
def INPUT_TYPES(s):
@@ -10,17 +138,25 @@ class QwenRectifiedFlowInverter:
"model": ("MODEL",),
"latent_image": ("LATENT",),
"positive": ("CONDITIONING", ),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Seed for the random velocity perturbation."}),
"steps": ("INT", {"default": 20, "min": 1, "max": 1000}),
"inversion_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"velocity_amplification": ("FLOAT", {
"default": 0.0,
"min": -0.9, # -0.9 dampens velocity to 10% of original
"max": 2.0, # 2.0 triples the velocity
"min": -0.9,
"max": 2.0,
"step": 0.05,
"tooltip": "Amplifies (>0) or dampens (<0) the predicted velocity at each step. Deterministic effect."
"tooltip": "Deterministically amplifies (>0) or dampens (<0) the predicted velocity at each step."
}),
"velocity_perturb_strength": ("FLOAT", {
"default": 0.0,
"min": 0.0,
"max": 5.0,
"step": 0.01,
"tooltip": "Adds seeded random noise to the velocity. >1.0 will heavily randomize the path."
}),
"internal_precision": (["bfloat16", "float16", "force_float32"], {"default": "bfloat16"}),
"debug_print": (["disable", "enable"], {"default": "disable"}),
"normalize_output": (["enable", "disable"], {"default": "enable", "tooltip": "Crucial for high inversion strengths. Rescales the output latent to a standard distribution (mean=0, std=1) to prevent the sampler from collapsing."}),
}
}
@@ -28,7 +164,10 @@ class QwenRectifiedFlowInverter:
FUNCTION = "invert"
CATEGORY = "Qwen/Sampling"
def invert(self, model, latent_image, positive, steps, inversion_strength, velocity_amplification, internal_precision, debug_print):
def invert(self, model, latent_image, positive, seed, steps, inversion_strength, velocity_amplification, velocity_perturb_strength, internal_precision, normalize_output):
# For simplicity in the function signature, we'll get debug_print from kwargs if it exists
debug_print = "disable" # Not including in UI for final version to keep it clean
if inversion_strength == 0.0:
return (latent_image,)
@@ -65,12 +204,29 @@ class QwenRectifiedFlowInverter:
if predicted_velocity.dim() == 4:
predicted_velocity = predicted_velocity.unsqueeze(2)
# --- NEW: Velocity Amplification ---
# This is a deterministic modification of the velocity vector.
# 1. Apply deterministic amplification
if velocity_amplification != 0.0:
predicted_velocity = predicted_velocity * (1.0 + velocity_amplification)
# --- END NEW ---
# 2. Apply stochastic (seeded) perturbation
if velocity_perturb_strength > 0:
# Use a different seed for each step to ensure uncorrelated noise
step_seed = seed + i
generator = torch.Generator(device=device).manual_seed(step_seed)
velocity_std = torch.std(predicted_velocity)
noise = torch.randn(
predicted_velocity.shape,
generator=generator,
device=device,
dtype=native_dtype
)
perturbation = noise * velocity_std * velocity_perturb_strength
predicted_velocity = predicted_velocity + perturbation
# 3. Apply the final, modified velocity
dt = t_next - t_current
x_t = x_t + predicted_velocity * dt
@@ -81,9 +237,231 @@ class QwenRectifiedFlowInverter:
pbar.update(1)
if normalize_output == "enable":
for i in range(x_t.shape[0]):
latent_item = x_t[i]
mean = torch.mean(latent_item)
std = torch.std(latent_item)
if std > 1e-5:
x_t[i] = (latent_item - mean) / std
else:
x_t[i] = latent_item - mean
out = latent_image.copy()
out["samples"] = x_t.cpu()
return (out, )
NODE_CLASS_MAPPINGS = { "QwenRectifiedFlowInverter": QwenRectifiedFlowInverter }
NODE_DISPLAY_NAME_MAPPINGS = { "QwenRectifiedFlowInverter": "Qwen Rectified Flow Inverter" }
class LatentGaussianBlur:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"latent": ("LATENT",),
"sigma": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
"blur_mode": (["Spatial Only", "Spatial and Channel"], {"default": "Spatial Only"}),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "blur_latent"
CATEGORY = "Latent/Filter"
def blur_latent(self, latent, sigma, blur_mode):
if sigma == 0.0:
return (latent,)
device = comfy.model_management.get_torch_device()
latent_tensor = latent["samples"].clone().to(device)
is_5d = latent_tensor.dim() == 5
if is_5d:
b, c, t, h, w = latent_tensor.shape
latent_4d = latent_tensor.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w)
else:
latent_4d = latent_tensor
if blur_mode == "Spatial Only":
kernel_size = int(sigma * 6) + 1
if kernel_size % 2 == 0: kernel_size += 1
blurred_4d = TF.gaussian_blur(latent_4d, kernel_size=kernel_size, sigma=sigma)
else: # "Spatial and Channel" mode
channels = latent_4d.shape[1]
latent_5d_for_conv = latent_4d.unsqueeze(1)
k_size_spatial = int(sigma * 6) + 1
if k_size_spatial % 2 == 0: k_size_spatial += 1
k_size_channel = min(channels, k_size_spatial)
if k_size_channel % 2 == 0: k_size_channel -= 1
if k_size_channel < 1: k_size_channel = 1
ax_c = torch.linspace(-(k_size_channel - 1) / 2., (k_size_channel - 1) / 2., k_size_channel, device=device)
ax_s = torch.linspace(-(k_size_spatial - 1) / 2., (k_size_spatial - 1) / 2., k_size_spatial, device=device)
gauss_c = torch.exp(-0.5 * torch.square(ax_c / sigma))
gauss_s = torch.exp(-0.5 * torch.square(ax_s / sigma))
# --- FIX: Construct 3D kernel using reshaping and broadcasting ---
# Reshape 1D vectors to the correct orientation
gauss_c = gauss_c.view(k_size_channel, 1, 1)
gauss_h = gauss_s.view(1, k_size_spatial, 1)
gauss_w = gauss_s.view(1, 1, k_size_spatial)
# Multiply them. Broadcasting handles the expansion to create the 3D kernel.
kernel_3d = gauss_c * gauss_h * gauss_w
# --- END FIX ---
kernel_3d /= torch.sum(kernel_3d)
kernel_3d = kernel_3d.view(1, 1, k_size_channel, k_size_spatial, k_size_spatial)
padding = (k_size_channel // 2, k_size_spatial // 2, k_size_spatial // 2)
blurred_5d = F.conv3d(latent_5d_for_conv, kernel_3d, padding=padding)
blurred_4d = blurred_5d.squeeze(1)
if is_5d:
blurred_tensor = blurred_4d.view(b, t, c, h, w).permute(0, 2, 1, 3, 4)
else:
blurred_tensor = blurred_4d
out = latent.copy()
out["samples"] = blurred_tensor.cpu()
return (out,)
class LatentAddNoise:
"""
Adds a configurable amount of seeded random noise to a latent tensor.
The strength is relative to the standard deviation of the input latent.
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"latent": ("LATENT",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"strength": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 10.0,
"step": 0.01,
"tooltip": "Strength of the noise. 1.0 adds noise with the same standard deviation as the latent."
}),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "add_noise"
CATEGORY = "Latent/Noise"
def add_noise(self, latent, seed, strength):
# If strength is zero, do nothing and return the original latent.
if strength == 0.0:
return (latent,)
# Get the device and the latent tensor
device = comfy.model_management.get_torch_device()
latent_tensor = latent["samples"].clone().to(device)
# Create a seeded random number generator
generator = torch.Generator(device=device).manual_seed(seed)
# Generate noise with the same shape, device, and dtype as the latent
noise = torch.randn(
latent_tensor.shape,
generator=generator,
device=device,
dtype=latent_tensor.dtype
)
# Scale the noise. A strength of 1.0 will make the noise have the
# same standard deviation as the original latent's signal.
latent_std = torch.std(latent_tensor)
scaled_noise = noise * latent_std * strength
# Add the scaled noise to the original latent
noised_latent = latent_tensor + scaled_noise
# Package the result for output
out = latent.copy()
out["samples"] = noised_latent.cpu()
return (out,)
class LatentForwardDiffusion:
"""
Applies the 'natural' forward diffusion process to a clean latent.
This produces a statistically 'perfect' noisy latent that samplers expect.
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"model": ("MODEL",),
"latent": ("LATENT",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"noise_strength": ("FLOAT", {
"default": 0.8,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "The point in the schedule to noise to. Must match the KSampler's effective start step."
}),
}
}
RETURN_TYPES = ("LATENT",)
FUNCTION = "add_scheduled_noise"
CATEGORY = "Latent/Noise"
def add_scheduled_noise(self, model, latent, seed, steps, noise_strength):
if noise_strength == 0.0:
return (latent,)
device = comfy.model_management.get_torch_device()
latent_tensor = latent["samples"].clone().to(device)
# Get the noise schedule (sigmas) from the model
sampler = comfy.samplers.KSampler(model, steps=steps, device=device)
sigmas = sampler.sigmas
# Find the sigma corresponding to our desired noise strength
# The start step is calculated from the end of the schedule
start_step = steps - int(steps * noise_strength)
if start_step >= len(sigmas):
start_step = len(sigmas) - 1
sigma = sigmas[start_step].to(device)
# Generate pure Gaussian noise
generator = torch.Generator(device=device).manual_seed(seed)
noise = torch.randn(
latent_tensor.shape,
generator=generator,
device=device,
dtype=latent_tensor.dtype
)
# Apply the forward diffusion formula: x_t = x_0 + sigma * noise
# This is a simplified form for noise schedules where signal scaling is 1
noised_latent = latent_tensor + noise * sigma
out = latent.copy()
out["samples"] = noised_latent.cpu()
return (out,)
# --------------------------------------------------------------------
# ---------------------- NODE MAPPINGS -------------------------------
# --------------------------------------------------------------------
NODE_CLASS_MAPPINGS = {
"QwenRectifiedFlowInverter": QwenRectifiedFlowInverter,
"LatentGaussianBlur": LatentGaussianBlur,
"LatentAddNoise": LatentAddNoise,
"LatentForwardDiffusion": LatentForwardDiffusion,
"LatentHybridInverter": LatentHybridInverter,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"QwenRectifiedFlowInverter": "Qwen Rectified Flow Inverter",
"LatentGaussianBlur": "Latent Gaussian Blur",
"LatentAddNoise": "Add Latent Noise (Seeded)",
"LatentForwardDiffusion": "Forward Diffusion (Add Scheduled Noise)",
"LatentHybridInverter": "Latent Hybrid Inverter (Qwen)",
}