[diffusion] add doc and fix qwen-image (#96)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

docs/api/README.md

@@ -11,6 +11,7 @@ Main entry points for vLLM-omni inference and serving.
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
+- [vllm_omni.entrypoints.omni.Omni][]
 - [vllm_omni.entrypoints.omni_llm.OmniLLM][]
 - [vllm_omni.entrypoints.omni_llm.OmniStageLLM][]
 - [vllm_omni.entrypoints.omni_stage.OmniStage][]
@@ -30,6 +31,7 @@ Input data structures for multi-modal inputs.
 
 Engine classes for offline and online inference.
 
+- [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
@@ -81,6 +83,8 @@ Configuration classes.
 
 Worker classes and model runners for distributed inference.
 
+- [vllm_omni.diffusion.worker.gpu_worker.GPUWorker][]
+- [vllm_omni.diffusion.worker.gpu_worker.WorkerProc][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]

docs/models/supported_models.md

@@ -4,7 +4,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
 
 ## Model Implementation
 
-If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm-omni/model_executor/models>.
+If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm-omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
 
 ## List of Supported Models for Nvidia GPU
 
@@ -17,7 +17,8 @@ th {
 
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
-| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
+| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
+| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 
 
 ## List of Supported Models for Ascend NPU

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -20,6 +20,7 @@ from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
+from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
@@ -38,7 +39,10 @@ def get_qwen_image_post_process_func(
     od_config: OmniDiffusionConfig,
 ):
     model_name = od_config.model
-    model_path = download_weights_from_hf_specific(model_name, None, ["*"])
+    if os.path.exists(model_name):
+        model_path = model_name
+    else:
+        model_path = download_weights_from_hf_specific(model_name, None, ["*"])
     vae_config_path = os.path.join(model_path, "vae/config.json")
     with open(vae_config_path) as f:
         vae_config = json.load(f)
@@ -237,24 +241,18 @@ class QwenImagePipeline(
     ):
         super().__init__()
         self.device = get_local_device()
-        # self.hf_config = vllm_config.model_config.hf_config
-        # TODO support user custom folders
-        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained("Qwen/Qwen-Image", subfolder="scheduler")
+        model = od_config.model
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model, subfolder="scheduler")
         logger.info("Loaded Qwen-Image scheduler successfully")
-        from vllm_omni.diffusion.utils.torch_utils import (
-            set_default_torch_dtype,
-        )  # to avoid ci error
 
         with set_default_torch_dtype(torch.bfloat16):
-            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                "Qwen/Qwen-Image", subfolder="text_encoder"
-            )
+            self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(model, subfolder="text_encoder")
             logger.info("Loaded Qwen-Image text encoder successfully")
-            self.vae = AutoencoderKLQwenImage.from_pretrained("Qwen/Qwen-Image", subfolder="vae").to(self.device)
+            self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae").to(self.device)
             logger.info("Loaded Qwen-Image VAE successfully")
             self.transformer = QwenImageTransformer2DModel()
             logger.info("Initialized Qwen-Image transformer successfully.")
-            self.tokenizer = Qwen2Tokenizer.from_pretrained("Qwen/Qwen-Image", subfolder="tokenizer")
+            self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer")
             logger.info("Loaded Qwen-Image tokenizer successfully.")
 
         self.stage = None

vllm_omni/entrypoints/omni.py

@@ -1,14 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-# unify OmniLLM and OmniDiffusion entrypoints
-
 from vllm_omni.diffusion.omni_diffusion import OmniDiffusion
 from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
 from vllm_omni.entrypoints.omni_llm import OmniLLM
 
 
 class Omni:
+    """Unified entrypoint for both LLM and Diffusion models for better usability."""
+
     def __init__(self, *args, **kwargs):
         model = args[0] if args else kwargs.get("model", "")
         if is_diffusion_model(model):