From c8b1af780acae5c00f6a9079a74163bb8a60dedb Mon Sep 17 00:00:00 2001 From: Ken Simpson Date: Sun, 28 Sep 2025 17:45:37 -0700 Subject: [PATCH] Document node behaviour and clean comments --- README.md | 116 ++++++++++++++-------------------------------------- __init__.py | 40 ++---------------- 2 files changed, 34 insertions(+), 122 deletions(-) diff --git a/README.md b/README.md index 42b200c..48ddef5 100644 --- a/README.md +++ b/README.md @@ -1,108 +1,52 @@ -# ComfyUI Qwen Rectified Flow Tools +# ComfyUI Flow Matching Inverter Nodes -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. +Custom nodes for ComfyUI that target rectified-flow models (for example, Qwen). The nodes cover three core tasks: -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. +- integrate the model's velocity field to produce a partially inverted latent, +- add noise using the same schedule as the sampler, and +- perform small latent-domain cleanups. -## Features +Each node is available on its own, and the hybrid node chains them together for a single-drop workflow. -- **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. +## Node overview -## Nodes Included +### Qwen Rectified Flow Inverter -### 1. Qwen Rectified Flow Inverter +Runs the rectified-flow ODE backwards to re-noise an encoded latent. Supports deterministic amplification, seeded stochastic perturbation, optional output normalisation, and shares the same interface as the sampler (`steps`, `conditioning`, etc.). -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. +### Latent Gaussian Blur -- **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. +Applies a Gaussian blur directly in latent space. `Spatial Only` blurs each channel independently; `Spatial and Channel` performs a joint 3D blur across channels for stronger smoothing. -### 2. Latent Gaussian Blur +### Add Latent Noise (Seeded) -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. +Adds seeded Gaussian noise scaled by the input latent's standard deviation. Useful whenever a latent needs a controlled bump in noise without touching the schedule. -- **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. +### Forward Diffusion (Add Scheduled Noise) -### 3. Add Latent Noise (Seeded) +Uses the model's noise schedule (via `KSampler`) to add the amount of noise that corresponds to a given progress value. This is the "anchor" latent that samplers expect when starting from a given denoise level. -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. +### Latent Hybrid Inverter (Qwen) -- **Key Inputs:** - - `strength`: How much noise to add, relative to the latent's existing signal strength. - - `seed`: Makes the random noise pattern repeatable. +Convenience node that calls the inverter and forward diffusion nodes internally, then blends their outputs using spherical interpolation. It keeps separate seeds for the creative and anchor paths and exposes the blend weight so you can bias towards either latent. -### 4. Forward Diffusion (Add Scheduled Noise) +## Example workflows -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. +### Manual composition -- **Key Inputs:** - - `noise_strength`: The target noise level (0.0 to 1.0). Should match your KSampler's `denoise` value. +1. Encode an image to a latent. +2. Feed the latent into `Qwen Rectified Flow Inverter` (pick strength/amplification/perturbation). +3. Feed the clean latent into `Forward Diffusion` with the same strength but a different seed. +4. Optionally clean either latent with `Latent Gaussian Blur` or `Add Latent Noise`. +5. Blend the creative and anchor latents (e.g. `Latent Slerp`) and sample with `denoise` equal to the chosen strength. -### 5. Latent Hybrid Inverter (Qwen) +### Hybrid node quickstart -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. +1. Encode an image to a latent. +2. Run the latent through `Latent Hybrid Inverter` with your preferred strength and blend factor. +3. Sample the result with `denoise` matching the passed strength. ## 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. +1. Clone or download this repository into `ComfyUI/custom_nodes/`. +2. Restart ComfyUI to load the nodes (look under `Qwen/Sampling` and `Latent`). diff --git a/__init__.py b/__init__.py index 052f44c..11bff60 100644 --- a/__init__.py +++ b/__init__.py @@ -11,27 +11,25 @@ def slerp(val, low, high): """ dims = low.shape - # Reshape to handle batches correctly + # Flatten per batch so we can SLERP across arbitrary shapes low = low.reshape(dims[0], -1) high = high.reshape(dims[0], -1) low_norm = F.normalize(low, p=2, dim=1) high_norm = F.normalize(high, p=2, dim=1) - # In case we divide by zero + # Guard against zero-length vectors low_norm[torch.isnan(low_norm)] = 0.0 high_norm[torch.isnan(high_norm)] = 0.0 omega = torch.acos((low_norm * high_norm).sum(1)) so = torch.sin(omega) - # This handles the case where the angle is very small - # and prevents division by zero. + # Fallback to lerp when the angle is tiny to avoid numerical issues cond = so > 1e-3 res = torch.zeros_like(low) - # For items where the angle is large enough, do SLERP if cond.any(): omega_cond = omega[cond].unsqueeze(1) so_cond = so[cond].unsqueeze(1) @@ -40,11 +38,8 @@ def slerp(val, low, high): term2 = (torch.sin(val * omega_cond) / so_cond) * high[cond] res[cond] = term1 + term2 - # For items where the angle is too small, do linear interpolation if not cond.all(): res[~cond] = (1.0 - val) * low[~cond] + val * high[~cond] - # --- END FIX --- - return res.reshape(dims) class LatentHybridInverter: @@ -195,13 +190,10 @@ class QwenRectifiedFlowInverter: if predicted_velocity.dim() == 4: predicted_velocity = predicted_velocity.unsqueeze(2) - # 1. Apply deterministic amplification if velocity_amplification != 0.0: predicted_velocity = predicted_velocity * (1.0 + velocity_amplification) - # 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) @@ -217,7 +209,6 @@ class QwenRectifiedFlowInverter: 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 @@ -273,7 +264,7 @@ class LatentGaussianBlur: 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 + else: channels = latent_4d.shape[1] latent_5d_for_conv = latent_4d.unsqueeze(1) @@ -288,14 +279,10 @@ class LatentGaussianBlur: 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) @@ -339,18 +326,13 @@ class LatentAddNoise: 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, @@ -358,15 +340,11 @@ class LatentAddNoise: 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() @@ -406,19 +384,15 @@ class LatentForwardDiffusion: 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, @@ -427,18 +401,12 @@ class LatentForwardDiffusion: 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,