reuse prepro orchestrator

Author: ryanontheinside

src/streamdiffusion/modules/controlnet_module.py

@@ -12,6 +12,7 @@
 from streamdiffusion.preprocessing.preprocessing_orchestrator import (
     PreprocessingOrchestrator,
 )
+from streamdiffusion.preprocessing.orchestrator_user import OrchestratorUser
 
 
 @dataclass
@@ -23,7 +24,7 @@ class ControlNetConfig:
     preprocessor_params: Optional[Dict[str, Any]] = None
 
 
-class ControlNetModule:
+class ControlNetModule(OrchestratorUser):
     """ControlNet module that provides a UNet hook for residual conditioning.
 
     Responsibilities in this step (3):
@@ -57,9 +58,8 @@ def install(self, stream) -> None:
         self.device = stream.device
         self.dtype = stream.dtype
         if self._preprocessing_orchestrator is None:
-            self._preprocessing_orchestrator = PreprocessingOrchestrator(
-                device=self.device, dtype=self.dtype, max_workers=4
-            )
+            # Enforce shared orchestrator via base helper (raises if missing)
+            self.attach_orchestrator(stream)
         # Register UNet hook
         stream.unet_hooks.append(self.build_unet_hook())
         # Expose controlnet collections so existing updater can find them

src/streamdiffusion/modules/ipadapter_module.py

@@ -6,6 +6,7 @@
 
 from streamdiffusion.hooks import EmbedsCtx, EmbeddingHook, StepCtx, UnetKwargsDelta, UnetHook
 import os
+from streamdiffusion.preprocessing.orchestrator_user import OrchestratorUser
 
 
 @dataclass
@@ -23,7 +24,7 @@ class IPAdapterConfig:
     scale: float = 1.0
 
 
-class IPAdapterModule:
+class IPAdapterModule(OrchestratorUser):
     """IP-Adapter embedding hook provider.
 
     Produces an embedding hook that concatenates cached image tokens (from
@@ -98,6 +99,9 @@ def install(self, stream) -> None:
         logger = __import__('logging').getLogger(__name__)
         style_key = self.config.style_image_key or "ipadapter_main"
 
+        # Attach shared orchestrator to ensure consistent reuse across modules
+        self.attach_orchestrator(stream)
+
         # Validate required paths
         if not self.config.ipadapter_model_path or not self.config.image_encoder_path:
             raise ValueError("IPAdapterModule.install: ipadapter_model_path and image_encoder_path are required")
@@ -128,7 +132,7 @@ def install(self, stream) -> None:
         )
         self.ipadapter = ipadapter
 
-        # Register embedding preprocessor for this style key
+        # Register embedding preprocessor for this style key 
         embedding_preprocessor = IPAdapterEmbeddingPreprocessor(
             ipadapter=ipadapter,
             device=stream.device,

src/streamdiffusion/preprocessing/orchestrator_user.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from .preprocessing_orchestrator import PreprocessingOrchestrator
+
+
+class OrchestratorUser:
+    """
+    Minimal base class to attach a shared PreprocessingOrchestrator from the stream.
+    No convenience methods; strictly enforces presence of a shared orchestrator on stream.
+    """
+
+    _preprocessing_orchestrator: Optional[PreprocessingOrchestrator] = None
+
+    def attach_orchestrator(self, stream) -> None:
+        orchestrator = getattr(stream, 'preprocessing_orchestrator', None)
+        if orchestrator is None:
+            # Lazy-create on stream once, on first user that needs it
+            orchestrator = PreprocessingOrchestrator(device=stream.device, dtype=stream.dtype, max_workers=4)
+            setattr(stream, 'preprocessing_orchestrator', orchestrator)
+        self._preprocessing_orchestrator = orchestrator
+
+

src/streamdiffusion/stream_parameter_updater.py

@@ -6,6 +6,7 @@
 
 import logging
 logger = logging.getLogger(__name__)
+from .preprocessing.orchestrator_user import OrchestratorUser
 
 class CacheStats:
     """Helper class to track cache statistics"""
@@ -20,7 +21,7 @@ def record_miss(self):
         self.misses += 1
 
 
-class StreamParameterUpdater:
+class StreamParameterUpdater(OrchestratorUser):
     def __init__(self, stream_diffusion, wrapper=None, normalize_prompt_weights: bool = True, normalize_seed_weights: bool = True):
         self.stream = stream_diffusion
         self.wrapper = wrapper  # Reference to wrapper for accessing pipeline structure
@@ -40,11 +41,15 @@ def __init__(self, stream_diffusion, wrapper=None, normalize_prompt_weights: boo
         self._seed_cache_stats = CacheStats()
 
 
+        # Attach shared orchestrator once (lazy-creates on stream if absent)
+        self.attach_orchestrator(self.stream)
+
         # IPAdapter embedding preprocessing
         self._embedding_preprocessors = []
         self._embedding_cache: Dict[str, Tuple[torch.Tensor, torch.Tensor]] = {}
         self._current_style_images: Dict[str, Any] = {}
-        self._embedding_orchestrator = None
+        # Use the shared orchestrator attached via OrchestratorUser
+        self._embedding_orchestrator = self._preprocessing_orchestrator
     def get_cache_info(self) -> Dict:
         """Get cache statistics for monitoring performance."""
         total_requests = self._prompt_cache_stats.hits + self._prompt_cache_stats.misses
@@ -100,12 +105,9 @@ def register_embedding_preprocessor(self, preprocessor: Any, style_image_key: st
             style_image_key: Unique key for the style image this preprocessor handles
         """
         if self._embedding_orchestrator is None:
-            from .preprocessing.preprocessing_orchestrator import PreprocessingOrchestrator
-            self._embedding_orchestrator = PreprocessingOrchestrator(
-                device=self.stream.device,
-                dtype=self.stream.dtype,
-                max_workers=4
-            )
+            # Ensure orchestrator is present
+            self.attach_orchestrator(self.stream)
+            self._embedding_orchestrator = self._preprocessing_orchestrator
 
         self._embedding_preprocessors.append((preprocessor, style_image_key))