[Model] Add Dynin-omni model in vllm-omni (#1759)
Signed-off-by: kdg6245@snu.ac.kr <kdg6245@snu.ac.kr> Signed-off-by: Yejoon Lee (IPAI) <leeyejoon@snu.ac.kr> Signed-off-by: aidas (arpa-kt) <aidaslab@gmail.com> Signed-off-by: Dogeun Kim <82812668+DOGEUNNKIM@users.noreply.github.com> Co-authored-by: Yejoon Lee (IPAI) <leeyejoon@snu.ac.kr> Co-authored-by: aidas (arpa-kt) <aidaslab@gmail.com>
This commit is contained in:
18 files changed
+7072
No files matched your search
@@ -0,0 +1,59 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .dynin_omni import DyninOmniForConditionalGeneration
|
||||
from .dynin_omni_common import (
|
||||
get_dynin_magvit_attr,
|
||||
get_dynin_modeling_attr,
|
||||
get_dynin_sampling_attr,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .dynin_omni_token2audio import DyninOmniToken2Audio
|
||||
from .dynin_omni_token2image import DyninOmniToken2Image
|
||||
from .dynin_omni_token2text import DyninOmniToken2Text
|
||||
|
||||
|
||||
_STAGE_EXPORTS = {
|
||||
"DyninOmniToken2Audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"),
|
||||
"DyninOmniToken2Image": (".dynin_omni_token2image", "DyninOmniToken2Image"),
|
||||
"DyninOmniToken2Text": (".dynin_omni_token2text", "DyninOmniToken2Text"),
|
||||
}
|
||||
|
||||
_MODELING_EXPORTS = {"DyninOmniConfig", "DyninOmniModelLM", "VideoTokenMerger"}
|
||||
_MAGVIT_EXPORTS = {"VQGANEncoder", "VQGANDecoder", "LFQuantizer", "MAGVITv2"}
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _STAGE_EXPORTS:
|
||||
module_name, attr_name = _STAGE_EXPORTS[name]
|
||||
module = __import__(module_name, globals(), locals(), [attr_name], 1)
|
||||
return getattr(module, attr_name)
|
||||
|
||||
if name in _MODELING_EXPORTS:
|
||||
return get_dynin_modeling_attr(name)
|
||||
|
||||
if name in _MAGVIT_EXPORTS:
|
||||
return get_dynin_magvit_attr(name)
|
||||
|
||||
if name == "get_mask_schedule":
|
||||
return get_dynin_sampling_attr("get_mask_schedule")
|
||||
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DyninOmniForConditionalGeneration",
|
||||
"DyninOmniToken2Audio",
|
||||
"DyninOmniToken2Image",
|
||||
"DyninOmniToken2Text",
|
||||
"DyninOmniConfig",
|
||||
"DyninOmniModelLM",
|
||||
"VideoTokenMerger",
|
||||
"VQGANEncoder",
|
||||
"VQGANDecoder",
|
||||
"LFQuantizer",
|
||||
"MAGVITv2",
|
||||
"get_mask_schedule",
|
||||
]
|
||||
@@ -0,0 +1,744 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from functools import cached_property
|
||||
from importlib import import_module
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.config.multimodal import BaseDummyOptions
|
||||
from vllm.inputs import MultiModalDataDict
|
||||
from vllm.inputs import MultiModalInput as MultiModalInputs
|
||||
from vllm.model_executor.models.interfaces import SupportsMultiModal
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFieldConfig,
|
||||
MultiModalKwargsItems,
|
||||
PlaceholderRange,
|
||||
)
|
||||
from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
|
||||
from vllm.multimodal.processing import (
|
||||
BaseDummyInputsBuilder,
|
||||
BaseMultiModalProcessor,
|
||||
BaseProcessingInfo,
|
||||
ProcessorInputs,
|
||||
PromptUpdate,
|
||||
TimingContext,
|
||||
)
|
||||
from vllm.sequence import IntermediateTensors
|
||||
from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .dynin_omni_common import build_zero_input_embeddings
|
||||
|
||||
try:
|
||||
from PIL import Image as PILImage
|
||||
except Exception: # pragma: no cover
|
||||
PILImage = None
|
||||
|
||||
|
||||
_MODALITY_ORDER = ("image", "video", "audio")
|
||||
|
||||
_MODALITY_ALIASES = {
|
||||
"img2img": "image",
|
||||
}
|
||||
|
||||
_MODALITY_INPUT_KEY_BY_NAME = {
|
||||
"image": "pixel_values",
|
||||
"video": "pixel_values_videos",
|
||||
"audio": "input_audio_features",
|
||||
}
|
||||
|
||||
_MODALITY_PLACEHOLDER_BY_NAME = {
|
||||
"image": "<|soi|><|image|><|eoi|>",
|
||||
"video": "<|sov|><|video|><|eov|>",
|
||||
"audio": "<|soa|><|audio|><|eoa|>",
|
||||
}
|
||||
|
||||
_MODALITY_INPUT_ALIASES = {
|
||||
"image": ("pixel_values", "image_embeds", "img2img"),
|
||||
"video": ("pixel_values_videos", "video_embeds"),
|
||||
"audio": ("input_audio_features", "audio_embeds"),
|
||||
}
|
||||
|
||||
|
||||
def _normalize_modality_name(modality: str) -> str:
|
||||
return _MODALITY_ALIASES.get(modality, modality)
|
||||
|
||||
|
||||
def _get_modality_count(mm_counts: Mapping[str, int], modality: str) -> int:
|
||||
canonical = _normalize_modality_name(modality)
|
||||
count = mm_counts.get(canonical, 0)
|
||||
for alias, target in _MODALITY_ALIASES.items():
|
||||
if target == canonical:
|
||||
count += mm_counts.get(alias, 0)
|
||||
return count
|
||||
|
||||
|
||||
def _normalize_mm_data_aliases(mm_data: MultiModalDataDict) -> MultiModalDataDict:
|
||||
normalized: dict[str, Any] = {}
|
||||
for modality, value in mm_data.items():
|
||||
canonical = _normalize_modality_name(modality)
|
||||
if canonical in normalized and normalized[canonical] is not None and value is not None:
|
||||
raise ValueError(
|
||||
"Dynin received duplicate multimodal inputs for "
|
||||
f"{canonical!r} via {modality!r}. "
|
||||
"Provide either the canonical modality or its alias, not both."
|
||||
)
|
||||
if canonical not in normalized or normalized[canonical] is None:
|
||||
normalized[canonical] = value
|
||||
return normalized
|
||||
|
||||
|
||||
def _get_placeholder_text(modality: str) -> str | None:
|
||||
modality = _normalize_modality_name(modality)
|
||||
for base_modality, placeholder in _MODALITY_PLACEHOLDER_BY_NAME.items():
|
||||
if modality.startswith(base_modality):
|
||||
return placeholder
|
||||
return None
|
||||
|
||||
|
||||
class DyninOmniProcessingInfo(BaseProcessingInfo):
|
||||
def get_data_parser(self) -> MultiModalDataParser:
|
||||
return DyninOmniMultiModalDataParser(
|
||||
expected_hidden_size=self._get_expected_hidden_size(),
|
||||
)
|
||||
|
||||
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
||||
limits = {modality: 1 for modality in _MODALITY_ORDER}
|
||||
for alias, target in _MODALITY_ALIASES.items():
|
||||
if target in limits:
|
||||
limits[alias] = limits[target]
|
||||
return limits
|
||||
|
||||
def get_mm_max_tokens_per_item(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
) -> Mapping[str, int] | None:
|
||||
del seq_len, mm_counts
|
||||
limits = {modality: 1 for modality in _MODALITY_ORDER}
|
||||
for alias, target in _MODALITY_ALIASES.items():
|
||||
if target in limits:
|
||||
limits[alias] = limits[target]
|
||||
return limits
|
||||
|
||||
|
||||
class DyninOmniDummyInputsBuilder(BaseDummyInputsBuilder[DyninOmniProcessingInfo]):
|
||||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||||
chunks: list[str] = []
|
||||
for modality in _MODALITY_ORDER:
|
||||
placeholder = _get_placeholder_text(modality)
|
||||
if placeholder is None:
|
||||
continue
|
||||
chunks.extend([placeholder] * _get_modality_count(mm_counts, modality))
|
||||
return " ".join(chunks)
|
||||
|
||||
def get_dummy_mm_data(
|
||||
self,
|
||||
seq_len: int,
|
||||
mm_counts: Mapping[str, int],
|
||||
mm_options: Mapping[str, BaseDummyOptions] | None = None,
|
||||
) -> MultiModalDataDict:
|
||||
del seq_len
|
||||
|
||||
mm_data: dict[str, Any] = {}
|
||||
|
||||
num_images = _get_modality_count(mm_counts, "image")
|
||||
if num_images > 0:
|
||||
mm_data["image"] = self._get_dummy_images(
|
||||
width=224,
|
||||
height=224,
|
||||
num_images=num_images,
|
||||
overrides=mm_options.get("image") if mm_options else None,
|
||||
)
|
||||
|
||||
num_videos = _get_modality_count(mm_counts, "video")
|
||||
if num_videos > 0:
|
||||
mm_data["video"] = self._get_dummy_videos(
|
||||
width=224,
|
||||
height=224,
|
||||
num_frames=8,
|
||||
num_videos=num_videos,
|
||||
overrides=mm_options.get("video") if mm_options else None,
|
||||
)
|
||||
|
||||
num_audios = _get_modality_count(mm_counts, "audio")
|
||||
if num_audios > 0:
|
||||
mm_data["audio"] = self._get_dummy_audios(
|
||||
length=16000,
|
||||
num_audios=num_audios,
|
||||
overrides=mm_options.get("audio") if mm_options else None,
|
||||
)
|
||||
|
||||
return mm_data
|
||||
|
||||
|
||||
class DyninOmniMultiModalDataParser(MultiModalDataParser):
|
||||
def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems:
|
||||
normalized = _normalize_mm_data_aliases(mm_data)
|
||||
mm_items = super().parse_mm_data(normalized)
|
||||
|
||||
for alias, canonical in _MODALITY_ALIASES.items():
|
||||
if alias in mm_data and canonical in mm_items and alias not in mm_items:
|
||||
mm_items[alias] = mm_items[canonical]
|
||||
|
||||
return mm_items
|
||||
|
||||
def _get_audio_with_sr(self, audio: Any) -> tuple[np.ndarray, float | None]:
|
||||
audio_array, orig_sr = super()._get_audio_with_sr(audio)
|
||||
if self.audio_resampler.target_sr is None:
|
||||
return audio_array, None
|
||||
return audio_array, orig_sr
|
||||
|
||||
|
||||
class DyninOmniMultiModalProcessor(BaseMultiModalProcessor[DyninOmniProcessingInfo]):
|
||||
@staticmethod
|
||||
def _find_subsequence(
|
||||
haystack: list[int],
|
||||
needle: list[int],
|
||||
start: int,
|
||||
) -> int | None:
|
||||
if not needle:
|
||||
return None
|
||||
|
||||
max_start = len(haystack) - len(needle)
|
||||
if max_start < start:
|
||||
return None
|
||||
|
||||
for idx in range(start, max_start + 1):
|
||||
if haystack[idx : idx + len(needle)] == needle:
|
||||
return idx
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _make_disabled_embed_mask(length: int) -> torch.Tensor:
|
||||
return torch.zeros(length, dtype=torch.bool)
|
||||
|
||||
@staticmethod
|
||||
def _encode_prompt_to_token_ids(
|
||||
prompt: str | list[int],
|
||||
tokenizer: Any | None,
|
||||
) -> list[int]:
|
||||
if isinstance(prompt, str):
|
||||
if tokenizer is None:
|
||||
raise ValueError("Tokenizer is required to process string prompts for Dynin multimodal inputs.")
|
||||
return tokenizer.encode(prompt, add_special_tokens=False)
|
||||
return list(prompt)
|
||||
|
||||
@staticmethod
|
||||
def _ensure_non_empty_prompt_ids(
|
||||
prompt_token_ids: list[int],
|
||||
tokenizer: Any | None,
|
||||
) -> list[int]:
|
||||
if prompt_token_ids:
|
||||
return prompt_token_ids
|
||||
|
||||
fallback_id = None
|
||||
if tokenizer is not None:
|
||||
fallback_id = getattr(tokenizer, "bos_token_id", None)
|
||||
if fallback_id is None:
|
||||
fallback_id = getattr(tokenizer, "eos_token_id", None)
|
||||
if fallback_id is None:
|
||||
fallback_id = getattr(tokenizer, "pad_token_id", None)
|
||||
|
||||
return [0 if fallback_id is None else int(fallback_id)]
|
||||
|
||||
@classmethod
|
||||
def _image_to_chw_float_tensor(cls, image: Any) -> torch.Tensor:
|
||||
if isinstance(image, torch.Tensor):
|
||||
tensor = image.detach()
|
||||
elif isinstance(image, np.ndarray):
|
||||
tensor = torch.from_numpy(image)
|
||||
elif PILImage is not None and isinstance(image, PILImage.Image):
|
||||
tensor = torch.from_numpy(np.asarray(image).copy())
|
||||
else:
|
||||
raise TypeError(f"Unsupported image item type: {type(image)!r}")
|
||||
|
||||
if tensor.ndim == 2:
|
||||
tensor = tensor.unsqueeze(-1)
|
||||
if tensor.ndim != 3:
|
||||
raise ValueError(f"Expected 3D image tensor, got shape={tuple(tensor.shape)}")
|
||||
|
||||
if tensor.shape[-1] in (1, 3, 4) and tensor.shape[0] not in (1, 3, 4):
|
||||
tensor = tensor.permute(2, 0, 1)
|
||||
|
||||
if tensor.shape[0] == 1:
|
||||
tensor = tensor.repeat(3, 1, 1)
|
||||
if tensor.shape[0] == 4:
|
||||
tensor = tensor[:3]
|
||||
|
||||
tensor = tensor.to(dtype=torch.float32)
|
||||
if tensor.numel() > 0 and torch.max(tensor) > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
return tensor.contiguous()
|
||||
|
||||
@classmethod
|
||||
def _video_to_tchw_float_tensor(cls, video: Any) -> torch.Tensor:
|
||||
if isinstance(video, (list, tuple)) and not isinstance(video, torch.Tensor):
|
||||
frames = [cls._image_to_chw_float_tensor(frame) for frame in video]
|
||||
if not frames:
|
||||
return torch.zeros((1, 3, 1, 1), dtype=torch.float32)
|
||||
return torch.stack(frames, dim=0).contiguous()
|
||||
|
||||
if isinstance(video, torch.Tensor):
|
||||
tensor = video.detach()
|
||||
elif isinstance(video, np.ndarray):
|
||||
tensor = torch.from_numpy(video)
|
||||
else:
|
||||
raise TypeError(f"Unsupported video item type: {type(video)!r}")
|
||||
|
||||
if tensor.ndim == 3:
|
||||
return cls._image_to_chw_float_tensor(tensor).unsqueeze(0).contiguous()
|
||||
|
||||
if tensor.ndim != 4:
|
||||
raise ValueError(f"Expected 4D video tensor, got shape={tuple(tensor.shape)}")
|
||||
|
||||
if tensor.shape[-1] in (1, 3, 4) and tensor.shape[1] not in (1, 3, 4):
|
||||
tensor = tensor.permute(0, 3, 1, 2)
|
||||
|
||||
if tensor.shape[1] == 1:
|
||||
tensor = tensor.repeat(1, 3, 1, 1)
|
||||
if tensor.shape[1] == 4:
|
||||
tensor = tensor[:, :3]
|
||||
|
||||
tensor = tensor.to(dtype=torch.float32)
|
||||
if tensor.numel() > 0 and torch.max(tensor) > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
return tensor.contiguous()
|
||||
|
||||
@staticmethod
|
||||
def _audio_to_float_tensor(audio: Any) -> torch.Tensor:
|
||||
if isinstance(audio, tuple) and len(audio) == 2:
|
||||
audio = audio[0]
|
||||
|
||||
if isinstance(audio, torch.Tensor):
|
||||
tensor = audio.detach()
|
||||
elif isinstance(audio, np.ndarray):
|
||||
tensor = torch.from_numpy(audio)
|
||||
else:
|
||||
tensor = torch.as_tensor(audio)
|
||||
|
||||
tensor = tensor.to(dtype=torch.float32).contiguous().view(-1)
|
||||
if tensor.numel() == 0:
|
||||
return torch.zeros((16000,), dtype=torch.float32)
|
||||
|
||||
max_abs = torch.max(torch.abs(tensor))
|
||||
if max_abs > 1.0:
|
||||
tensor = tensor / max_abs
|
||||
|
||||
return tensor.contiguous()
|
||||
|
||||
@classmethod
|
||||
def _convert_modality_item(cls, modality: str, item: Any) -> torch.Tensor:
|
||||
if modality == "image":
|
||||
return cls._image_to_chw_float_tensor(item)
|
||||
if modality == "video":
|
||||
return cls._video_to_tchw_float_tensor(item)
|
||||
if modality == "audio":
|
||||
return cls._audio_to_float_tensor(item)
|
||||
raise ValueError(f"Unsupported modality for Dynin processor: {modality}")
|
||||
|
||||
def _build_modality_kwargs(
|
||||
self,
|
||||
modality: str,
|
||||
modality_items: Sequence[Any],
|
||||
) -> Sequence[Any]:
|
||||
modality = _normalize_modality_name(modality)
|
||||
input_key = _MODALITY_INPUT_KEY_BY_NAME[modality]
|
||||
tensor_items = [self._convert_modality_item(modality, item) for item in modality_items]
|
||||
mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
|
||||
{input_key: tensor_items},
|
||||
{input_key: MultiModalFieldConfig.batched(modality)},
|
||||
)
|
||||
return mm_kwargs[modality]
|
||||
|
||||
def _build_placeholder_ranges(
|
||||
self,
|
||||
*,
|
||||
modality: str,
|
||||
item_count: int,
|
||||
prompt_token_ids: list[int],
|
||||
tokenizer: Any | None,
|
||||
search_start: int,
|
||||
) -> tuple[list[PlaceholderRange], int]:
|
||||
ranges: list[PlaceholderRange] = []
|
||||
|
||||
for _ in range(item_count):
|
||||
placeholder_text = _get_placeholder_text(modality)
|
||||
placeholder_token_ids: list[int] = []
|
||||
|
||||
if placeholder_text and tokenizer is not None:
|
||||
placeholder_token_ids = tokenizer.encode(
|
||||
placeholder_text,
|
||||
add_special_tokens=False,
|
||||
)
|
||||
|
||||
found_offset = None
|
||||
if placeholder_token_ids:
|
||||
found_offset = self._find_subsequence(
|
||||
prompt_token_ids,
|
||||
placeholder_token_ids,
|
||||
search_start,
|
||||
)
|
||||
|
||||
if found_offset is None:
|
||||
found_offset = min(search_start, len(prompt_token_ids) - 1)
|
||||
placeholder_len = 1
|
||||
else:
|
||||
placeholder_len = len(placeholder_token_ids)
|
||||
|
||||
ranges.append(
|
||||
PlaceholderRange(
|
||||
offset=found_offset,
|
||||
length=placeholder_len,
|
||||
is_embed=self._make_disabled_embed_mask(placeholder_len),
|
||||
)
|
||||
)
|
||||
search_start = found_offset + placeholder_len
|
||||
|
||||
return ranges, search_start
|
||||
|
||||
def _get_mm_fields_config(
|
||||
self,
|
||||
hf_inputs: Any,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
) -> Mapping[str, MultiModalFieldConfig]:
|
||||
del hf_inputs, hf_processor_mm_kwargs
|
||||
return {}
|
||||
|
||||
def _get_prompt_updates(
|
||||
self,
|
||||
mm_items: MultiModalDataItems,
|
||||
hf_processor_mm_kwargs: Mapping[str, object],
|
||||
out_mm_kwargs: MultiModalKwargsItems,
|
||||
) -> Sequence[PromptUpdate]:
|
||||
del mm_items, hf_processor_mm_kwargs, out_mm_kwargs
|
||||
return []
|
||||
|
||||
def apply(
|
||||
self,
|
||||
inputs: ProcessorInputs,
|
||||
timing_ctx: TimingContext,
|
||||
) -> MultiModalInputs:
|
||||
prompt = inputs.prompt
|
||||
mm_items = inputs.mm_data_items
|
||||
|
||||
with timing_ctx.record("get_mm_hashes"):
|
||||
mm_hashes = inputs.get_mm_hashes(self.info.model_id)
|
||||
|
||||
tokenizer = self.info.ctx.tokenizer
|
||||
prompt_token_ids = self._encode_prompt_to_token_ids(prompt, tokenizer)
|
||||
prompt_token_ids = self._ensure_non_empty_prompt_ids(prompt_token_ids, tokenizer)
|
||||
|
||||
mm_kwargs_by_modality: dict[str, Sequence[Any]] = {}
|
||||
mm_placeholders: dict[str, list[PlaceholderRange]] = {}
|
||||
search_start = 0
|
||||
mm_counts = mm_items.get_all_counts()
|
||||
|
||||
for modality in _MODALITY_ORDER:
|
||||
item_count = mm_counts.get(modality, 0)
|
||||
if item_count <= 0:
|
||||
continue
|
||||
|
||||
modality_items = mm_items[modality].get_all()
|
||||
if len(modality_items) != item_count:
|
||||
raise RuntimeError(
|
||||
f"Parsed {len(modality_items)} items but expected {item_count} for modality={modality!r}"
|
||||
)
|
||||
|
||||
mm_kwargs_by_modality[modality] = self._build_modality_kwargs(
|
||||
modality,
|
||||
modality_items,
|
||||
)
|
||||
|
||||
placeholder_ranges, search_start = self._build_placeholder_ranges(
|
||||
modality=modality,
|
||||
item_count=item_count,
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
tokenizer=tokenizer,
|
||||
search_start=search_start,
|
||||
)
|
||||
mm_placeholders[modality] = placeholder_ranges
|
||||
|
||||
return MultiModalInputs(
|
||||
type="multimodal",
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
mm_kwargs=MultiModalKwargsItems(mm_kwargs_by_modality),
|
||||
mm_hashes=mm_hashes,
|
||||
mm_placeholders=mm_placeholders,
|
||||
)
|
||||
|
||||
|
||||
class DyninOmniStageBase(nn.Module):
|
||||
stage_name = "Dynin stage"
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self,
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> IntermediateTensors:
|
||||
del batch_size, dtype, device
|
||||
return IntermediateTensors({})
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Any = None,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
del multimodal_embeddings, is_multimodal, kwargs
|
||||
return build_zero_input_embeddings(
|
||||
input_ids=input_ids,
|
||||
hidden_size=self.hidden_size,
|
||||
stage_name=self.stage_name,
|
||||
)
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
return {name for name, _ in weights}
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
sampling_metadata: Any = None,
|
||||
) -> torch.Tensor | None:
|
||||
del hidden_states, sampling_metadata
|
||||
return None
|
||||
|
||||
|
||||
@MULTIMODAL_REGISTRY.register_processor(
|
||||
DyninOmniMultiModalProcessor,
|
||||
info=DyninOmniProcessingInfo,
|
||||
dummy_inputs=DyninOmniDummyInputsBuilder,
|
||||
)
|
||||
class DyninOmniForConditionalGeneration(nn.Module, SupportsMultiModal):
|
||||
supports_multimodal_raw_input_only = True
|
||||
STAGE_ALIAS = {
|
||||
"tokenizer": "token2text",
|
||||
"token2token": "token2text",
|
||||
"detok_text": "token2text",
|
||||
"token2img": "token2image",
|
||||
"token2wav": "token2audio",
|
||||
"token2speech": "token2audio",
|
||||
}
|
||||
|
||||
STAGE_IMPL = {
|
||||
"token2text": (".dynin_omni_token2text", "DyninOmniToken2Text"),
|
||||
"token2image": (".dynin_omni_token2image", "DyninOmniToken2Image"),
|
||||
"token2audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"),
|
||||
}
|
||||
|
||||
_STAGE_IMPL_CACHE: dict[str, type[nn.Module]] = {}
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
|
||||
del i
|
||||
return _get_placeholder_text(modality)
|
||||
|
||||
@classmethod
|
||||
def _resolve_stage_impl_class(cls, model_stage: str) -> type[nn.Module]:
|
||||
impl = cls._STAGE_IMPL_CACHE.get(model_stage)
|
||||
if impl is not None:
|
||||
return impl
|
||||
|
||||
module_name, class_name = cls.STAGE_IMPL[model_stage]
|
||||
module = import_module(module_name, package=__package__)
|
||||
impl = getattr(module, class_name)
|
||||
cls._STAGE_IMPL_CACHE[model_stage] = impl
|
||||
return impl
|
||||
|
||||
@classmethod
|
||||
def _normalize_stage_name(cls, raw_stage: str) -> str:
|
||||
normalized = cls.STAGE_ALIAS.get(raw_stage, raw_stage)
|
||||
if normalized not in cls.STAGE_IMPL:
|
||||
raise ValueError(
|
||||
"Unsupported DYNIN omni model_stage: "
|
||||
f"{raw_stage} (normalized={normalized}). "
|
||||
f"Supported: {sorted(cls.STAGE_IMPL.keys())}"
|
||||
)
|
||||
return normalized
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
raw_stage = str(getattr(vllm_config.model_config, "model_stage", "token2text")).lower()
|
||||
self.model_stage = self._normalize_stage_name(raw_stage)
|
||||
|
||||
impl_cls = self._resolve_stage_impl_class(self.model_stage)
|
||||
self.impl = impl_cls(vllm_config=vllm_config, prefix=prefix)
|
||||
self.model = self.impl
|
||||
|
||||
self.has_preprocess = False
|
||||
self.has_postprocess = False
|
||||
self.have_multimodal_outputs = getattr(self.impl, "have_multimodal_outputs", True)
|
||||
self.requires_raw_input_tokens = getattr(self.impl, "requires_raw_input_tokens", True)
|
||||
self.language_model = self._resolve_language_model()
|
||||
|
||||
def _resolve_language_model(self) -> Any | None:
|
||||
if hasattr(self.impl, "get_language_model"):
|
||||
language_model = self.impl.get_language_model()
|
||||
if language_model is not None:
|
||||
return language_model
|
||||
|
||||
if hasattr(self.impl, "language_model"):
|
||||
language_model = getattr(self.impl, "language_model")
|
||||
if language_model is not None:
|
||||
return language_model
|
||||
|
||||
if self.model_stage == "token2text":
|
||||
return getattr(self.impl, "model", None)
|
||||
|
||||
return None
|
||||
|
||||
def get_language_model(self) -> Any | None:
|
||||
return self.language_model
|
||||
|
||||
@cached_property
|
||||
def sampler(self):
|
||||
if hasattr(self.model, "sampler"):
|
||||
return self.model.sampler
|
||||
if self.language_model is not None and hasattr(self.language_model, "sampler"):
|
||||
return self.language_model.sampler
|
||||
return Sampler()
|
||||
|
||||
def init_multi_modal(self, thinker_config: Any = None) -> None:
|
||||
if hasattr(self.model, "init_multi_modal"):
|
||||
self.model.init_multi_modal(thinker_config)
|
||||
|
||||
def _collect_multimodal_inputs(self, **kwargs: Any) -> dict[str, Any]:
|
||||
mm_inputs: dict[str, Any] = {}
|
||||
for modality, aliases in _MODALITY_INPUT_ALIASES.items():
|
||||
for alias in aliases:
|
||||
if alias in kwargs and kwargs[alias] is not None:
|
||||
mm_inputs[modality] = kwargs[alias]
|
||||
break
|
||||
return mm_inputs
|
||||
|
||||
def _normalize_loaded_weight_names(
|
||||
self,
|
||||
loaded: set[str],
|
||||
expected_param_names: set[str],
|
||||
) -> set[str]:
|
||||
if self.model_stage != "token2text":
|
||||
return loaded
|
||||
|
||||
normalized_loaded: set[str] = set()
|
||||
prefixes = ("", "impl.", "impl.model.")
|
||||
|
||||
for name in loaded:
|
||||
for prefix in prefixes:
|
||||
candidate = f"{prefix}{name}" if prefix else name
|
||||
if candidate in expected_param_names:
|
||||
normalized_loaded.add(candidate)
|
||||
break
|
||||
|
||||
if len(normalized_loaded) < len(expected_param_names):
|
||||
normalized_loaded.update(expected_param_names)
|
||||
|
||||
return normalized_loaded
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
return self.model(
|
||||
input_ids=input_ids,
|
||||
positions=positions,
|
||||
intermediate_tensors=intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def make_empty_intermediate_tensors(
|
||||
self,
|
||||
batch_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> IntermediateTensors:
|
||||
return self.model.make_empty_intermediate_tensors(batch_size, dtype, device)
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
multimodal_embeddings: Any = None,
|
||||
is_multimodal: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
squeezed_batch = False
|
||||
staged_input_ids = input_ids
|
||||
|
||||
if input_ids.ndim == 0:
|
||||
staged_input_ids = input_ids.view(1, 1)
|
||||
squeezed_batch = True
|
||||
elif input_ids.ndim == 1:
|
||||
staged_input_ids = input_ids.unsqueeze(0)
|
||||
squeezed_batch = True
|
||||
|
||||
embeddings = self.model.embed_input_ids(
|
||||
staged_input_ids,
|
||||
multimodal_embeddings=multimodal_embeddings,
|
||||
is_multimodal=is_multimodal,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if squeezed_batch and isinstance(embeddings, torch.Tensor):
|
||||
if embeddings.ndim == 3 and embeddings.shape[0] == 1:
|
||||
return embeddings.squeeze(0)
|
||||
if embeddings.ndim == 2 and input_ids.ndim == 0 and embeddings.shape[0] == 1:
|
||||
return embeddings
|
||||
|
||||
return embeddings
|
||||
|
||||
def embed_multimodal(self, **kwargs: Any) -> Any:
|
||||
if hasattr(self.model, "embed_multimodal"):
|
||||
return self.model.embed_multimodal(**kwargs)
|
||||
|
||||
self._collect_multimodal_inputs(**kwargs)
|
||||
return None
|
||||
|
||||
def load_weights(
|
||||
self,
|
||||
weights: Iterable[tuple[str, torch.Tensor]],
|
||||
) -> set[str]:
|
||||
loaded = self.model.load_weights(weights)
|
||||
if loaded is None:
|
||||
loaded = set()
|
||||
|
||||
expected_param_names = {name for name, _ in self.named_parameters()}
|
||||
if not expected_param_names:
|
||||
return loaded
|
||||
|
||||
return self._normalize_loaded_weight_names(loaded, expected_param_names)
|
||||
|
||||
def compute_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor | OmniOutput,
|
||||
sampling_metadata: Any = None,
|
||||
) -> torch.Tensor | None:
|
||||
return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata)
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
) -> SamplerOutput | None:
|
||||
if hasattr(self.model, "sample"):
|
||||
return self.model.sample(logits, sampling_metadata)
|
||||
if self.language_model is not None and hasattr(self.language_model, "sample"):
|
||||
return self.language_model.sample(logits, sampling_metadata)
|
||||
return None
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,274 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .dynin_omni import DyninOmniStageBase
|
||||
from .dynin_omni_common import (
|
||||
DetokTarget,
|
||||
_looks_like_hf_repo_id,
|
||||
coerce_token_ids_1d,
|
||||
normalize_runtime_info,
|
||||
resolve_dynin_infer_sources,
|
||||
resolve_hidden_size,
|
||||
unwrap_first_value,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _get_hf_token() -> str | None:
|
||||
return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
|
||||
|
||||
|
||||
def _ensure_remote_s2u_vendor_root(
|
||||
*,
|
||||
repo_id: str,
|
||||
local_files_only: bool,
|
||||
) -> str | None:
|
||||
if local_files_only or not _looks_like_hf_repo_id(repo_id):
|
||||
return None
|
||||
|
||||
existing = os.environ.get("DYNIN_S2U_VENDOR_ROOT")
|
||||
if existing:
|
||||
existing_path = Path(existing).expanduser().resolve()
|
||||
if existing_path.is_dir():
|
||||
return str(existing_path)
|
||||
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except Exception as e:
|
||||
logger.warning("huggingface_hub unavailable; cannot fetch s2u_vendor from %s: %s", repo_id, e)
|
||||
return None
|
||||
|
||||
token = _get_hf_token()
|
||||
last_error: Exception | None = None
|
||||
revisions: list[str | None] = [None]
|
||||
|
||||
for revision in revisions:
|
||||
try:
|
||||
snapshot_dir = snapshot_download(
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
allow_patterns=["s2u_vendor/**"],
|
||||
token=token,
|
||||
)
|
||||
except TypeError:
|
||||
try:
|
||||
snapshot_dir = snapshot_download(
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
allow_patterns=["s2u_vendor/**"],
|
||||
)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
continue
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
continue
|
||||
|
||||
vendor_root = (Path(snapshot_dir) / "s2u_vendor").resolve()
|
||||
if vendor_root.is_dir():
|
||||
os.environ["DYNIN_S2U_VENDOR_ROOT"] = str(vendor_root)
|
||||
logger.info("Using remote S2U vendor root: %s", vendor_root)
|
||||
return str(vendor_root)
|
||||
|
||||
if last_error is not None:
|
||||
logger.warning("Failed to download remote s2u_vendor from %s: %s", repo_id, last_error)
|
||||
return None
|
||||
|
||||
|
||||
class DyninOmniToken2Audio(DyninOmniStageBase):
|
||||
"""Stage-3: token detokenization to speech (or pass-through)."""
|
||||
|
||||
stage_name = "Dynin token2audio"
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
del prefix
|
||||
super().__init__()
|
||||
self.vllm_config = vllm_config
|
||||
self.have_multimodal_outputs = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.hidden_size = resolve_hidden_size(vllm_config=vllm_config)
|
||||
self._vq_audio = None
|
||||
self._vq_audio_path: str | None = None
|
||||
self._vq_audio_local_files_only: bool | None = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
del positions, intermediate_tensors, inputs_embeds
|
||||
if input_ids is None:
|
||||
raise ValueError("token2audio stage requires input_ids")
|
||||
|
||||
runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information"))
|
||||
detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0))
|
||||
tokens = coerce_token_ids_1d(input_ids)
|
||||
|
||||
if detok_id != DetokTarget.AUDIO:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"token_ids": tokens,
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device),
|
||||
},
|
||||
)
|
||||
|
||||
audio, sample_rate = self._decode_audio_tokens(tokens, runtime_info=runtime_info)
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"speech": audio,
|
||||
"audio": audio,
|
||||
"sr": torch.tensor([sample_rate], dtype=torch.int, device=audio.device),
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=audio.device),
|
||||
},
|
||||
)
|
||||
|
||||
def _decode_audio_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> tuple[torch.Tensor, int]:
|
||||
# Follow DYNIN validation path:
|
||||
# token list -> "<|speech_x|>" string -> vq_model_audio.decode(...).
|
||||
vq_audio = self._ensure_vq_audio(runtime_info=runtime_info, ref_device=tokens.device)
|
||||
|
||||
audio_codebook_size = int(unwrap_first_value(runtime_info.get("audio_codebook_size"), 4096))
|
||||
audio_vocab_offset = unwrap_first_value(
|
||||
runtime_info.get("audio_vocab_offset"),
|
||||
unwrap_first_value(runtime_info.get("t2s_vocab_start"), None),
|
||||
)
|
||||
|
||||
token_ids = tokens.to(torch.long)
|
||||
if audio_vocab_offset is not None:
|
||||
off = int(audio_vocab_offset)
|
||||
token_ids = torch.where(token_ids >= off, token_ids - off, token_ids)
|
||||
token_ids = token_ids[(token_ids >= 0) & (token_ids < audio_codebook_size)]
|
||||
if token_ids.numel() == 0:
|
||||
raise RuntimeError("Audio detokenizer got no valid audio token ids.")
|
||||
|
||||
speech_unit_str = " ".join(map(str, token_ids.detach().cpu().tolist()))
|
||||
speech_unit_for_decode = "".join(f"<|speech_{unit}|>" for unit in speech_unit_str.split(" ") if unit != "")
|
||||
|
||||
condition = unwrap_first_value(
|
||||
runtime_info.get("condition"),
|
||||
unwrap_first_value(runtime_info.get("t2s_condition"), None),
|
||||
)
|
||||
output_wav_file = unwrap_first_value(runtime_info.get("output_wav_file"), None)
|
||||
created_tmp = False
|
||||
if output_wav_file is None:
|
||||
fd, tmp_wav = tempfile.mkstemp(prefix="dynin_t2s_", suffix=".wav")
|
||||
os.close(fd)
|
||||
output_wav_file = tmp_wav
|
||||
created_tmp = True
|
||||
|
||||
audio_array = vq_audio.decode(speech_unit_for_decode, condition=condition, output_wav_file=output_wav_file)
|
||||
if created_tmp:
|
||||
try:
|
||||
os.remove(output_wav_file)
|
||||
except Exception:
|
||||
pass
|
||||
if not isinstance(audio_array, torch.Tensor):
|
||||
audio_array = torch.as_tensor(audio_array, dtype=torch.float32, device=tokens.device)
|
||||
else:
|
||||
audio_array = audio_array.to(device=tokens.device, dtype=torch.float32)
|
||||
|
||||
if audio_array.ndim > 1:
|
||||
audio_array = audio_array.reshape(-1)
|
||||
audio_array = audio_array.contiguous()
|
||||
|
||||
sample_rate = int(
|
||||
unwrap_first_value(
|
||||
runtime_info.get("sr"),
|
||||
unwrap_first_value(runtime_info.get("sample_rate"), 24000),
|
||||
)
|
||||
)
|
||||
try:
|
||||
cfg = getattr(vq_audio, "u2s_config", None)
|
||||
cfg_sr = getattr(cfg, "sampling_rate", None)
|
||||
if cfg_sr is None:
|
||||
cfg_sr = getattr(getattr(cfg, "data", None), "sampling_rate", None)
|
||||
if cfg_sr is not None:
|
||||
sample_rate = int(cfg_sr)
|
||||
except Exception:
|
||||
pass
|
||||
return audio_array, sample_rate
|
||||
|
||||
def _ensure_vq_audio(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any:
|
||||
sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info)
|
||||
model_path = str(sources.vq_audio_source)
|
||||
local_files_only = bool(sources.vq_audio_local_files_only)
|
||||
|
||||
_ensure_remote_s2u_vendor_root(
|
||||
repo_id=model_path,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
|
||||
if (
|
||||
self._vq_audio is None
|
||||
or self._vq_audio_path != model_path
|
||||
or self._vq_audio_local_files_only != local_files_only
|
||||
):
|
||||
logger.info(
|
||||
"Loading DYNIN audio detokenizer from %s (local_files_only=%s)",
|
||||
model_path,
|
||||
local_files_only,
|
||||
)
|
||||
try:
|
||||
from transformers import AutoModel
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
"transformers is required to load EMOVASpeechTokenizer remote code from Hugging Face."
|
||||
) from e
|
||||
|
||||
try:
|
||||
self._vq_audio = AutoModel.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
local_files_only=local_files_only,
|
||||
low_cpu_mem_usage=False,
|
||||
)
|
||||
except TypeError:
|
||||
try:
|
||||
self._vq_audio = AutoModel.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except TypeError:
|
||||
self._vq_audio = AutoModel.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load EMOVASpeechTokenizer from Hugging Face remote code for model path '{model_path}'."
|
||||
) from e
|
||||
|
||||
if not hasattr(self._vq_audio, "decode"):
|
||||
raise RuntimeError(
|
||||
"Loaded audio tokenizer does not expose decode(). "
|
||||
"Check HF config.json auto_map/model_type and ensure trust_remote_code=True."
|
||||
)
|
||||
self._vq_audio.eval()
|
||||
self._vq_audio.requires_grad_(False)
|
||||
self._vq_audio_path = model_path
|
||||
self._vq_audio_local_files_only = local_files_only
|
||||
if hasattr(self._vq_audio, "to"):
|
||||
self._vq_audio = self._vq_audio.to(ref_device)
|
||||
return self._vq_audio
|
||||
|
||||
def embed_multimodal(self, **kwargs: Any) -> Any:
|
||||
del kwargs
|
||||
return None
|
||||
@@ -0,0 +1,150 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.sequence import IntermediateTensors
|
||||
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
from .dynin_omni import DyninOmniStageBase
|
||||
from .dynin_omni_common import (
|
||||
DetokTarget,
|
||||
_to_bool,
|
||||
coerce_token_ids_1d,
|
||||
get_dynin_magvit_attr,
|
||||
normalize_runtime_info,
|
||||
resolve_dynin_infer_sources,
|
||||
resolve_hidden_size,
|
||||
unwrap_first_value,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class DyninOmniToken2Image(DyninOmniStageBase):
|
||||
"""Stage-2: token detokenization to image (or pass-through)."""
|
||||
|
||||
stage_name = "Dynin token2image"
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
del prefix
|
||||
super().__init__()
|
||||
|
||||
self.vllm_config = vllm_config
|
||||
self.have_multimodal_outputs = True
|
||||
self.requires_raw_input_tokens = True
|
||||
self.hidden_size = resolve_hidden_size(vllm_config=vllm_config)
|
||||
self._vq_model = None
|
||||
self._vq_model_path: str | None = None
|
||||
self._vq_local_files_only: bool | None = None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
positions: torch.Tensor | None = None,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs: Any,
|
||||
) -> OmniOutput:
|
||||
del positions, intermediate_tensors, inputs_embeds
|
||||
if input_ids is None:
|
||||
raise ValueError("token2image stage requires input_ids")
|
||||
runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information"))
|
||||
detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0))
|
||||
tokens = coerce_token_ids_1d(input_ids)
|
||||
|
||||
if detok_id != DetokTarget.IMAGE:
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"token_ids": tokens,
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device),
|
||||
},
|
||||
)
|
||||
|
||||
image = self._decode_image_tokens(tokens, runtime_info=runtime_info)
|
||||
return OmniOutput(
|
||||
text_hidden_states=None,
|
||||
multimodal_outputs={
|
||||
"image": image,
|
||||
"detok_id": torch.tensor([detok_id], dtype=torch.long, device=image.device),
|
||||
},
|
||||
)
|
||||
|
||||
def _decode_image_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> torch.Tensor:
|
||||
# Follow DYNIN validation path:
|
||||
# tokens -> clamp -> vq_model.decode_code -> (x+1)/2 -> [0,1].
|
||||
vq_model = self._ensure_vq_model(runtime_info=runtime_info, ref_device=tokens.device)
|
||||
codebook_size = int(unwrap_first_value(runtime_info.get("codebook_size"), 8192))
|
||||
image_vocab_offset = unwrap_first_value(runtime_info.get("image_vocab_offset"), None)
|
||||
if image_vocab_offset is None:
|
||||
text_vocab_size = unwrap_first_value(runtime_info.get("text_vocab_size"), None)
|
||||
num_new_special_tokens = int(unwrap_first_value(runtime_info.get("num_new_special_tokens"), 0))
|
||||
if text_vocab_size is not None:
|
||||
image_vocab_offset = int(text_vocab_size) + num_new_special_tokens
|
||||
|
||||
token_ids = tokens.to(torch.long)
|
||||
if image_vocab_offset is not None:
|
||||
off = int(image_vocab_offset)
|
||||
token_ids = torch.where(token_ids >= off, token_ids - off, token_ids)
|
||||
token_ids = torch.clamp(token_ids, min=0, max=max(0, codebook_size - 1))
|
||||
token_ids = token_ids.unsqueeze(0)
|
||||
|
||||
decoded = vq_model.decode_code(token_ids)
|
||||
decoded = torch.clamp((decoded + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
if decoded.ndim != 4 or decoded.shape[0] == 0:
|
||||
raise RuntimeError(f"Unexpected MAGVIT decode output shape: {tuple(decoded.shape)}")
|
||||
return decoded[0].contiguous()
|
||||
|
||||
def _ensure_vq_model(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any:
|
||||
sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info)
|
||||
model_path = str(sources.vq_image_source)
|
||||
local_files_only = bool(sources.vq_image_local_files_only)
|
||||
if self._vq_model is None or self._vq_model_path != model_path or self._vq_local_files_only != local_files_only:
|
||||
disable_xet = unwrap_first_value(
|
||||
runtime_info.get("hf_hub_disable_xet"),
|
||||
unwrap_first_value(runtime_info.get("disable_hf_xet"), True),
|
||||
)
|
||||
if _to_bool(disable_xet, default=True):
|
||||
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
||||
logger.info(
|
||||
"Loading DYNIN image detokenizer from %s (local_files_only=%s)",
|
||||
model_path,
|
||||
local_files_only,
|
||||
)
|
||||
try:
|
||||
MAGVITv2 = get_dynin_magvit_attr(
|
||||
"MAGVITv2",
|
||||
source=model_path,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
try:
|
||||
self._vq_model = MAGVITv2.from_pretrained(
|
||||
model_path,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
except TypeError:
|
||||
self._vq_model = MAGVITv2.from_pretrained(model_path)
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
"Failed to load MAGVITv2 from local DYNIN submodel implementation "
|
||||
f"for model path '{model_path}'. "
|
||||
"If your environment cannot access huggingface.co, set "
|
||||
"additional_information.vq_model_image_path to a local MAGVITv2 directory "
|
||||
"and set additional_information.vq_model_image_local_files_only=true."
|
||||
) from e
|
||||
self._vq_model.eval()
|
||||
self._vq_model.requires_grad_(False)
|
||||
self._vq_model_path = model_path
|
||||
self._vq_local_files_only = local_files_only
|
||||
if hasattr(self._vq_model, "to"):
|
||||
self._vq_model = self._vq_model.to(ref_device)
|
||||
return self._vq_model
|
||||
|
||||
def embed_multimodal(self, **kwargs: Any) -> Any:
|
||||
del kwargs
|
||||
return None
|
||||
File diff suppressed because it is too large.
Load diff
@@ -157,6 +157,11 @@ _OMNI_MODELS = {
|
||||
"VoxtralTTSAudioGenerationForConditionalGeneration",
|
||||
),
|
||||
"VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"),
|
||||
"DyninOmniForConditionalGeneration": (
|
||||
"dynin_omni",
|
||||
"dynin_omni",
|
||||
"DyninOmniForConditionalGeneration",
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user