[Refact] Refactor diffusion output streaming (#5166)

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

tests/diffusion/test_diffusion_model_runner.py

@@ -98,6 +98,29 @@ class _ChunkStepPipeline:
         return output
 
 
+class _FinalOnlyStepPipeline:
+    device = torch.device("cpu")
+    supports_step_execution = True
+
+    def prepare_encode(self, state):
+        state.prompt_embeds = torch.zeros(1, 1, 1)
+        state.latents = torch.zeros(1, 1)
+        state.timesteps = torch.tensor([1.0])
+        state.step_index = 0
+        return state
+
+    def denoise_step(self, input_batch, states):
+        del states
+        return torch.ones_like(input_batch.latents)
+
+    def step_scheduler(self, state, noise_pred):
+        state.latents = noise_pred
+        state.step_index += 1
+
+    def post_decode(self, state):
+        return DiffusionOutput(output=state.latents.clone())
+
+
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -228,6 +251,35 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     assert fourth.get_request_output("req").finished is True
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_stepwise_streaming_decodes_final_only_pipeline(monkeypatch):
+    """Step streaming still returns a final output when no chunk boundaries exist."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _FinalOnlyStepPipeline()
+    runner.od_config.streaming_output = True
+    runner.od_config.step_execution = True
+    req = _make_request()
+    req.request_id = "req"
+    req.sampling_params.num_inference_steps = 1
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req)],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+
+    output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output).get_request_output("req")
+
+    assert output.finished is True
+    assert output.result is not None
+    assert torch.equal(output.result.output, torch.ones(1, 1))
+
+
 @pytest.mark.core_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):