diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py
index e0b6d0e4a..94f3cf9b8 100644
--- a/vllm_omni/__init__.py
+++ b/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",
diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py
index 4a498829b..85be92fac 100644
--- a/vllm_omni/engine/output_processor.py
+++ b/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)
diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py
index 6719c62d1..0a451725e 100644
--- a/vllm_omni/entrypoints/omni_llm.py
+++ b/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
diff --git a/vllm_omni/model_executor/models/__init__.py b/vllm_omni/model_executor/models/__init__.py
index e69de29bb..27e699e12 100644
--- a/vllm_omni/model_executor/models/__init__.py
+++ b/vllm_omni/model_executor/models/__init__.py
@@ -0,0 +1 @@
+from .registry import OmniModelRegistry
\ No newline at end of file