Initial commit: ComfyUI-ContextualRepulsion custom node

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.gitignore

@@ -0,0 +1 @@
+__pycache__/

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

__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"]

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()

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)