[Feat] DiffusionEngine Support async batch infer (#2729)

Signed-off-by: Semmer <semmer@live.cn>
Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/diffusion/test_diffusion_step_pipeline.py

@@ -77,12 +77,12 @@ class _StepPipeline:
         self.prepare_calls += 1
         state.timesteps = [torch.tensor(10), torch.tensor(5)]
         state.latents = torch.tensor([0.0])
+        state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
         return state
 
-    def denoise_step(self, state, **kwargs):
-        del state, kwargs
+    def denoise_step(self, input_batch, **kwargs):
         self.denoise_calls += 1
-        return torch.tensor([1.0])
+        return torch.full_like(input_batch.prompt_embeds, fill_value=0.5)
 
     def step_scheduler(self, state, noise_pred, **kwargs):
         del noise_pred, kwargs
@@ -145,6 +145,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
         state.step_index = 0
         state.scheduler = self.scheduler
         state.do_true_cfg = self.mode == "cfg"
+        state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
         return state
 
     def denoise_step(self, state, **kwargs):
@@ -291,6 +292,7 @@ def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
     engine.scheduler = scheduler
     engine.execute_fn = execute_fn
     engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
     engine.abort_queue = queue.Queue()
     return engine
 
@@ -327,10 +329,11 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
             raise ValueError(f"Unsupported distributed test mode: {mode}")
 
         runner = _make_distributed_runner(mode, device)
-        output = DiffusionModelRunner.execute_stepwise(
+        result = DiffusionModelRunner.execute_stepwise(
             runner,
             _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
         )
+        output = result.get_req_output("req-1")
 
         assert output.finished is True
         assert output.result is not None
@@ -354,14 +357,16 @@ class TestRunner:
         req = _make_step_request()
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
-        first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        first = result.get_req_output("req-1")
         assert first.req_id == "req-1"
         assert first.step_index == 1
         assert first.finished is False
         assert first.result is None
         assert "req-1" in runner.state_cache
 
-        second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
+        result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
+        second = result.get_req_output("req-1")
         assert second.req_id == "req-1"
         assert second.step_index == 2
         assert second.finished is True
@@ -393,8 +398,8 @@ class TestRunner:
             num_waiting_reqs=0,
         )
 
-        with pytest.raises(ValueError, match="batch_size=1"):
-            DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+        result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+        assert len(result) == 2
 
     def test_rejects_missing_cached_state(self):
         runner = _make_runner()
@@ -408,8 +413,8 @@ class TestRunner:
         req = _make_step_request()
         monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
 
-        output = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-
+        result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        output = result.get_req_output("req-1")
         assert output.req_id == "req-1"
         assert output.step_index == 0
         assert output.finished is True
@@ -580,7 +585,7 @@ class TestEngine:
     )
     def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture):
         scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
         engine = _make_engine(scheduler, execute_fn=execute_fn)
 
         output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
@@ -590,7 +595,7 @@ class TestEngine:
 
     def test_step_execution_completes(self, mocker: MockerFixture):
         scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
         engine = _make_engine(scheduler)
         request = _make_engine_request("req-step", num_inference_steps=2)
 
@@ -616,7 +621,7 @@ class TestEngine:
 
     def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture):
         scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
         engine = _make_engine(scheduler)
         request = _make_engine_request("req-stop", num_inference_steps=4)
 
@@ -641,7 +646,7 @@ class TestEngine:
 
     def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture):
         scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
         engine = _make_engine(scheduler)
         request = _make_engine_request("req-mid", num_inference_steps=4)
 
@@ -668,7 +673,7 @@ class TestEngine:
 
     def test_finished_step_without_result_returns_error(self, mocker: MockerFixture):
         scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
         engine = _make_engine(
             scheduler,
             execute_fn=lambda _: RunnerOutput(