diff --git a/README.md b/README.md new file mode 100644 index 0000000..42b200c --- /dev/null +++ b/README.md @@ -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. diff --git a/__init__.py b/__init__.py index d5c6899..ccbf937 100644 --- a/__init__.py +++ b/__init__.py @@ -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)", +}