[Diffusion] Qwen image support (#82)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import multiprocessing as mp
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.registry import get_diffusion_post_process_func
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.scheduler import scheduler
+from vllm_omni.diffusion.worker.gpu_worker import WorkerProc
+
+logger = init_logger(__name__)
+
+
+class DiffusionEngine:
+    """The diffusion engine for vLLM-Omni diffusion models."""
+
+    def __init__(self, od_config: OmniDiffusionConfig):
+        """Initialize the diffusion engine.
+
+        Args:
+            config: The configuration for the diffusion engine.
+        """
+        self.od_config = od_config
+
+        self.post_process_func = get_diffusion_post_process_func(od_config)
+
+        self._make_client()
+
+    def step(self, requests: list[OmniDiffusionRequest]):
+        try:
+            output = self.add_req_and_wait_for_response(requests)
+            if output.error:
+                raise Exception(f"{output.error}")
+            logger.info("Generation completed successfully.")
+            return self.post_process_func(output.output)
+        except Exception as e:
+            logger.error(f"Generation failed: {e}")
+            return None
+
+    @staticmethod
+    def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine":
+        """Factory method to create a DiffusionEngine instance.
+
+        Args:
+            config: The configuration for the diffusion engine.
+
+        Returns:
+            An instance of DiffusionEngine.
+        """
+        return DiffusionEngine(config)
+
+    def _make_client(self):
+        # TODO rename it
+        scheduler.initialize(self.od_config)
+
+        # Get the broadcast handle from the initialized scheduler
+        broadcast_handle = scheduler.get_broadcast_handle()
+
+        processes, result_handle = self._launch_workers(
+            broadcast_handle=broadcast_handle,
+        )
+
+        if result_handle is not None:
+            scheduler.initialize_result_queue(result_handle)
+        else:
+            logger.error("Failed to get result queue handle from workers")
+
+    def _launch_workers(self, broadcast_handle):
+        od_config = self.od_config
+        logger.info("Starting server...")
+
+        num_gpus = od_config.num_gpus
+        mp.set_start_method("spawn", force=True)
+        processes = []
+
+        # Launch all worker processes
+        scheduler_pipe_readers = []
+        scheduler_pipe_writers = []
+
+        for i in range(num_gpus):
+            reader, writer = mp.Pipe(duplex=False)
+            scheduler_pipe_writers.append(writer)
+            process = mp.Process(
+                target=WorkerProc.worker_main,
+                args=(
+                    i,  # rank
+                    od_config,
+                    writer,
+                    broadcast_handle,
+                ),
+                name=f"DiffusionWorker-{i}",
+                daemon=True,
+            )
+            scheduler_pipe_readers.append(reader)
+            process.start()
+            processes.append(process)
+
+        # Wait for all workers to be ready
+        scheduler_infos = []
+        result_handle = None
+        for writer in scheduler_pipe_writers:
+            writer.close()
+
+        for i, reader in enumerate(scheduler_pipe_readers):
+            try:
+                data = reader.recv()
+            except EOFError:
+                logger.error(f"Rank {i} scheduler is dead. Please check if there are relevant logs.")
+                processes[i].join()
+                logger.error(f"Exit code: {processes[i].exitcode}")
+                raise
+
+            if data["status"] != "ready":
+                raise RuntimeError("Initialization failed. Please see the error messages above.")
+
+            if i == 0:
+                result_handle = data.get("result_handle")
+
+            scheduler_infos.append(data)
+            reader.close()
+
+        logger.debug("All workers are ready")
+
+        return processes, result_handle
+
+    def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
+        return scheduler.add_req(requests)