41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""
|
|
Utilities for resolving real models to their tiny model equivalents.
|
|
"""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
|
|
from vllm_omni.diffusion.data import resolve_model_class_name
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def resolve_tiny_model_path(model: str) -> str:
|
|
"""Given a real model name/path, resolve it to a tiny model path.
|
|
|
|
Raises ValueError if the pipeline class cannot be determined (invalid
|
|
model). Returns the original model path if no tiny builder exists yet."""
|
|
pipeline_class = resolve_model_class_name(model)
|
|
if pipeline_class is None:
|
|
raise ValueError(
|
|
f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
|
|
)
|
|
|
|
test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
|
|
if test_opts is None:
|
|
logger.warning(
|
|
"No tiny model builder for pipeline %s (model: %s). Using original model.",
|
|
pipeline_class,
|
|
model,
|
|
)
|
|
return model
|
|
|
|
model_path = test_opts.builder()
|
|
if test_opts.checkpoint_filename is not None:
|
|
return str(Path(model_path) / test_opts.checkpoint_filename)
|
|
return model_path
|