[Core] Supports stage abstraction in the diffusion model (#391)

Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig
 from vllm_omni.diffusion.registry import get_diffusion_post_process_func, get_diffusion_pre_process_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import Scheduler, scheduler
+from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.platform_utils import get_diffusion_worker_class
 
 logger = init_logger(__name__)
@@ -84,12 +85,86 @@ class DiffusionEngine:
                 raise Exception(f"{output.error}")
             logger.info("Generation completed successfully.")
 
+            if output.output is None:
+                logger.warning("Output is None, returning empty OmniRequestOutput")
+                # Return empty output for the first request
+                if len(requests) > 0:
+                    request = requests[0]
+                    request_id = request.request_id or ""
+                    prompt = request.prompt
+                    if isinstance(prompt, list):
+                        prompt = prompt[0] if prompt else None
+                    return OmniRequestOutput.from_diffusion(
+                        request_id=request_id,
+                        images=[],
+                        prompt=prompt,
+                        metrics={},
+                        latents=None,
+                    )
+                return None
+
             postprocess_start_time = time.time()
-            result = self.post_process_func(output.output) if self.post_process_func is not None else output.output
+            images = self.post_process_func(output.output) if self.post_process_func is not None else output.output
             postprocess_time = time.time() - postprocess_start_time
             logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds")
 
-            return result
+            # Convert to OmniRequestOutput format
+            # Ensure images is a list
+            if not isinstance(images, list):
+                images = [images] if images is not None else []
+
+            # Handle single request or multiple requests
+            if len(requests) == 1:
+                # Single request: return single OmniRequestOutput
+                request = requests[0]
+                request_id = request.request_id or ""
+                prompt = request.prompt
+                if isinstance(prompt, list):
+                    prompt = prompt[0] if prompt else None
+
+                metrics = {}
+                if output.trajectory_timesteps is not None:
+                    metrics["trajectory_timesteps"] = output.trajectory_timesteps
+
+                return OmniRequestOutput.from_diffusion(
+                    request_id=request_id,
+                    images=images,
+                    prompt=prompt,
+                    metrics=metrics,
+                    latents=output.trajectory_latents,
+                )
+            else:
+                # Multiple requests: return list of OmniRequestOutput
+                # Split images based on num_outputs_per_prompt for each request
+                results = []
+                image_idx = 0
+
+                for request in requests:
+                    request_id = request.request_id or ""
+                    prompt = request.prompt
+                    if isinstance(prompt, list):
+                        prompt = prompt[0] if prompt else None
+
+                    # Get images for this request
+                    num_outputs = request.num_outputs_per_prompt
+                    request_images = images[image_idx : image_idx + num_outputs] if image_idx < len(images) else []
+                    image_idx += num_outputs
+
+                    metrics = {}
+                    if output.trajectory_timesteps is not None:
+                        metrics["trajectory_timesteps"] = output.trajectory_timesteps
+
+                    results.append(
+                        OmniRequestOutput.from_diffusion(
+                            request_id=request_id,
+                            images=request_images,
+                            prompt=prompt,
+                            metrics=metrics,
+                            latents=output.trajectory_latents,
+                        )
+                    )
+
+                return results
         except Exception as e:
             logger.error(f"Generation failed: {e}")
             return None