Merge pull request #35 from tzhouam/end2end_fixup

End2end fixup

vllm_omni/__init__.py

@@ -15,7 +15,7 @@ Architecture:
 __version__ = "0.1.0"
 __author__ = "vLLM-omni Team"
 __email__ = "hsliuustc@gmail.com"
-
+from . import patch
 from .config import (
     DiTCacheConfig,
     DiTConfig,
@@ -25,7 +25,7 @@ from .config import (
 )
 
 # Main entry points
-from .entrypoints.omni_llm import AsyncOmniLLM, OmniLLM
+from .entrypoints.omni_llm import OmniLLM
 
 __all__ = [
     # Version info
@@ -34,7 +34,7 @@ __all__ = [
     "__email__",
     # Main components
     "OmniLLM",
-    "AsyncOmniLLM",
+    # AsyncOmniLLM is not supported yet
     # Configuration
     "OmniStageConfig",
     "DiTConfig",

vllm_omni/engine/output_processor.py

@@ -242,7 +242,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
     ) -> OutputProcessorOutput:
         self._reqid_to_mm_type.clear()
         for eco in engine_core_outputs:
-            mm_type = (getattr(eco, "output_type", None) or "").lower()
+            mm_type = (self.engine_core_output_type or "").lower()
             if mm_type:
                 self._reqid_to_mm_type[eco.request_id] = mm_type
             self._route_and_normalize(eco)

vllm_omni/entrypoints/omni_llm.py

@@ -163,7 +163,7 @@ class OmniStageLLM(LLM):
             engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
         )
         self.llm_engine.output_processor = MultimodalOutputProcessor(
-            tokenizer=self.llm_engine.tokenizer, log_stats=self.llm_engine.log_stats
+            tokenizer=self.llm_engine.tokenizer, log_stats=self.llm_engine.log_stats, engine_core_output_type=engine_args.engine_output_type
         )
         self.llm_engine.processor = OmniProcessor(
             vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer

vllm_omni/model_executor/models/__init__.py

@@ -0,0 +1 @@
+from .registry import OmniModelRegistry