From ca612bfff0dd53e75fcd4d5b5079db07ff968fbd Mon Sep 17 00:00:00 2001 From: ParkSeongHyun Date: Wed, 1 Apr 2026 15:52:50 +0000 Subject: [PATCH] Initial commit: ComfyUI-ContextualRepulsion custom node Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 1 + README.md | 75 +++++++++++++++++++++++++ __init__.py | 11 ++++ nodes.py | 151 +++++++++++++++++++++++++++++++++++++++++++++++++++ repulsion.py | 64 ++++++++++++++++++++++ 5 files changed, 302 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 repulsion.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..8e3cb48 --- /dev/null +++ b/README.md @@ -0,0 +1,75 @@ +# ComfyUI-ContextualRepulsion + +A ComfyUI custom node that increases **batch diversity** in Diffusion Transformer (DiT) models by applying repulsion to text conditioning embeddings during the denoising process. + +Based on the paper **"On-the-fly Repulsion in the Contextual Space for Rich Diversity in Diffusion Transformers"** (SIGGRAPH 2026, [arXiv:2603.28762](https://arxiv.org/abs/2603.28762)). + +## What It Does + +When generating multiple images with the same prompt (batch_size >= 2), diffusion models tend to produce very similar outputs — similar compositions, poses, styles, and colors. This is known as **typicality bias**. + +This node breaks that pattern by pushing the text conditioning embeddings apart across the batch, so each image is guided toward a different interpretation of the same prompt. The result is a set of images that are **diverse in composition and style** while still faithfully following the prompt. + +## Installation + +Copy or clone this folder into your ComfyUI `custom_nodes` directory: + +``` +ComfyUI/custom_nodes/ComfyUI-ContextualRepulsion/ +``` + +No additional dependencies required — uses only PyTorch (already included with ComfyUI). + +## Usage + +1. Set **batch_size >= 2** in the `EmptyLatentImage` node (e.g., 4) +2. Add the **Contextual Repulsion (Diversity)** node (found under `model_patches`) +3. Connect it between your model loader and the sampler: + +``` +Model Load --> [Contextual Repulsion] --> KSampler --> VAEDecode --> Preview +``` + +## Parameters + +| Parameter | Default | Range | Description | +|-----------|---------|-------|-------------| +| `repulsion_scale` | 0.1 | 0.0 - 5.0 | Controls how strongly embeddings are pushed apart. Higher = more diversity. Start with **0.05 - 0.2** for subtle variation, **0.3 - 1.0** for strong diversity. | +| `timestep_cutoff` | 0.35 | 0.0 - 1.0 | Fraction of denoising steps during which repulsion is active. **0.35** = first 35% of steps. Early steps determine composition, so this is where repulsion matters most. | +| `gradient_steps` | 3 | 1 - 10 | Number of iterative gradient descent steps per denoising step. More steps = stronger per-step effect. **1-3** is usually sufficient. | + +### Tuning Tips + +- **Start conservative**: `repulsion_scale=0.1`, `timestep_cutoff=0.35`, `gradient_steps=3` +- **Want more diversity?** Increase `repulsion_scale` first (up to ~0.5), then try raising `timestep_cutoff` +- **Getting artifacts?** Lower `repulsion_scale` or reduce `timestep_cutoff` +- **Setting `repulsion_scale=0`** bypasses the node entirely (zero overhead) +- The node only affects **conditional** embeddings — unconditional embeddings are left unchanged so CFG works correctly + +## How It Works + +The method adapts the Contextual Space Repulsion framework to single-stream DiT architectures: + +1. **Intercepts** the text conditioning embeddings (`c_crossattn`) before each model forward pass +2. **Computes** a cosine similarity kernel matrix across all batch samples +3. **Calculates** the Vendi Score diversity loss (negative von Neumann entropy of the kernel) +4. **Applies** gradient descent to push embeddings apart, increasing their angular separation +5. **Only active** during the early denoising steps (controlled by `timestep_cutoff`) where global composition is determined + +The repulsion gradient is **normalized** relative to embedding magnitude, so the `repulsion_scale` parameter behaves consistently regardless of model architecture or embedding dimensions. + +## Compatibility + +- Designed and tested with **Anima DiT** (animayume series) +- Should work with any DiT model that uses `c_crossattn` text conditioning in ComfyUI +- Requires **batch_size >= 2** (the node is a no-op for single image generation) +- Works with CFG (any scale) and various samplers/schedulers + +## References + +- Paper: [On-the-fly Repulsion in the Contextual Space for Rich Diversity in Diffusion Transformers](https://arxiv.org/abs/2603.28762) +- Project page: [contextual-repulsion.github.io](https://contextual-repulsion.github.io/) + +## License + +MIT diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3b824b4 --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .nodes import ContextualRepulsion, comfy_entrypoint + +NODE_CLASS_MAPPINGS = { + "ContextualRepulsion": ContextualRepulsion, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ContextualRepulsion": "Contextual Repulsion (Diversity)", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "comfy_entrypoint"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..48f1235 --- /dev/null +++ b/nodes.py @@ -0,0 +1,151 @@ +from typing_extensions import override + +import torch + +import comfy.model_patcher +from comfy_api.latest import ComfyExtension, io + +from .repulsion import compute_repulsion_gradient + + +class ContextualRepulsion(io.ComfyNode): + """Applies contextual repulsion to text embeddings across a batch to increase + diversity of generated images, adapted from the Contextual Space Repulsion paper.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ContextualRepulsion", + display_name="Contextual Repulsion (Diversity)", + category="model_patches", + description=( + "Increases batch diversity by applying repulsion to text conditioning " + "embeddings. Based on 'On-the-fly Repulsion in the Contextual Space' " + "(SIGGRAPH 2026). Requires batch_size >= 2." + ), + inputs=[ + io.Model.Input("model"), + io.Float.Input( + "repulsion_scale", + default=0.1, + min=0.0, + max=5.0, + step=0.01, + round=0.001, + tooltip="Repulsion strength. 0.1 = shift embeddings by ~10% of their magnitude per gradient step. Start with 0.05-0.2, increase for more diversity.", + ), + io.Float.Input( + "timestep_cutoff", + default=0.35, + min=0.0, + max=1.0, + step=0.05, + round=0.01, + tooltip="Fraction of early denoising steps to apply repulsion. 0.35 = first 35% of steps.", + ), + io.Int.Input( + "gradient_steps", + default=3, + min=1, + max=10, + tooltip="Number of iterative gradient steps per model call. More steps = stronger effect.", + ), + ], + outputs=[ + io.Model.Output(), + ], + ) + + @classmethod + def execute(cls, model, repulsion_scale, timestep_cutoff, gradient_steps) -> io.NodeOutput: + if repulsion_scale == 0.0: + return io.NodeOutput(model) + + m = model.clone() + + # Closure state for sigma progress tracking + sigma_state = {"max": None, "min": None} + + def unet_wrapper(apply_model, params): + input_x = params["input"] + timestep = params["timestep"] + c = params["c"] + cond_or_uncond = params["cond_or_uncond"] + + crossattn = c.get("c_crossattn") + if crossattn is None or len(cond_or_uncond) < 1: + return apply_model(input_x, timestep, **c) + + # --- Sigma progress estimation --- + sigma = timestep[0].item() + if sigma_state["max"] is None or sigma > sigma_state["max"]: + sigma_state["max"] = sigma + if sigma_state["min"] is None or sigma < sigma_state["min"]: + sigma_state["min"] = sigma + + sigma_range = sigma_state["max"] - sigma_state["min"] + if sigma_range > 0: + progress = (sigma_state["max"] - sigma) / sigma_range + else: + progress = 0.0 # First step — always apply + + # Skip repulsion after cutoff + if progress > timestep_cutoff: + return apply_model(input_x, timestep, **c) + + # --- Identify conditional embeddings --- + # cond_or_uncond: list like [0, 1] where 0=cond, 1=uncond + # crossattn is batched: each chunk has batch_size samples + batch_chunks = len(cond_or_uncond) + total_samples = crossattn.shape[0] + chunk_size = total_samples // batch_chunks + + # Need at least 2 conditional samples for repulsion + if chunk_size < 2: + return apply_model(input_x, timestep, **c) + + # Find which chunk indices are conditional (cond_or_uncond == 0) + cond_indices = [i for i, v in enumerate(cond_or_uncond) if v == 0] + if not cond_indices: + return apply_model(input_x, timestep, **c) + + # Extract conditional embeddings + # Each chunk i occupies crossattn[i*chunk_size : (i+1)*chunk_size] + cond_slices = [] + for idx in cond_indices: + start = idx * chunk_size + end = start + chunk_size + cond_slices.append((start, end)) + + cond_emb = torch.cat([crossattn[s:e] for s, e in cond_slices], dim=0) + + # --- Apply iterative repulsion --- + for _ in range(gradient_steps): + grad = compute_repulsion_gradient(cond_emb) + cond_emb = cond_emb - repulsion_scale * grad + + # --- Reassemble crossattn --- + modified = crossattn.clone() + offset = 0 + for start, end in cond_slices: + size = end - start + modified[start:end] = cond_emb[offset:offset + size] + offset += size + + c["c_crossattn"] = modified + return apply_model(input_x, timestep, **c) + + m.set_model_unet_function_wrapper(unet_wrapper) + return io.NodeOutput(m) + + +class ContextualRepulsionExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + ContextualRepulsion, + ] + + +async def comfy_entrypoint() -> ContextualRepulsionExtension: + return ContextualRepulsionExtension() diff --git a/repulsion.py b/repulsion.py new file mode 100644 index 0000000..2bbd346 --- /dev/null +++ b/repulsion.py @@ -0,0 +1,64 @@ +import torch +import torch.nn.functional as F + + +def compute_repulsion_gradient(embeddings, eps_noise=1e-4): + """Compute the normalized gradient of the Vendi Score diversity loss. + + Based on "On-the-fly Repulsion in the Contextual Space for Rich Diversity + in Diffusion Transformers" (arXiv:2603.28762). + + The gradient is normalized per-sample so that repulsion_scale has a + consistent effect regardless of embedding dimensionality. + + Args: + embeddings: Tensor of shape [B, N, D] — batch of token sequences. + eps_noise: Small noise added to break symmetry when embeddings are identical. + + Returns: + Normalized gradient tensor of shape [B, N, D]. Subtract (scale * grad) + from embeddings to increase diversity. + """ + B, N, D = embeddings.shape + if B < 2: + return torch.zeros_like(embeddings) + + with torch.inference_mode(False): + with torch.enable_grad(): + # Clone into float32 with grad tracking + emb = embeddings.detach().float().clone().requires_grad_(True) + + # Break symmetry for initially identical embeddings + noise = torch.randn_like(emb) * eps_noise + emb_noisy = emb + noise + + # Flatten: [B, N, D] -> [B, N*D] + flat = emb_noisy.reshape(B, -1) + + # Cosine similarity kernel: [B, B] + flat_norm = F.normalize(flat, dim=-1) + K = flat_norm @ flat_norm.T + + # Normalized kernel + K_tilde = K / B + + # Eigenvalues (K_tilde is real symmetric positive semi-definite) + eigenvalues = torch.linalg.eigvalsh(K_tilde) + + # Von Neumann entropy loss (negative entropy — minimize to maximize diversity) + eigenvalues_clamped = eigenvalues.clamp(min=1e-8) + loss = -(eigenvalues_clamped * eigenvalues_clamped.log()).sum() + + loss.backward() + + grad = emb.grad # [B, N, D] + + # Normalize: scale each sample's gradient to have unit norm relative + # to the embedding norm, so repulsion_scale is interpretable as a + # fraction of embedding magnitude. + emb_norms = embeddings.float().reshape(B, -1).norm(dim=-1, keepdim=True).clamp(min=1e-8) # [B, 1] + grad_norms = grad.reshape(B, -1).norm(dim=-1, keepdim=True).clamp(min=1e-8) # [B, 1] + grad = grad.reshape(B, -1) * (emb_norms / grad_norms) + grad = grad.reshape(B, N, D) + + return grad.to(embeddings.dtype)