Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 49 additions & 1 deletion src/streamdiffusion/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,54 @@ def _extract_prepare_params(config: Dict[str, Any]) -> Dict[str, Any]:
return prepare_params


def dedupe_controlnet_configs(configs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Collapse duplicate ControlNet ``model_id`` entries, keeping the highest
``conditioning_scale`` per model_id and preserving first-occurrence order.

Mirrors the dedup semantic TouchDesigner's ``Cnblock()`` already applies on its
live-update path, so the same model loaded twice at startup (or requested twice in
a single live update) behaves the same way the course teaches: highest weight wins.
Ties (equal scale) keep the first occurrence. Used both at config-prepare time
(`_prepare_controlnet_configs`) and by the stream updater on incoming desired
configs, so the rule lives in exactly one place.
"""
best_by_model: Dict[str, Dict[str, Any]] = {}
order: List[str] = []
for cfg in configs:
model_id = cfg.get("model_id")
if model_id is None:
# No model_id to key on — keep as-is, never collapse.
sentinel_key = f"__no_model_id_{id(cfg)}__"
order.append(sentinel_key)
best_by_model[sentinel_key] = cfg
continue
existing = best_by_model.get(model_id)
if existing is None:
order.append(model_id)
best_by_model[model_id] = cfg
else:
existing_scale = existing.get("conditioning_scale", 1.0)
new_scale = cfg.get("conditioning_scale", 1.0)
if new_scale > existing_scale:
logger.info(
"dedupe_controlnet_configs: dropping duplicate ControlNet model_id=%s "
"(conditioning_scale=%s), keeping conditioning_scale=%s",
model_id,
existing_scale,
new_scale,
)
best_by_model[model_id] = cfg
else:
logger.info(
"dedupe_controlnet_configs: dropping duplicate ControlNet model_id=%s "
"(conditioning_scale=%s), keeping conditioning_scale=%s",
model_id,
new_scale,
existing_scale,
)
return [best_by_model[key] for key in order]


def _prepare_controlnet_configs(config: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Prepare ControlNet configurations for wrapper"""
controlnet_configs = []
Expand Down Expand Up @@ -283,7 +331,7 @@ def _prepare_controlnet_configs(config: Dict[str, Any]) -> List[Dict[str, Any]]:

controlnet_configs.append(controlnet_config)

return controlnet_configs
return dedupe_controlnet_configs(controlnet_configs)


def _prepare_ipadapter_configs(config: Dict[str, Any]) -> List[Dict[str, Any]]:
Expand Down
31 changes: 30 additions & 1 deletion src/streamdiffusion/stream_parameter_updater.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import torch
import torch.nn.functional as F

from .config import dedupe_controlnet_configs
from .param_schema import (
PromptInterpolationMethod,
SeedInterpolationMethod,
Expand Down Expand Up @@ -1479,12 +1480,32 @@ def _update_controlnet_config(self, desired_config: List[Dict[str, Any]]) -> Non
)
return

current_config = self._get_current_controlnet_config()
# Dedup the incoming desired config first. Without this, a caller that hands us
# two entries for a model that isn't currently loaded produces two add_controlnet
# calls below (existing_index is None for both, since current_models is only
# refreshed at the top of this method) — i.e. this method can *create* duplicates,
# not just fail to clean up ones created elsewhere.
desired_config = dedupe_controlnet_configs(desired_config)

# Simple approach: detect what changed and apply minimal updates
current_models = {
i: getattr(cn, "model_id", f"controlnet_{i}") for i, cn in enumerate(controlnet_pipeline.controlnets)
}

# Drop any duplicate model_ids already loaded in the pipeline (e.g. left over from
# a startup config that predates this dedup, or from a prior version of this
# method). Keep the first occurrence of each model_id, remove the rest — this is
# what lets an already-running stream self-heal. Must happen before the reorder
# below so current_models is recomputed over an already duplicate-free list.
seen_model_ids = set()
for i in reversed(range(len(controlnet_pipeline.controlnets))):
model_id = current_models.get(i, f"controlnet_{i}")
if model_id in seen_model_ids:
logger.info(f"_update_controlnet_config: Removing pre-existing duplicate ControlNet {model_id}")
controlnet_pipeline.remove_controlnet(i)
else:
seen_model_ids.add(model_id)

desired_models = {cfg["model_id"]: cfg for cfg in desired_config}

# Reorder to match desired order (module supports stable reordering)
Expand All @@ -1507,6 +1528,14 @@ def _update_controlnet_config(self, desired_config: List[Dict[str, Any]]) -> Non
logger.info(f"_update_controlnet_config: Removing ControlNet {model_id}")
controlnet_pipeline.remove_controlnet(i)

# Recompute current models/config after all removals above so indices line up —
# current_config captured before these mutations would be stale here and could
# read the wrong row (or raise IndexError) when used below.
current_models = {
i: getattr(cn, "model_id", f"controlnet_{i}") for i, cn in enumerate(controlnet_pipeline.controlnets)
}
current_config = self._get_current_controlnet_config()

# Add new controlnets and update existing ones
for desired_cfg in desired_config:
model_id = desired_cfg["model_id"]
Expand Down
Loading