update ar_dit_test

vllm_omni/engine/diffusion_engine.py

@@ -0,0 +1,132 @@
+
+from typing import Any, Optional
+
+from ..config import DiTConfig
+from ..worker.gpu_diffusion_worker import (
+    DiffusionGPUWorker,
+    DiffusionRunnerOutput,
+)
+
+
+class DiffusionEngine:
+    "Legacy placeholder for DiffusionEngine"
+
+    def __init__(
+        self,
+        dit_config: DiTConfig,
+        model_path: Optional[str] = None,
+        log_stats: bool = False,
+        multiprocess_mode: bool = False,
+    ) -> None:
+        self.dit_config = dit_config
+
+        model_cfg = getattr(dit_config, "model_config", None)
+        config_model_path = None
+        if model_cfg is not None:
+            if isinstance(model_cfg, dict):
+                config_model_path = model_cfg.get("model")
+            else:
+                config_model_path = getattr(model_cfg, "model", None)
+
+        self.model_path = model_path or config_model_path
+        self.model_config = dit_config.model_config
+        self.cache_config = dit_config.dit_cache_config
+
+        self.height = dit_config.height
+        self.width = dit_config.width
+        self.num_inference_steps = dit_config.num_inference_steps
+        self.guidance_scale = dit_config.guidance_scale
+        
+        self.log_stats = log_stats
+        self.multiprocess_mode = multiprocess_mode
+
+    def load_model(self):
+        raise NotImplementedError
+
+    def generate(self, *args, **kwargs):
+        """Generate images based on the input prompt."""
+        # Placeholder implementation
+        # In a full implementation, this would use the loaded model to generate images
+        raise NotImplementedError("This method should be implemented in subclasses.")
+
+
+class DiffusersPipelineEngine(DiffusionEngine):
+    """Thin executor that delegates to a diffusion worker."""
+
+    # crucial: must match signature of parent class
+
+    def __init__(
+        self,
+        dit_config: DiTConfig,
+        model_path: Optional[str] = None,
+        log_stats: bool = False,
+        multiprocess_mode: bool = False,
+    ) -> None:
+        super().__init__(dit_config, model_path, log_stats, multiprocess_mode)
+
+        resolved_model_path = self.model_path or None
+        if resolved_model_path is None:
+            raise ValueError("Model path must be provided for DiffusersPipelineEngine.")
+
+        device_cfg = getattr(dit_config, "device_config", None)
+        model_cfg = getattr(dit_config, "model_config", None)
+
+        device = None
+        dtype = None
+
+        if device_cfg:
+            if isinstance(device_cfg, dict):
+                device = device_cfg.get("device")
+                dtype = device_cfg.get("dtype", None)
+            else:
+                device = getattr(device_cfg, "device", None)
+                dtype = getattr(device_cfg, "dtype", None)
+
+        if dtype is None and model_cfg:
+            if isinstance(model_cfg, dict):
+                dtype = model_cfg.get("dtype")
+            else:
+                dtype = getattr(model_cfg, "dtype", None)
+
+        self.worker = DiffusionGPUWorker(
+            resolved_model_path,
+            pipeline_name=dit_config.diffusers_pipeline,
+            device=device,
+            dtype=dtype,
+        )
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        negative_prompt: Optional[str] = None,
+        seed: Optional[int] = None,
+        image: Optional[Any] = None,
+    ) -> DiffusionRunnerOutput:
+        height = int(height) if height is not None else self.height
+        width = int(width) if width is not None else self.width
+        num_steps = (
+            int(num_inference_steps)
+            if num_inference_steps is not None
+            else self.num_inference_steps
+        )
+        guidance = (
+            float(guidance_scale)
+            if guidance_scale is not None
+            else self.guidance_scale
+        )
+
+        return self.worker.generate(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance,
+            negative_prompt=negative_prompt,
+            seed=seed,
+            image=image,
+        )