update ar_dit_test

vllm_omni/executor/diffusers_executor.py

@@ -0,0 +1,64 @@
+"""
+Diffusers-backed executor for DiT stages.
+
+Provides a minimal wrapper around HuggingFace diffusers pipelines to support
+text-to-image generation as the DiT stage backend.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from vllm_omni.worker.gpu_diffusion_worker import (
+    DiffusionGPUWorker,
+    DiffusionRunnerOutput,
+)
+
+
+class DiffusersPipelineExecutor:
+    """Thin executor that delegates to a diffusion worker."""
+
+    def __init__(
+        self,
+        model_path: str,
+        device: Optional[str] = None,
+        dtype: Optional[str] = None,
+        pipeline_name: Optional[str] = None,
+    ) -> None:
+        try:
+            import torch  # noqa: F401
+            import diffusers  # noqa: F401
+        except Exception as e:  # pragma: no cover
+            raise ImportError(
+                "diffusers and torch are required for DiT diffusers backend."
+            ) from e
+
+        self.worker = DiffusionGPUWorker(
+            model_path,
+            pipeline_name=pipeline_name,
+            device=device,
+            dtype=dtype,
+        )
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        height: int = 512,
+        width: int = 512,
+        num_inference_steps: int = 30,
+        guidance_scale: float = 5.0,
+        negative_prompt: Optional[str] = None,
+        seed: Optional[int] = None,
+        image: Optional[Any] = None,
+    ) -> DiffusionRunnerOutput:
+        return self.worker.generate(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            negative_prompt=negative_prompt,
+            seed=seed,
+            image=image,
+        )