577 lines
23 KiB
Python
577 lines
23 KiB
Python
"""
|
|
Core OmniLLM and AsyncOmniLLM classes for multi-stage processing.
|
|
"""
|
|
|
|
from typing import List, Dict, Any, Optional, Union, Callable
|
|
from vllm.entrypoints.llm import LLM
|
|
from vllm.v1.engine.async_llm import AsyncLLM
|
|
from vllm.v1.engine.llm_engine import LLMEngine
|
|
from vllm.outputs import RequestOutput, LoRARequest
|
|
|
|
from ..config import OmniStageConfig
|
|
from .stage_manager import StageManager
|
|
from ..engine.output_processor import MultimodalOutputProcessor
|
|
from ..engine.diffusion_engine import DiffusersPipelineEngine
|
|
|
|
|
|
|
|
class OmniLLM(LLM):
|
|
"""Extended LLM supporting multiple engines and stage-based processing."""
|
|
|
|
def __init__(
|
|
self,
|
|
stage_configs: List[OmniStageConfig],
|
|
log_stats: bool = False,
|
|
**kwargs
|
|
):
|
|
# Use the first stage's model as the default model for LLM
|
|
default_model = stage_configs[0].model_path if stage_configs else "test-model"
|
|
|
|
# Track whether we should launch engines in multiprocess mode.
|
|
self.multiprocess_mode = kwargs.pop("multiprocess_mode", False)
|
|
|
|
# Fix configuration validation issues
|
|
# Ensure max_num_batched_tokens is at least as large as max_model_len
|
|
if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs:
|
|
if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']:
|
|
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
|
|
elif 'max_model_len' in kwargs:
|
|
# If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len
|
|
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
|
|
else:
|
|
# Set reasonable defaults to avoid validation errors
|
|
kwargs['max_model_len'] = 2048
|
|
kwargs['max_num_batched_tokens'] = 2048
|
|
|
|
super().__init__(model=default_model, **kwargs)
|
|
self.stage_configs = stage_configs
|
|
self.log_stats = log_stats
|
|
self.stage_manager = StageManager(stage_configs, log_stats)
|
|
self.output_processor = MultimodalOutputProcessor()
|
|
self._initialize_stage_engines()
|
|
|
|
def _initialize_stage_engines(self) -> None:
|
|
"""Initialize LLMEngine instances for each stage."""
|
|
self.stage_manager.initialize_engines()
|
|
|
|
def generate(
|
|
self,
|
|
stage_args_list: Optional[List[Dict[str, Any]]] = None,
|
|
use_tqdm: Union[bool, Callable[..., Any]] = True,
|
|
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
|
priority: Optional[List[int]] = None,
|
|
*,
|
|
prompt: Optional[str] = None,
|
|
stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None,
|
|
) -> List[RequestOutput]:
|
|
"""Main generation interface - orchestrates multi-stage processing."""
|
|
if stage_args_list is None:
|
|
if prompt is None:
|
|
raise ValueError(
|
|
"prompt must be provided when stage_args_list is not supplied"
|
|
)
|
|
stage_args_list = self._build_stage_args_from_config(
|
|
prompt, stage_overrides or {}
|
|
)
|
|
|
|
if len(stage_args_list) != len(self.stage_configs):
|
|
raise ValueError(
|
|
f"Number of stage arguments ({len(stage_args_list)}) must match "
|
|
f"number of stage configs ({len(self.stage_configs)})"
|
|
)
|
|
|
|
# Process through each stage sequentially
|
|
current_output = None
|
|
|
|
for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)):
|
|
stage_engine = self.stage_manager.get_engine(i)
|
|
|
|
# Prepare input for this stage
|
|
processed_input = self._process_stage_inputs(
|
|
stage_config, stage_args or {}, current_output
|
|
)
|
|
|
|
# Execute stage
|
|
stage_output = self._execute_stage(
|
|
stage_engine, processed_input, lora_request, priority, stage_config
|
|
)
|
|
|
|
# Update for next stage
|
|
current_output = stage_output
|
|
stage_config.stage_output = stage_output
|
|
|
|
# Process final output
|
|
final_output = self.output_processor.process_output(current_output)
|
|
return final_output
|
|
|
|
def _build_stage_args_from_config(
|
|
self,
|
|
prompt: str,
|
|
stage_overrides: Dict[int, Dict[str, Any]],
|
|
) -> List[Dict[str, Any]]:
|
|
"""Derive per-stage argument dictionaries from configuration defaults."""
|
|
stage_args: List[Dict[str, Any]] = []
|
|
for stage_config in self.stage_configs:
|
|
combined: Dict[str, Any] = dict(stage_config.default_stage_args or {})
|
|
override = stage_overrides.get(stage_config.stage_id)
|
|
if override:
|
|
combined.update(override)
|
|
if stage_config.engine_type == "AR":
|
|
combined["prompt"] = prompt
|
|
stage_args.append(combined)
|
|
return stage_args
|
|
|
|
def _process_stage_inputs(
|
|
self,
|
|
stage_config: OmniStageConfig,
|
|
stage_args: Dict[str, Any],
|
|
previous_output: Optional[Any]
|
|
) -> Dict[str, Any]:
|
|
"""Prepare input for specific stage."""
|
|
if stage_config.engine_type == "AR":
|
|
return self._process_ar_inputs(stage_config, stage_args, previous_output)
|
|
elif stage_config.engine_type == "DiT":
|
|
return self._process_dit_inputs(stage_config, stage_args, previous_output)
|
|
else:
|
|
raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}")
|
|
|
|
def _process_ar_inputs(
|
|
self,
|
|
stage_config: OmniStageConfig,
|
|
stage_args: Dict[str, Any],
|
|
previous_output: Optional[Any]
|
|
) -> Dict[str, Any]:
|
|
"""Process inputs for AR stage."""
|
|
combined = dict(stage_config.default_stage_args or {})
|
|
combined.update(stage_args)
|
|
combined.setdefault("prompt", "")
|
|
combined.setdefault("max_tokens", 100)
|
|
combined.setdefault("temperature", 0.7)
|
|
|
|
# If we have previous output (e.g., from a previous AR stage),
|
|
# we might want to use it as context
|
|
if previous_output is not None:
|
|
# Extract text from previous output if available
|
|
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
|
last_output = previous_output.outputs[-1]
|
|
if hasattr(last_output, 'text'):
|
|
combined["prompt"] = last_output.text + " " + combined["prompt"]
|
|
|
|
return combined
|
|
|
|
def _process_dit_inputs(
|
|
self,
|
|
stage_config: OmniStageConfig,
|
|
stage_args: Dict[str, Any],
|
|
previous_output: Optional[Any]
|
|
) -> Dict[str, Any]:
|
|
"""Process inputs for DiT stage."""
|
|
combined = dict(stage_config.default_stage_args or {})
|
|
combined.update(stage_args)
|
|
|
|
dit = stage_config.dit_config
|
|
if dit is not None:
|
|
combined.setdefault("height", getattr(dit, "height", 512))
|
|
combined.setdefault("width", getattr(dit, "width", 512))
|
|
combined.setdefault(
|
|
"num_inference_steps", getattr(dit, "num_inference_steps", 50)
|
|
)
|
|
combined.setdefault(
|
|
"guidance_scale", getattr(dit, "guidance_scale", 7.5)
|
|
)
|
|
else:
|
|
combined.setdefault("height", 512)
|
|
combined.setdefault("width", 512)
|
|
combined.setdefault("num_inference_steps", 50)
|
|
combined.setdefault("guidance_scale", 7.5)
|
|
|
|
# Handle image inputs if present
|
|
if "image" in stage_args:
|
|
# For now, we'll pass the image path directly
|
|
# In a full implementation, this would involve VAE encoding
|
|
combined["image"] = stage_args["image"]
|
|
|
|
# If we have previous output from an AR stage, we might want to use it
|
|
if previous_output is not None:
|
|
# Extract text from previous AR output
|
|
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
|
last_output = previous_output.outputs[-1]
|
|
if hasattr(last_output, 'text'):
|
|
combined["prompt"] = last_output.text
|
|
|
|
combined.setdefault("prompt", stage_args.get("prompt", ""))
|
|
|
|
return combined
|
|
|
|
def _execute_stage(
|
|
self,
|
|
stage_engine: Optional[Union[LLMEngine, DiffusersPipelineEngine]],
|
|
processed_input: Dict[str, Any],
|
|
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
|
priority: Optional[List[int]] = None,
|
|
stage_config: Optional[OmniStageConfig] = None,
|
|
) -> Any:
|
|
"""Execute a single stage."""
|
|
# DiT via diffusers backend
|
|
if stage_config and stage_config.engine_type == "DiT":
|
|
dit = stage_config.dit_config
|
|
if dit and getattr(dit, "use_diffusers", False):
|
|
# Lazy-init executor per stage
|
|
if not hasattr(self, "_dit_engines"):
|
|
self._dit_engines = {}
|
|
exec_inst = self._dit_engines.get(stage_config.stage_id)
|
|
if exec_inst is None:
|
|
exec_inst = DiffusersPipelineEngine(
|
|
dit_config=dit,
|
|
model_path=stage_config.model_path,
|
|
log_stats=self.log_stats,
|
|
multiprocess_mode=self.multiprocess_mode,
|
|
)
|
|
|
|
self._dit_engines[stage_config.stage_id] = exec_inst
|
|
|
|
return exec_inst.generate(
|
|
prompt=processed_input.get("prompt", ""),
|
|
height=processed_input.get("height", getattr(dit, "height", 512)),
|
|
width=processed_input.get("width", getattr(dit, "width", 512)),
|
|
num_inference_steps=processed_input.get(
|
|
"num_inference_steps", getattr(dit, "num_inference_steps", 30)
|
|
),
|
|
guidance_scale=processed_input.get(
|
|
"guidance_scale", getattr(dit, "guidance_scale", 5.0)
|
|
),
|
|
negative_prompt=processed_input.get("negative_prompt"),
|
|
seed=processed_input.get("seed"),
|
|
image=processed_input.get("image"),
|
|
)
|
|
|
|
# Use the parent LLM's generate method for AR text generation
|
|
prompt = processed_input.get("prompt", "")
|
|
max_tokens = processed_input.get("max_tokens", 100)
|
|
temperature = processed_input.get("temperature", 0.7)
|
|
|
|
# Generate using the base LLM class
|
|
from vllm.sampling_params import SamplingParams
|
|
|
|
sampling_params = SamplingParams(
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=processed_input.get("top_p", 1.0),
|
|
frequency_penalty=processed_input.get("frequency_penalty", 0.0),
|
|
presence_penalty=processed_input.get("presence_penalty", 0.0),
|
|
stop=processed_input.get("stop", None)
|
|
)
|
|
|
|
# Use the parent class's generate method
|
|
outputs = super().generate([prompt], sampling_params)
|
|
|
|
# Return the first output (we're processing one prompt at a time)
|
|
if outputs:
|
|
return outputs[0]
|
|
else:
|
|
# Fallback to mock output if generation fails
|
|
from vllm.outputs import RequestOutput, CompletionOutput
|
|
|
|
mock_output = CompletionOutput(
|
|
index=0,
|
|
text="Generation failed",
|
|
token_ids=[],
|
|
cumulative_logprob=0.0,
|
|
logprobs=None,
|
|
finish_reason="error"
|
|
)
|
|
|
|
return RequestOutput(
|
|
request_id="fallback_request",
|
|
prompt=prompt,
|
|
prompt_token_ids=[],
|
|
prompt_logprobs=None,
|
|
outputs=[mock_output],
|
|
finished=True
|
|
)
|
|
|
|
|
|
class AsyncOmniLLM(LLM):
|
|
"""Extended LLM class supporting multiple engines and stage-based processing."""
|
|
|
|
def __init__(
|
|
self,
|
|
stage_configs: List[OmniStageConfig],
|
|
log_stats: bool = False,
|
|
**kwargs
|
|
):
|
|
# Use the first stage's model for the base LLM
|
|
if stage_configs and stage_configs[0].model_path:
|
|
model = stage_configs[0].model_path
|
|
else:
|
|
model = "Qwen/Qwen3-0.6B"
|
|
|
|
# Fix configuration validation issues
|
|
# Ensure max_num_batched_tokens is at least as large as max_model_len
|
|
if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs:
|
|
if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']:
|
|
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
|
|
elif 'max_model_len' in kwargs:
|
|
# If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len
|
|
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
|
|
else:
|
|
# Set reasonable defaults to avoid validation errors
|
|
kwargs['max_model_len'] = 2048
|
|
kwargs['max_num_batched_tokens'] = 2048
|
|
|
|
super().__init__(model=model, **kwargs)
|
|
self.stage_configs = stage_configs
|
|
self.log_stats = log_stats
|
|
self.stage_manager = StageManager(stage_configs, log_stats)
|
|
self.output_processor = MultimodalOutputProcessor()
|
|
|
|
def _initialize_async_stage_engines(self) -> None:
|
|
"""Initialize AsyncLLM instances for each stage."""
|
|
self.stage_manager.initialize_async_engines()
|
|
|
|
async def generate_async(
|
|
self,
|
|
stage_args_list: Optional[List[Dict[str, Any]]] = None,
|
|
use_tqdm: Union[bool, Callable[..., Any]] = True,
|
|
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
|
priority: Optional[List[int]] = None,
|
|
*,
|
|
prompt: Optional[str] = None,
|
|
stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None,
|
|
) -> List[RequestOutput]:
|
|
"""Async generation interface - orchestrates multi-stage processing."""
|
|
|
|
if stage_args_list is None:
|
|
if prompt is None:
|
|
raise ValueError(
|
|
"prompt must be provided when stage_args_list is not supplied"
|
|
)
|
|
stage_args_list = self._build_stage_args_from_config(
|
|
prompt, stage_overrides or {}
|
|
)
|
|
|
|
if len(stage_args_list) != len(self.stage_configs):
|
|
raise ValueError(
|
|
f"Number of stage arguments ({len(stage_args_list)}) must match "
|
|
f"number of stage configs ({len(self.stage_configs)})"
|
|
)
|
|
|
|
# Process through each stage sequentially
|
|
current_output = None
|
|
|
|
for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)):
|
|
stage_engine = self.stage_manager.get_async_engine(i)
|
|
|
|
# Prepare input for this stage
|
|
processed_input = self._process_stage_inputs(
|
|
stage_config, stage_args or {}, current_output
|
|
)
|
|
|
|
# Execute stage asynchronously
|
|
stage_output = await self._execute_stage_async(
|
|
stage_engine, processed_input, lora_request, priority, stage_config
|
|
)
|
|
|
|
# Update for next stage
|
|
current_output = stage_output
|
|
stage_config.stage_output = stage_output
|
|
|
|
# Process final output
|
|
final_output = self.output_processor.process_output(current_output)
|
|
return final_output
|
|
|
|
def _build_stage_args_from_config(
|
|
self,
|
|
prompt: str,
|
|
stage_overrides: Dict[int, Dict[str, Any]],
|
|
) -> List[Dict[str, Any]]:
|
|
stage_args: List[Dict[str, Any]] = []
|
|
for stage_config in self.stage_configs:
|
|
combined: Dict[str, Any] = dict(stage_config.default_stage_args or {})
|
|
override = stage_overrides.get(stage_config.stage_id)
|
|
if override:
|
|
combined.update(override)
|
|
if stage_config.engine_type == "AR":
|
|
combined["prompt"] = prompt
|
|
stage_args.append(combined)
|
|
return stage_args
|
|
|
|
def _process_stage_inputs(
|
|
self,
|
|
stage_config: OmniStageConfig,
|
|
stage_args: Dict[str, Any],
|
|
previous_output: Optional[Any]
|
|
) -> Dict[str, Any]:
|
|
"""Prepare input for specific stage (same as OmniLLM)."""
|
|
if stage_config.engine_type == "AR":
|
|
return self._process_ar_inputs(stage_config, stage_args, previous_output)
|
|
elif stage_config.engine_type == "DiT":
|
|
return self._process_dit_inputs(stage_config, stage_args, previous_output)
|
|
else:
|
|
raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}")
|
|
|
|
def _process_ar_inputs(
|
|
self,
|
|
stage_config: OmniStageConfig,
|
|
stage_args: Dict[str, Any],
|
|
previous_output: Optional[Any]
|
|
) -> Dict[str, Any]:
|
|
"""Process inputs for AR stage (same as OmniLLM)."""
|
|
combined = dict(stage_config.default_stage_args or {})
|
|
combined.update(stage_args)
|
|
combined.setdefault("prompt", "")
|
|
combined.setdefault("max_tokens", 100)
|
|
combined.setdefault("temperature", 0.7)
|
|
|
|
if previous_output is not None:
|
|
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
|
last_output = previous_output.outputs[-1]
|
|
if hasattr(last_output, 'text'):
|
|
combined["prompt"] = last_output.text + " " + combined["prompt"]
|
|
|
|
return combined
|
|
|
|
def _process_dit_inputs(
|
|
self,
|
|
stage_config: OmniStageConfig,
|
|
stage_args: Dict[str, Any],
|
|
previous_output: Optional[Any]
|
|
) -> Dict[str, Any]:
|
|
"""Process inputs for DiT stage (same as OmniLLM)."""
|
|
combined = dict(stage_config.default_stage_args or {})
|
|
combined.update(stage_args)
|
|
|
|
dit = stage_config.dit_config
|
|
if dit is not None:
|
|
combined.setdefault("height", getattr(dit, "height", 512))
|
|
combined.setdefault("width", getattr(dit, "width", 512))
|
|
combined.setdefault(
|
|
"num_inference_steps", getattr(dit, "num_inference_steps", 50)
|
|
)
|
|
combined.setdefault(
|
|
"guidance_scale", getattr(dit, "guidance_scale", 7.5)
|
|
)
|
|
else:
|
|
combined.setdefault("height", 512)
|
|
combined.setdefault("width", 512)
|
|
combined.setdefault("num_inference_steps", 50)
|
|
combined.setdefault("guidance_scale", 7.5)
|
|
|
|
if "image" in stage_args:
|
|
combined["image"] = stage_args["image"]
|
|
|
|
if previous_output is not None:
|
|
if hasattr(previous_output, 'outputs') and previous_output.outputs:
|
|
last_output = previous_output.outputs[-1]
|
|
if hasattr(last_output, 'text'):
|
|
combined["prompt"] = last_output.text
|
|
|
|
combined.setdefault("prompt", stage_args.get("prompt", ""))
|
|
|
|
return combined
|
|
|
|
async def _execute_stage_async(
|
|
self,
|
|
stage_engine: AsyncLLM,
|
|
processed_input: Dict[str, Any],
|
|
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
|
|
priority: Optional[List[int]] = None,
|
|
stage_config: Optional[OmniStageConfig] = None,
|
|
) -> Any:
|
|
"""Execute a single stage asynchronously."""
|
|
# DiT via diffusers backend (sync call inside async for MVP)
|
|
if stage_config and stage_config.engine_type == "DiT":
|
|
dit = stage_config.dit_config
|
|
if dit and getattr(dit, "use_diffusers", False):
|
|
if not hasattr(self, "_dit_engines"):
|
|
self._dit_engines = {}
|
|
exec_inst = self._dit_engines.get(stage_config.stage_id)
|
|
if exec_inst is None:
|
|
from vllm_omni.engine.diffusion_engine import (
|
|
DiffusersPipelineEngine,
|
|
)
|
|
|
|
pipeline_name = getattr(dit, "diffusers_pipeline", None)
|
|
device_cfg = getattr(dit, "device_config", None)
|
|
model_cfg = getattr(dit, "model_config", None)
|
|
if isinstance(device_cfg, dict):
|
|
device = device_cfg.get("device")
|
|
dtype = device_cfg.get("dtype")
|
|
else:
|
|
device = getattr(device_cfg, "device", None)
|
|
dtype = getattr(device_cfg, "dtype", None)
|
|
|
|
if dtype is None:
|
|
if isinstance(model_cfg, dict):
|
|
dtype = model_cfg.get("dtype")
|
|
else:
|
|
dtype = getattr(model_cfg, "dtype", None)
|
|
|
|
exec_inst = DiffusersPipelineEngine(
|
|
model_path=stage_config.model_path,
|
|
pipeline_name=pipeline_name,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
self._dit_engines[stage_config.stage_id] = exec_inst
|
|
|
|
return exec_inst.generate(
|
|
prompt=processed_input.get("prompt", ""),
|
|
height=processed_input.get("height", getattr(dit, "height", 512)),
|
|
width=processed_input.get("width", getattr(dit, "width", 512)),
|
|
num_inference_steps=processed_input.get(
|
|
"num_inference_steps", getattr(dit, "num_inference_steps", 30)
|
|
),
|
|
guidance_scale=processed_input.get(
|
|
"guidance_scale", getattr(dit, "guidance_scale", 5.0)
|
|
),
|
|
negative_prompt=processed_input.get("negative_prompt"),
|
|
seed=processed_input.get("seed"),
|
|
image=processed_input.get("image"),
|
|
)
|
|
|
|
# Use the parent LLM's generate method for AR text generation
|
|
prompt = processed_input.get("prompt", "")
|
|
max_tokens = processed_input.get("max_tokens", 100)
|
|
temperature = processed_input.get("temperature", 0.7)
|
|
|
|
# Generate using the base LLM class
|
|
from vllm.sampling_params import SamplingParams
|
|
|
|
sampling_params = SamplingParams(
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=processed_input.get("top_p", 1.0),
|
|
frequency_penalty=processed_input.get("frequency_penalty", 0.0),
|
|
presence_penalty=processed_input.get("presence_penalty", 0.0),
|
|
stop=processed_input.get("stop", None)
|
|
)
|
|
|
|
# Use the parent class's generate method
|
|
outputs = super().generate([prompt], sampling_params)
|
|
|
|
# Return the first output (we're processing one prompt at a time)
|
|
if outputs:
|
|
return outputs[0]
|
|
else:
|
|
# Fallback to mock output if generation fails
|
|
from vllm.outputs import RequestOutput, CompletionOutput
|
|
|
|
mock_output = CompletionOutput(
|
|
index=0,
|
|
text="Generation failed",
|
|
token_ids=[],
|
|
cumulative_logprob=0.0,
|
|
logprobs=None,
|
|
finish_reason="error"
|
|
)
|
|
|
|
return RequestOutput(
|
|
request_id="fallback_request",
|
|
prompt=prompt,
|
|
prompt_token_ids=[],
|
|
prompt_logprobs=None,
|
|
outputs=[mock_output],
|
|
finished=True
|
|
)
|