diff --git a/vllm_omni/engine/diffusion_engine.py b/vllm_omni/engine/diffusion_engine.py new file mode 100644 index 000000000..9009c98f5 --- /dev/null +++ b/vllm_omni/engine/diffusion_engine.py @@ -0,0 +1,133 @@ +"""Diffusion engine adapters for diffusers-backed DiT stages.""" + +from __future__ import annotations + +from typing import Any, Optional + +from ..config import DiTConfig +from ..worker.gpu_diffusion_worker import ( + DiffusionGPUWorker, + DiffusionRunnerOutput, +) + + +class DiffusionEngine: + """Base diffusion engine storing shared DiT configuration.""" + + 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") + or model_cfg.get("model_path") + ) + else: + config_model_path = ( + getattr(model_cfg, "model", None) + or getattr(model_cfg, "model_path", 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 generate(self, *args, **kwargs): # pragma: no cover - placeholder + raise NotImplementedError("This method should be implemented in subclasses.") + + +class DiffusersPipelineEngine(DiffusionEngine): + """Adapter that invokes a diffusers pipeline via DiffusionGPUWorker.""" + + 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 + 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") + 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, + ) \ No newline at end of file