update diffusion_engine

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,