ParkSeongHyunandClaude Opus 4.6
ca612bfff0
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
152 lines
5.5 KiB
Python
152 lines
5.5 KiB
Python
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()
|