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

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
• Project page: contextual-repulsion.github.io

License

MIT