diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py index c7d254eb9..cedd11068 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py @@ -1,12 +1,8 @@ from collections.abc import Callable from enum import Enum, auto -from typing import ( - Any, - Literal, - NotRequired, - TypeAlias, - TypedDict, -) +from typing import Any, Literal, TypeAlias + +from typing_extensions import NotRequired, TypedDict AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"] diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/comfyui/test_comfyui_integration.py index f532c350f..e15098b4e 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/tests/comfyui/test_comfyui_integration.py @@ -10,7 +10,7 @@ import multiprocessing import time import traceback from collections.abc import Iterable, Sequence -from enum import StrEnum, auto +from enum import Enum from types import SimpleNamespace from typing import Any, NamedTuple @@ -57,15 +57,15 @@ class SamplingCase(NamedTuple): lora: dict | None = None -class SamplingKind(StrEnum): - IMAGE_NONE = auto() - IMAGE_DIFFUSION_SINGLE = auto() - UNDERSTANDING_NONE = auto() - UNDERSTANDING_AR_LIST = auto() - TTS_NONE = auto() - TTS_DIFFUSION_SINGLE = auto() - VIDEO_NONE = auto() - VIDEO_DIFFUSION_SINGLE = auto() +class SamplingKind(str, Enum): + IMAGE_NONE = "image_none" + IMAGE_DIFFUSION_SINGLE = "image_diffusion_single" + UNDERSTANDING_NONE = "understanding_none" + UNDERSTANDING_AR_LIST = "understanding_ar_list" + TTS_NONE = "tts_none" + TTS_DIFFUSION_SINGLE = "tts_diffusion_single" + VIDEO_NONE = "video_none" + VIDEO_DIFFUSION_SINGLE = "video_diffusion_single" # Pre-defined arguments to be used in function calls during the tests diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 63ec9f858..5ba3fea68 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -16,20 +16,10 @@ Key invariants tested: - voice/speaker parameter compatibility in chat completions """ -import enum import json from unittest.mock import MagicMock import pytest - -# Python 3.10 compat: StrEnum was added in 3.11 -if not hasattr(enum, "StrEnum"): - - class _StrEnum(str, enum.Enum): - """Minimal StrEnum backport for Python 3.10.""" - - enum.StrEnum = _StrEnum # type: ignore[attr-defined] - from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponseStreamChoice, diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py index da599c2c2..49845fa65 100644 --- a/tests/model_tests/diffusion/config_types.py +++ b/tests/model_tests/diffusion/config_types.py @@ -3,7 +3,7 @@ Common definitions for controlling what tests run where. """ from collections.abc import Callable -from enum import StrEnum, auto +from enum import Enum from typing import NamedTuple, TypeAlias from pytest import MarkDecorator @@ -15,26 +15,26 @@ from vllm_omni.entrypoints.omni import Omni TinyDiffusionBuilder: TypeAlias = Callable[[], str] -class DiffusionAccs(StrEnum): +class DiffusionAccs(str, Enum): """Supported acceleration types / test settings for Diffusion Models.""" - HSDP = auto() - TEA_CACHE = auto() - CACHE_DIT = auto() - SEQUENCE_PARALLEL = auto() - CFG_PARALLEL = auto() - TENSOR_PARALLEL = auto() - CPU_OFFLOAD = auto() - LAYERWISE_OFFLOAD = auto() - VAE_PATCH_PARALLEL = auto() + HSDP = "hsdp" + TEA_CACHE = "tea_cache" + CACHE_DIT = "cache_dit" + SEQUENCE_PARALLEL = "sequence_parallel" + CFG_PARALLEL = "cfg_parallel" + TENSOR_PARALLEL = "tensor_parallel" + CPU_OFFLOAD = "cpu_offload" + LAYERWISE_OFFLOAD = "layerwise_offload" + VAE_PATCH_PARALLEL = "vae_patch_parallel" -class DiffusionTasks(StrEnum): +class DiffusionTasks(str, Enum): """Supported tasks for Diffusion Models.""" - TEXT_TO_IMAGE = auto() - IMAGE_TO_IMAGE = auto() - TEXT_TO_VIDEO = auto() + TEXT_TO_IMAGE = "text_to_image" + IMAGE_TO_IMAGE = "image_to_image" + TEXT_TO_VIDEO = "text_to_video" # Text to audio, etc should be added here as needed diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 61d3d12a3..f0344c994 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -1,24 +1,16 @@ import copy import pprint from dataclasses import asdict, dataclass, field -from typing import Any, TypeAlias, TypedDict +from typing import Any, TypeAlias -from vllm.inputs import PromptType +import torch +from typing_extensions import NotRequired, TypedDict +from vllm.inputs import EmbedsPrompt, PromptType, TextPrompt, TokensPrompt +from vllm.inputs.engine import TokensInput from vllm.sampling_params import SamplingParams from vllm_omni.lora.request import LoRARequest -try: - from typing import NotRequired -except ImportError: - # Python < 3.11: use typing_extensions - from typing_extensions import NotRequired - - -import torch -from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt -from vllm.inputs.engine import TokensInput - class OmniTextPrompt(TextPrompt): """Text prompt with optional embeddings and stage payloads. diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py index 743669454..7cdbc1f19 100644 --- a/vllm_omni/outputs/output_modality.py +++ b/vllm_omni/outputs/output_modality.py @@ -8,23 +8,8 @@ for type-safe multimodal output routing and tensor merging. from __future__ import annotations import re -import sys from enum import Enum, Flag, auto -if sys.version_info >= (3, 11): - from enum import StrEnum -else: - # TODO: Remove this Python 3.10 compatibility shim when Python 3.10 support is dropped. - class StrEnum(str, Enum): - """``enum.StrEnum`` for Python 3.10, which predates it. - - ``__str__`` is taken from ``str`` so ``str(member)`` yields the member - value rather than ``"ClassName.MEMBER"``, matching 3.11 semantics. - """ - - __str__ = str.__str__ - - _MODALITY_ALIASES: dict[str, str] = { "speech": "audio", "images": "image", @@ -38,13 +23,17 @@ _MODALITY_ALIASES: dict[str, str] = { } -class OutputModalityNames(StrEnum): +class OutputModalityNames(str, Enum): """Keys for output modalities. TODO: (Alex) Integrate this with the big-flag enum below + throughout the code for better type safety (currently only used for output processor). """ + # Keep the string conversion provided by ``enum.StrEnum`` while using the + # Python 3.10-compatible ``str, Enum`` form. + __str__ = str.__str__ + TEXT = "text" IMAGE = "image" AUDIO = "audio"