Use runtime latent batch for batch-step range mode

This commit is contained in:
bbc-s
2026-05-14 15:27:27 +02:00
parent 9e11e499f7
commit b514c39ced
2 changed files with 90 additions and 39 deletions
+4 -4
View File
@@ -11,11 +11,11 @@ Inputs:
- `step` (used only in `range`)
- `direction` (used only in `range`)
- `manual_values` (used only in `manual_values`)
- `batch_size` (used only in `range_step_by_batch`)
`range_step_by_batch` formula:
- if `batch_size = 1`: value = `start`
- else: `auto_step = (stop - start) / (batch_size - 1)`
- uses the runtime latent batch from `EmptySD3LatentImage`
- if runtime batch is 1: value = `start`
- else: `auto_step = (stop - start) / (runtime_batch - 1)`
Examples:
- start=-4, stop=5, batch_size=10 -> step=1
@@ -28,4 +28,4 @@ Outputs:
## Wiring
1. `CheckpointLoaderSimple.model` -> `Per Sample LoRA Loader (Single Pass).model`
2. Output `model` from this node -> `KSampler.model`
3. Keep `EmptySD3LatentImage.batch_size` aligned with this node's `batch_size` when using `range_step_by_batch`.
3. Set `EmptySD3LatentImage.batch_size` normally; `range_step_by_batch` reads the runtime batch during sampling.
+86 -35
View File
@@ -109,7 +109,6 @@ class PerSampleLoraLoader:
"step": ("FLOAT", {"default": 0.1, "min": 0.000001, "max": 100.0, "step": 0.01}),
"direction": (["increment", "decrement"], {"default": "increment"}),
"manual_values": ("STRING", {"default": "0.1,0.2,0.3,0.4", "multiline": True}),
"batch_size": ("INT", {"default": 10, "min": 1, "max": 8192, "step": 1}),
}
}
@@ -146,35 +145,65 @@ class PerSampleLoraLoader:
return values
@staticmethod
def _build_range_step_by_batch(start: float, stop: float, batch_size: int) -> List[float]:
if batch_size <= 0:
raise ValueError("batch_size must be > 0")
if batch_size == 1:
return [start]
auto_step = (stop - start) / float(batch_size - 1)
return [start + (auto_step * i) for i in range(batch_size)]
def _build_range_step_by_batch_tensor(start: float, stop: float, batch_size: int, device, dtype):
if batch_size <= 1:
return torch.tensor([start], device=device, dtype=dtype)
return torch.linspace(start, stop, steps=batch_size, device=device, dtype=dtype)
@staticmethod
def _attach_per_sample_multiplier(adapter, per_sample_weights: List[float]):
def _install_runtime_batch_tracker(model, runtime_state):
previous_wrapper = model.model_options.get("model_function_wrapper")
def wrapper(model_apply, args):
cond_chunks = len(args.get("cond_or_uncond", [])) or 1
runtime_state["batch_size"] = max(1, int(args["input"].shape[0]) // cond_chunks)
if previous_wrapper is not None:
return previous_wrapper(model_apply, args)
return model_apply(args["input"], args["timestep"], **args["c"])
model.set_model_unet_function_wrapper(wrapper)
@staticmethod
def _expand_to_actual_batch(weights: torch.Tensor, actual_batch: int) -> torch.Tensor:
logical_batch = int(weights.shape[0])
if logical_batch == actual_batch:
return weights
if logical_batch > 0 and actual_batch > logical_batch:
repeats = (actual_batch + logical_batch - 1) // logical_batch
return weights.repeat(repeats)[:actual_batch]
return weights[:actual_batch]
@staticmethod
def _attach_per_sample_multiplier(
adapter,
per_sample_weights: List[float] | None,
runtime_state,
dynamic_range: tuple[float, float] | None,
base_strength: float,
):
if not hasattr(adapter, "h"):
return
original_h = adapter.h
weights_cpu = torch.tensor(per_sample_weights, dtype=torch.float32)
weights_cpu = None
if per_sample_weights is not None:
weights_cpu = torch.tensor(per_sample_weights, dtype=torch.float32)
def patched_h(x, base_out):
batch = x.shape[0]
w = weights_cpu.to(device=x.device, dtype=base_out.dtype)
n = w.shape[0]
if batch == n:
w_batch = w
elif n > 0 and batch > n:
repeats = (batch + n - 1) // n
w_batch = w.repeat(repeats)[:batch]
actual_batch = int(x.shape[0])
if dynamic_range is not None:
logical_batch = int(runtime_state.get("batch_size", actual_batch))
start, stop = dynamic_range
weights = PerSampleLoraLoader._build_range_step_by_batch_tensor(
start, stop, logical_batch, x.device, base_out.dtype
)
deltas = weights - base_strength
w_batch = PerSampleLoraLoader._expand_to_actual_batch(deltas, actual_batch)
else:
w_batch = w[:batch]
w = weights_cpu.to(device=x.device, dtype=base_out.dtype)
w_batch = PerSampleLoraLoader._expand_to_actual_batch(w, actual_batch)
old_multiplier = getattr(adapter, "multiplier", 1.0)
view_shape = [batch] + [1] * (base_out.ndim - 1)
view_shape = [actual_batch] + [1] * (base_out.ndim - 1)
adapter.multiplier = w_batch.view(*view_shape)
try:
return original_h(x, base_out)
@@ -193,19 +222,24 @@ class PerSampleLoraLoader:
step,
direction,
manual_values,
batch_size,
):
if mode == "manual_values":
values = self._parse_manual_values(manual_values)
elif mode == "range_step_by_batch":
values = self._build_range_step_by_batch(start, stop, batch_size)
values = None
else:
values = self._build_range(start, stop, step, direction)
if len(values) == 0:
if values is not None and len(values) == 0:
raise ValueError("No weights were generated.")
weights = values
base_strength = float(min(weights))
delta_weights = [float(w - base_strength) for w in weights]
dynamic_range = None
if mode == "range_step_by_batch":
base_strength = float(min(start, stop))
delta_weights = None
dynamic_range = (float(start), float(stop))
else:
weights = values
base_strength = float(min(weights))
delta_weights = [float(w - base_strength) for w in weights]
lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
lora_file = comfy.utils.load_torch_file(lora_path, safe_load=True)
@@ -224,8 +258,12 @@ class PerSampleLoraLoader:
regular_patches[key] = patch_data
new_model = model.clone()
runtime_state = {}
if dynamic_range is not None:
self._install_runtime_batch_tracker(new_model, runtime_state)
# Baseline identical to normal LoRA application (model-only):
# apply full LoRA first using strength = first requested weight.
# apply full LoRA first using strength = lowest requested weight.
# Then bypass adds per-sample deltas on top of this baseline.
new_model.add_patches(loaded, base_strength)
@@ -235,7 +273,13 @@ class PerSampleLoraLoader:
model_sd_keys = set(new_model.model.state_dict().keys())
for key, adapter in bypass_patches.items():
if key in model_sd_keys:
self._attach_per_sample_multiplier(adapter, delta_weights)
self._attach_per_sample_multiplier(
adapter,
delta_weights,
runtime_state,
dynamic_range,
base_strength,
)
model_manager.add_adapter(key, adapter, strength=1.0)
matched_bypass += 1
model_injections = model_manager.create_injections(new_model.model)
@@ -248,13 +292,20 @@ class PerSampleLoraLoader:
"This LoRA/model pair is not compatible with this per-sample bypass method."
)
used_weights = ", ".join([f"{v:.6g}" for v in weights])
used_deltas = ", ".join([f"{v:.6g}" for v in delta_weights])
used = (
f"weights=[{used_weights}] deltas=[{used_deltas}] base={base_strength:.6g} "
f"loaded_total={len(loaded)} bypass={len(bypass_patches)} regular={len(regular_patches)} "
f"matched_bypass={matched_bypass}"
)
if dynamic_range is not None:
used = (
f"weights=runtime_linspace(start={start:.6g}, stop={stop:.6g}, batch=EmptySD3LatentImage) "
f"base={base_strength:.6g} loaded_total={len(loaded)} bypass={len(bypass_patches)} "
f"regular={len(regular_patches)} matched_bypass={matched_bypass}"
)
else:
used_weights = ", ".join([f"{v:.6g}" for v in weights])
used_deltas = ", ".join([f"{v:.6g}" for v in delta_weights])
used = (
f"weights=[{used_weights}] deltas=[{used_deltas}] base={base_strength:.6g} "
f"loaded_total={len(loaded)} bypass={len(bypass_patches)} regular={len(regular_patches)} "
f"matched_bypass={matched_bypass}"
)
return (new_model, used)