diff --git a/vllm_omni/engine/diffusion_engine.py b/vllm_omni/engine/diffusion_engine.py index ff5875cea..ca9d4e51c 100644 --- a/vllm_omni/engine/diffusion_engine.py +++ b/vllm_omni/engine/diffusion_engine.py @@ -1,5 +1,6 @@ from typing import Any, Optional +from diffusers import DiffusionPipeline from ..config import DiTConfig from ..worker.gpu_diffusion_worker import ( @@ -14,21 +15,10 @@ class DiffusionEngine: def __init__( self, dit_config: DiTConfig, - model_path: Optional[str] = None, log_stats: bool = False, - multiprocess_mode: 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 @@ -40,9 +30,12 @@ class DiffusionEngine: self.log_stats = log_stats self.multiprocess_mode = multiprocess_mode - def load_model(self): - raise NotImplementedError - + def load_model(self) -> DiffusionPipeline: + """Load the diffusion model pipeline.""" + # Placeholder implementation + # In a full implementation, this would load the actual model + return DiffusionPipeline.from_pretrained(self.model_config.model_name) + def generate(self, *args, **kwargs): """Generate images based on the input prompt.""" # Placeholder implementation @@ -58,74 +51,33 @@ class DiffusersPipelineEngine(DiffusionEngine): def __init__( self, dit_config: DiTConfig, - model_path: Optional[str] = None, log_stats: bool = False, - multiprocess_mode: 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) + super().__init__(dit_config, log_stats, multiprocess_mode) self.worker = DiffusionGPUWorker( - resolved_model_path, + dit_config.model_config.model_path, pipeline_name=dit_config.diffusers_pipeline, - device=device, - dtype=dtype, + device=dit_config.device_config.device if dit_config.device_config else None, + dtype=dit_config.model_config.dtype if dit_config.model_config else None ) 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, + height=self.height, + width=self.width, + num_inference_steps=self.num_inference_steps, + guidance_scale=self.guidance_scale, negative_prompt=negative_prompt, seed=seed, image=image,