[Refact] Refactor diffusion output streaming (#5166)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -11,6 +11,7 @@ import threading
 import time
 from collections.abc import AsyncGenerator, Iterable
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -33,7 +34,6 @@ from vllm_omni.diffusion.io_support import (
     supports_multimodal_input,
 )
 from vllm_omni.diffusion.output_formatter import (
-    DiffusionStepTimings,
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
     normalize_diffusion_postprocess_output,
@@ -57,6 +57,7 @@ logger = init_logger(__name__)
 
 __all__ = [
     "DiffusionEngine",
+    "DiffusionExecutionMode",
     "_RpcTask",
     "_move_tensor_tree_to_cpu",
     "get_dummy_run_num_frames",
@@ -106,6 +107,13 @@ def supports_request_batch(od_config: OmniDiffusionConfig) -> bool:
     return bool(getattr(model_cls, "supports_request_batch", False))
 
 
+def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
+    try:
+        return max(1, int(getattr(od_config, "max_num_seqs", 1)))
+    except (TypeError, ValueError):
+        return 1
+
+
 def _move_tensor_tree_to_cpu(value: object) -> object:
     if isinstance(value, torch.Tensor):
         return value.cpu() if value.device.type != "cpu" else value
@@ -130,6 +138,11 @@ class _RpcTask:
     future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future)
 
 
+class DiffusionExecutionMode(str, Enum):
+    REQUEST_BATCH = "request_batch"
+    STEP_BATCH = "step_batch"
+
+
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
@@ -149,27 +162,66 @@ class DiffusionEngine:
         """Initialize the diffusion engine.
 
         Args:
-            config: The configuration for the diffusion engine.
+            od_config: The configuration for the diffusion engine.
+            scheduler: Optional scheduler override for tests or custom engine
+                integrations. When omitted, the engine selects a scheduler
+                from the resolved execution mode.
         """
         self.od_config = od_config
 
+        self._init_process_hooks(od_config)
+        self.execution_mode = self._resolve_execution_mode(od_config)
+        self._init_executor(od_config)
+        self._init_scheduler(od_config, scheduler)
+        self._init_runtime_state()
+        self._init_execute_fn()
+        self._log_execution_mode(od_config)
+
+    def _init_process_hooks(self, od_config: OmniDiffusionConfig) -> None:
         self.post_process_func = get_diffusion_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
         # Cache whether the model-specific postprocess accepts request-level
         # sampling params so step() can support both legacy and extended hooks.
         self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")
 
+    def _resolve_execution_mode(self, od_config: OmniDiffusionConfig) -> DiffusionExecutionMode:
         self.step_execution = bool(getattr(od_config, "step_execution", False))
-        if self.od_config.streaming_output and not self.step_execution:
+        if od_config.streaming_output and not self.step_execution:
             logger.warning("streaming_output=True requires step_execution=True; enabling step execution.")
-            self.od_config.step_execution = True
+            od_config.step_execution = True
             self.step_execution = True
 
+        if self.step_execution:
+            self.supports_request_batch = False
+            return DiffusionExecutionMode.STEP_BATCH
+
+        self.supports_request_batch = supports_request_batch(od_config)
+        if not self.supports_request_batch and _max_num_seqs(od_config) > 1:
+            raise ValueError(
+                f"{getattr(od_config, 'model_class_name', None)!r} does not support request-level batching. "
+                "Use max_num_seqs=1 for serial request execution, or choose a pipeline with "
+                "supports_request_batch=True."
+            )
+        return DiffusionExecutionMode.REQUEST_BATCH
+
+    def _init_executor(self, od_config: OmniDiffusionConfig) -> None:
         executor_class = DiffusionExecutor.get_class(od_config)
         self.executor = executor_class(od_config)
-        self.scheduler: BaseScheduler = scheduler or (StepScheduler() if self.step_execution else RequestScheduler())
+
+    def _init_scheduler(
+        self,
+        od_config: OmniDiffusionConfig,
+        scheduler: BaseScheduler | None = None,
+    ) -> None:
+        if scheduler is not None:
+            self.scheduler = scheduler
+        elif self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
+            self.scheduler = StepScheduler()
+        else:
+            self.scheduler = RequestScheduler()
         self.scheduler.initialize(od_config)
-        self.supports_request_batch = False if self.step_execution else supports_request_batch(od_config)
+
+    def _init_runtime_state(self) -> None:
         self.main_loop: asyncio.AbstractEventLoop | None = None
         self.stop_event: threading.Event | None = None
         self.worker_thread: threading.Thread | None = None
@@ -181,33 +233,26 @@ class DiffusionEngine:
         # busy loop rather than via mutual exclusion.
         self._rpc_lock = threading.RLock()
         self._cv = threading.Condition(self._rpc_lock)
-        self._out_queue: dict[str, asyncio.Future] = {}
-        self._out_queue_streaming: dict[str, asyncio.Queue[DiffusionOutput]] = {}
+        self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {}
         self._closed = False
         self._shutdown_complete = False
         self.abort_queue: queue.Queue[str] = queue.Queue()
         self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
-        if self.step_execution:
-            self.execute_fn = self.executor.execute_step
-        elif self.supports_request_batch:
-            self.execute_fn = self.executor.execute_batch
-        else:
-            self.execute_fn = self.executor.execute_request
 
-        if self.supports_request_batch:
+    def _init_execute_fn(self) -> None:
+        if self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
+            self.execute_fn = self.executor.execute_step
+        else:
+            self.execute_fn = self.executor.execute_batch
+
+    def _log_execution_mode(self, od_config: OmniDiffusionConfig) -> None:
+        if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH:
             logger.info(
                 "[RequestBatch] engine init max_num_seqs=%s max_wait_ms=%s",
                 getattr(od_config, "max_num_seqs", None),
                 getattr(od_config, "request_batch_max_wait_ms", None),
             )
 
-        try:
-            self._dummy_run()
-        except Exception as e:
-            logger.error(f"Dummy run failed: {e}")
-            self.close()
-            raise e
-
     async def _check_and_start_background_loop(self):
         if self._closed:
             raise RuntimeError("DiffusionEngine is closed.")
@@ -227,24 +272,6 @@ class DiffusionEngine:
             self.worker_thread.start()
             self._loop_started = True
 
-    async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
-        await self._check_and_start_background_loop()
-
-        diffusion_engine_start_time = time.perf_counter()
-
-        # Apply pre-processing if available
-        preprocess_time = 0.0
-        if self.pre_process_func is not None:
-            preprocess_start_time = time.perf_counter()
-            request = self.pre_process_func(request)
-            preprocess_time = time.perf_counter() - preprocess_start_time
-            logger.debug("Pre-processing completed in %.4f seconds", preprocess_time)
-
-        exec_start_time = time.perf_counter()
-        output = await self.async_add_req_and_wait_for_response(request)
-        exec_total_time = time.perf_counter() - exec_start_time
-        return self.postprocess_output(request, output, diffusion_engine_start_time, preprocess_time, exec_total_time)
-
     async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[list[OmniRequestOutput], None]:
         await self._check_and_start_background_loop()
 
@@ -261,19 +288,50 @@ class DiffusionEngine:
         generator = self.async_add_req_and_stream_response(request)
         async for output in generator:
             exec_total_time = time.perf_counter() - exec_start_time
-            yield self.postprocess_output(
-                request, output, diffusion_engine_start_time, preprocess_time, exec_total_time
+            postprocess_start_time = time.perf_counter()
+            formatted_outputs = self.postprocess_output(request, output)
+            postprocess_time = time.perf_counter() - postprocess_start_time
+            step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
+            logger.debug(
+                "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, "
+                "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
+                preprocess_time * 1000,
+                exec_total_time * 1000,
+                postprocess_time * 1000,
+                step_total_ms,
             )
+            for request_output in formatted_outputs:
+                request_output.metrics.update(
+                    {
+                        "preprocess_time_ms": preprocess_time * 1000,
+                        "diffusion_engine_exec_time_ms": exec_total_time * 1000,
+                        "diffusion_engine_total_time_ms": step_total_ms,
+                        "postprocess_time_ms": postprocess_time * 1000,
+                    }
+                )
+            yield formatted_outputs
+
+    async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+        """Deprecated compatibility wrapper over ``step_streaming()``.
+
+        Use ``step_streaming()`` for new callers. This method drains the
+        unified output stream and returns only the final output batch, matching
+        the historical non-streaming ``step()`` behavior.
+        """
+        logger.warning_once(
+            "DiffusionEngine.step() is deprecated; use step_streaming() and consume the final output batch instead."
+        )
+        final_output: list[OmniRequestOutput] | None = None
+        async for output in self.step_streaming(request):
+            final_output = output
+        return final_output or []
 
     def postprocess_output(
         self,
         request: OmniDiffusionRequest,
         output: DiffusionOutput,
-        diffusion_engine_start_time: float,
-        preprocess_time: float,
-        exec_total_time: float,
     ) -> list[OmniRequestOutput]:
-        """Convert a DiffusionOutput to a list of OmniRequestOutput, attaching profiling metrics."""
+        """Convert a DiffusionOutput to a list of OmniRequestOutput."""
         if output.aborted:
             raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
         if output.error:
@@ -297,7 +355,6 @@ class DiffusionEngine:
         if self.od_config.enable_cpu_offload:
             output_data = _move_tensor_tree_to_cpu(output_data)
 
-        postprocess_start_time = time.perf_counter()
         if self.post_process_func is not None:
             # Some video pipelines need request-level controls during
             # postprocess (for example worker-side frame interpolation).
@@ -309,18 +366,6 @@ class DiffusionEngine:
             outputs = output_data
 
         postprocess_output = normalize_diffusion_postprocess_output(outputs)
-        postprocess_time = time.perf_counter() - postprocess_start_time
-        logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
-
-        step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
-        logger.debug(
-            "DiffusionEngine.step breakdown: preprocess=%.2f ms, "
-            "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
-            preprocess_time * 1000,
-            exec_total_time * 1000,
-            postprocess_time * 1000,
-            step_total_ms,
-        )
 
         return format_diffusion_outputs(
             request=request,
@@ -328,12 +373,6 @@ class DiffusionEngine:
             diffusion_output=output,
             output_data=output_data,
             postprocess_output=postprocess_output,
-            timings=DiffusionStepTimings(
-                preprocess_time_s=preprocess_time,
-                exec_time_s=exec_total_time,
-                postprocess_time_s=postprocess_time,
-                total_time_ms=step_total_ms,
-            ),
         )
 
     def _busy_loop(self):
@@ -357,16 +396,13 @@ class DiffusionEngine:
                     # Only RPC / abort work pending; loop back to drain it.
                     continue
 
-                if self.supports_request_batch:
+                if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH:
                     self._wait_for_request_batch_admission_locked()
 
                 sched_output = self.scheduler.schedule()
 
             if sched_output.is_empty:
-                if self.od_config.streaming_output:
-                    self._handle_empty_streaming_requests(sched_output.finished_req_ids)
-                else:
-                    self._handle_finished_requests(sched_output.finished_req_ids, None)
+                self._emit_finished_outputs(sched_output.finished_req_ids, None)
                 continue
 
             try:
@@ -390,14 +426,7 @@ class DiffusionEngine:
             self._process_aborts_queue()
             self._process_rpc_queue()
             finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
-            if self.od_config.streaming_output:
-                self._handle_step_streaming_runner_output(
-                    finished_req_ids,
-                    sched_output.scheduled_request_ids,
-                    runner_output,
-                )
-            else:
-                self._handle_finished_requests(finished_req_ids, runner_output)
+            self._emit_outputs(finished_req_ids, sched_output.scheduled_request_ids, runner_output)
 
         # Engine is stopping: fail any RPCs still queued so callers don't hang.
         self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down."))
@@ -509,74 +538,62 @@ class DiffusionEngine:
             if not task.future.done():
                 task.future.set_exception(exc)
 
-    def _handle_finished_requests(
+    def _emit_finished_outputs(
         self,
         finished_ids: set[str],
         runner_output: BaseRunnerOutput | None = None,
         missing_result_error: str = "Diffusion execution finished without a final output",
-    ):
+    ) -> None:
         for rid in finished_ids:
-            with self._cv:
-                fut = self._out_queue.pop(rid, None)
-            if fut is None:
-                continue
             if runner_output is not None:
                 _output = runner_output.get_request_output(rid)
             else:
                 _output = None
             out = self._finalize_finished_request(rid, _output, missing_result_error)
-            self._complete_future(fut, out)
+            self._put_output(rid, out)
 
-    def _handle_empty_streaming_requests(
+    def _emit_outputs(
         self,
         finished_ids: set[str],
-        missing_result_error: str = "Diffusion streaming request finished without execution output.",
-    ) -> None:
-        """Mirrors `_handle_finished_requests()` in non-streaming mode when used for empty scheduler output."""
-        for rid in finished_ids:
-            out = self._finalize_finished_request(rid, None, missing_result_error=missing_result_error)
-            self._put_streaming_output_with_cv(rid, out)
-
-    def _handle_step_streaming_runner_output(
-        self,
-        finished_req_ids: set[str],
         scheduled_request_ids: list[str],
         runner_output: BaseRunnerOutput,
     ) -> None:
-        """
-        Deliver partial step-execution outputs in streaming mode.
+        """Emit output chunks for every request through the unified output stream."""
+        if self.execution_mode != DiffusionExecutionMode.STEP_BATCH:
+            self._emit_finished_outputs(finished_ids, runner_output)
+            return
 
-        Step execution returns one ``RunnerOutput`` per scheduled request per
-        engine tick. Most denoise steps have ``result=None``; chunk boundaries
-        return a ``DiffusionOutput`` that must be delivered even before the
-        scheduler marks the request finished.
-        """
         delivered_finished_req_ids: set[str] = set()
 
         # finished_ids may have some requests that are not scheduler in this round.
         # First handle this-round requests.
         for request_id in scheduled_request_ids:
             req_output = runner_output.get_request_output(request_id)
-            if request_id in finished_req_ids:
+            if request_id in finished_ids:
                 # This entire request is finished (this is the last chunk)
                 out = self._finalize_finished_request(
                     request_id,
                     req_output,
-                    missing_result_error="Diffusion streaming execution finished without a final output.",
+                    missing_result_error="Diffusion step execution finished without a final output.",
                 )
-                self._put_streaming_output_with_cv(request_id, out)
+                self._put_output(request_id, out)
                 delivered_finished_req_ids.add(request_id)
             elif req_output is not None and req_output.result is not None:
-                # This is a non-terminal chunk. So it is not in scheduler's finished_req_ids, but still need delivering.
-                self._put_streaming_output_with_cv(request_id, req_output.result)
+                # This is a non-terminal chunk. It is not in scheduler's
+                # finished_ids, but it still belongs on the request stream.
+                self._put_output(request_id, req_output.result)
 
         # Then handle other requests that are finished in this round.
-        for request_id in finished_req_ids - delivered_finished_req_ids:
+        for request_id in finished_ids - delivered_finished_req_ids:
             out = self._finalize_finished_request(
                 request_id,
-                missing_result_error="Diffusion streaming request finished without execution output.",
+                missing_result_error="Diffusion step request finished without execution output.",
             )
-            self._put_streaming_output_with_cv(request_id, out)
+            self._put_output(request_id, out)
+
+    def _has_output_stream(self, request_id: str) -> bool:
+        with self._cv:
+            return request_id in self._out_streams
 
     @staticmethod
     def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
@@ -622,44 +639,31 @@ class DiffusionEngine:
 
         Args:
             config: The configuration for the diffusion engine.
+            scheduler: Optional scheduler override. When omitted, the selected
+                engine chooses the scheduler from its execution mode.
 
         Returns:
             An instance of the resolved ``DiffusionEngine`` (sub)class.
         """
         engine_class = DiffusionEngine.resolve_engine_class(config)
-        return engine_class(config, scheduler=scheduler)
+        engine = engine_class(config, scheduler=scheduler)
+        engine.run_startup_warmup()
+        return engine
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         with self._cv:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
-            if not self.od_config.streaming_output:
-                fut = self.main_loop.create_future()
-                request_id = self.scheduler.add_request(request)
-                self._out_queue[request_id] = fut
-            else:
-                queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-                request_id = self.scheduler.add_request(request)
-                self._out_queue_streaming[request_id] = queue
+            queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
+            request_id = self.scheduler.add_request(request)
+            self._out_streams[request_id] = queue
             self._cv.notify_all()
 
         return request_id
 
-    async def get_result(self, request_id: str) -> DiffusionOutput:
-        fut = self._out_queue.get(request_id)
-
-        if fut is None:
-            raise RuntimeError(f"Request {request_id} not found in output queue.")
-        try:
-            return await fut
-        except Exception as e:
-            logger.error(f"Wait for response failed: {e}")
-            raise
-
-    async def get_streaming_result(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
-        """Mirrors `get_result()` in non-streaming mode."""
+    async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
         with self._cv:
-            queue = self._out_queue_streaming.get(request_id)
+            queue = self._out_streams.get(request_id)
         if queue is None:
             raise RuntimeError(f"Request {request_id} not found in output queue.")
         try:
@@ -672,22 +676,31 @@ class DiffusionEngine:
             logger.error(f"Wait for response failed: {e}")
             raise
         finally:
-            # In streaming mode, an output queue is maintained until the terminal chunk is met.
-            # So unlike the non-streaming mode where output Future is popped in `_handle_finished_requests` (immediately
-            # after the request is returned), the streaming mode needs to pop the output queue here (one layer above).
             with self._cv:
-                if self._out_queue_streaming.get(request_id) is queue:
-                    self._out_queue_streaming.pop(request_id, None)
-
-    async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
-        # No lock needed: add_request is already protected by self._cv, and
-        # all executor calls are serialized inside the busy loop.
-        request_id = self.add_request(request)
-        return await self.get_result(request_id)
+                if self._out_streams.get(request_id) is queue:
+                    self._out_streams.pop(request_id, None)
 
     def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]:
         request_id = self.add_request(request)
-        return self.get_streaming_result(request_id)
+        return self.get_output_stream(request_id)
+
+    async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
+        """Deprecated compatibility wrapper over ``async_add_req_and_stream_response()``.
+
+        Use ``async_add_req_and_stream_response()`` for new callers. This
+        method drains the unified output stream and returns only the final
+        ``DiffusionOutput``, matching the historical non-streaming behavior.
+        """
+        logger.warning_once(
+            "DiffusionEngine.async_add_req_and_wait_for_response() is deprecated; "
+            "use async_add_req_and_stream_response() and consume the final output instead."
+        )
+        final_output: DiffusionOutput | None = None
+        async for output in self.async_add_req_and_stream_response(request):
+            final_output = output
+        if final_output is None:
+            raise RuntimeError("Diffusion execution completed without an output.")
+        return final_output
 
     def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         with self._rpc_lock:
@@ -759,6 +772,14 @@ class DiffusionEngine:
             if is_start:
                 raise RuntimeError(f"Could not {action} profiler: {e}") from e
 
+    def run_startup_warmup(self) -> None:
+        try:
+            self._dummy_run()
+        except Exception as e:
+            logger.error(f"Dummy run failed: {e}")
+            self.close()
+            raise e
+
     def _dummy_run(self):
         """A dummy run to warm up the model."""
         num_inference_steps = 1
@@ -891,8 +912,8 @@ class DiffusionEngine:
     ) -> Any:
         """Async variant of :meth:`collective_rpc` for event-loop callers.
 
-        Mirrors :meth:`async_add_req_and_wait_for_response`: enqueue a task
-        keyed by a future and ``await`` the result without blocking the loop.
+        Enqueue a task keyed by a future and ``await`` the result without
+        blocking the loop.
         """
         await self._check_and_start_background_loop()
         task = self._submit_rpc(method, timeout, args, kwargs, unique_reply_rank)
@@ -905,51 +926,26 @@ class DiffusionEngine:
             task.future.cancel()
             raise TimeoutError(f"RPC call to {method} timed out.") from exc
 
-    def _complete_future(self, fut: asyncio.Future, output: DiffusionOutput) -> None:
-        if fut.done():
-            return
-
-        def _set_result() -> None:
-            if not fut.done():
-                fut.set_result(output)
-
-        try:
-            loop = fut.get_loop()
-        except AttributeError:
-            loop = self.main_loop
-
-        try:
-            running_loop = asyncio.get_running_loop()
-        except RuntimeError:
-            running_loop = None
-
-        if loop is not None and loop.is_running() and loop is not running_loop:
-            loop.call_soon_threadsafe(_set_result)
-        else:
-            _set_result()
-
-    def _put_streaming_queue_output(
+    def _put_queue_output(
         self,
         queue: asyncio.Queue[DiffusionOutput],
         output: DiffusionOutput,
     ) -> None:
-        """Append to streaming output queue in a safe event loop. Mirrors `_complete_future()` in non-streaming mode."""
         loop = self.main_loop
         if loop is not None and loop.is_running():
             loop.call_soon_threadsafe(queue.put_nowait, output)
         else:
             queue.put_nowait(output)
 
-    def _put_streaming_output_with_cv(self, request_id: str, output: DiffusionOutput) -> None:
+    def _put_output(self, request_id: str, output: DiffusionOutput) -> None:
         with self._cv:
-            queue = self._out_queue_streaming.get(request_id)
+            queue = self._out_streams.get(request_id)
         if queue is None:
             return
-        self._put_streaming_queue_output(queue, output)
+        self._put_queue_output(queue, output)
 
     def close(self) -> None:
-        pending_futures: list[asyncio.Future] = []
-        pending_streaming_queues: list[asyncio.Queue[DiffusionOutput]] = []
+        pending_streams: list[asyncio.Queue[DiffusionOutput]] = []
         with self._cv:
             if self._closed and self._shutdown_complete:
                 return
@@ -957,17 +953,13 @@ class DiffusionEngine:
                 self._closed = True
                 if self.stop_event is not None:
                     self.stop_event.set()
-                pending_futures = list(self._out_queue.values())
-                pending_streaming_queues = list(self._out_queue_streaming.values())
-                self._out_queue.clear()
-                self._out_queue_streaming.clear()
+                pending_streams = list(self._out_streams.values())
+                self._out_streams.clear()
                 self._cv.notify_all()
 
         closed_output = DiffusionOutput(error="DiffusionEngine is closed.")
-        for fut in pending_futures:
-            self._complete_future(fut, closed_output)
-        for streaming_queue in pending_streaming_queues:
-            self._put_streaming_queue_output(streaming_queue, closed_output)
+        for stream in pending_streams:
+            self._put_queue_output(stream, closed_output)
 
         worker_thread = self.worker_thread
         if worker_thread is not None: