[Tests] Further extend the tiny model testing framework and enable multiple models (#5353)

Signed-off-by: Nick Cao <ncao@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/helpers/tiny_model.py

@@ -4,11 +4,13 @@ Mirrors the component loading loop in DiffusionPipeline.from_pretrained,
 but monkeypatches from_pretrained on ModelMixin and PreTrainedModel to
 initialize with random weights instead of loading checkpoint files.
 Components with vendored configs use those; others load configs from
-the upstream HF model.
+the upstream HF model, optionally passed through a per-component transform
+function first.
 """
 
 import os
 import tempfile
+from collections.abc import Callable
 from pathlib import Path
 from unittest.mock import patch
 
@@ -16,7 +18,7 @@ import torch
 from diffusers import ModelMixin
 from diffusers.pipelines.pipeline_loading_utils import _get_pipeline_class, simple_get_class_obj
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from transformers import AutoConfig, PreTrainedModel
+from transformers import PretrainedConfig, PreTrainedModel
 
 TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
 
@@ -27,39 +29,73 @@ def _get_tiny_model_path(name: str) -> str:
     return path
 
 
-def _diffusers_from_config(cls, pretrained_model_name_or_path, **kwargs):
-    """Replacement for ModelMixin.from_pretrained that initializes random weights."""
-    subfolder = kwargs.get("subfolder")
-    load_kwargs = {"subfolder": subfolder} if subfolder else {}
-    config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
-    return cls.from_config(config)
-
-
-def _transformers_from_config(cls, pretrained_model_name_or_path, **kwargs):
-    """Replacement for PreTrainedModel.from_pretrained that initializes random weights."""
-    subfolder = kwargs.get("subfolder")
-    load_kwargs = {"subfolder": subfolder} if subfolder else {}
-    config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **load_kwargs)
-    return cls(config)
-
-
-def build_tiny_from_configs(pipeline_name: str, model_id: str, configs_dir: str | Path) -> str:
-    """Build a tiny diffusion model from vendored configs with random weights.
+def build_tiny_from_configs(
+    pipeline_name: str,
+    model_id: str,
+    configs_dir: str | Path | None = None,
+    transform: dict[str, Callable[[dict], dict]] | None = None,
+) -> str:
+    """Build a tiny diffusion model with random weights.
 
     Args:
         pipeline_name: Name of the pipeline (used as output directory name).
         model_id: HuggingFace model ID for loading upstream components
-            (tokenizer, scheduler) that don't have vendored configs.
-        configs_dir: Path to the directory containing vendored config files
+            (tokenizer, scheduler, and any component without a vendored
+            config or transform) that don't have vendored configs.
+        configs_dir: Path to a directory containing vendored config files
             (model_index.json and per-component config.json).
+        transform: Maps a component (subfolder) name to a function that
+            takes that component's raw config dict (fetched from model_id)
+            and returns a shrunk one, applied only when that component has
+            no vendored config in configs_dir. Lets each pipeline encode
+            exactly the shrink it needs (layer counts, and any component-
+            internal width that isn't read by another component), including
+            keeping sibling config fields in sync (e.g. layer_types matching
+            num_hidden_layers).
 
     Returns:
         Path to the saved tiny model directory with safetensors weights.
     """
     model_dir = _get_tiny_model_path(pipeline_name)
-    config_dir = Path(configs_dir)
+    config_dir = Path(configs_dir) if configs_dir is not None else None
+    transform = transform or {}
 
-    config_dict = DiffusionPipeline.load_config(config_dir)
+    if config_dir is None and not transform:
+        raise ValueError(
+            f"Neither configs_dir nor transform given for {pipeline_name!r}. Pass a "
+            "configs_dir with vendored configs, and/or transform={{name: fn}} to shrink "
+            f"specific components fetched from the upstream model {model_id!r}."
+        )
+
+    def _diffusers_from_config(cls, pretrained_model_name_or_path, **kwargs):
+        """Replacement for ModelMixin.from_pretrained that initializes random weights."""
+        subfolder = kwargs.get("subfolder")
+        load_kwargs = {"subfolder": subfolder} if subfolder else {}
+        config = cls.load_config(pretrained_model_name_or_path, **load_kwargs)
+        if subfolder in transform:
+            config = transform[subfolder](dict(config))
+        return cls.from_config(config)
+
+    def _transformers_from_config(cls, pretrained_model_name_or_path, **kwargs):
+        """Replacement for PreTrainedModel.from_pretrained that initializes random weights.
+
+        Uses the model's own config_class (every PreTrainedModel subclass
+        declares one) so the raw dict can be transformed before constructing
+        the config object.
+        """
+        subfolder = kwargs.get("subfolder")
+        load_kwargs = {"subfolder": subfolder} if subfolder else {}
+        config_dict, _ = PretrainedConfig.get_config_dict(pretrained_model_name_or_path, **load_kwargs)
+        if subfolder in transform:
+            config_dict = transform[subfolder](config_dict)
+        config = cls.config_class.from_dict(config_dict)
+        return cls(config)
+
+    # model_index.json comes from the vendored configs_dir when it's there, otherwise
+    # from model_id. Independent of the per-component checks below, so a configs_dir
+    # can vendor just one component without also needing its own model_index.json.
+    vendored_index = config_dir is not None and (config_dir / "model_index.json").exists()
+    config_dict = DiffusionPipeline.load_config(config_dir if vendored_index else model_id)
     pipeline_cls = _get_pipeline_class(DiffusionPipeline, config=config_dict)
 
     init_dict, _, _ = pipeline_cls.extract_init_dict(config_dict)
@@ -77,7 +113,7 @@ def build_tiny_from_configs(pipeline_name: str, model_id: str, configs_dir: str
             cls = simple_get_class_obj(library_name, class_name)
 
             # Use vendored config dir if available, otherwise upstream model
-            if (config_dir / name).exists():
+            if config_dir is not None and (config_dir / name).exists():
                 init_kwargs[name] = cls.from_pretrained(config_dir / name)
             else:
                 init_kwargs[name] = cls.from_pretrained(model_id, subfolder=name)