80 lines
3.3 KiB
Python
80 lines
3.3 KiB
Python
import inspect
|
|
from typing import Any
|
|
|
|
from vllm_omni.diffusion.data import TransformerConfig
|
|
|
|
|
|
def get_transformer_config_kwargs(
|
|
tf_model_config: TransformerConfig, model_class: type[Any] | None = None
|
|
) -> dict[str, Any]:
|
|
"""
|
|
This function extracts parameters from a TransformerConfig instance and filters out internal
|
|
diffusers metadata keys (those starting with '_') that should not be passed to model initialization.
|
|
Also filters out parameters that are not accepted by the model's __init__ method (e.g., pooled_projection_dim
|
|
for QwenImageTransformer2DModel).
|
|
|
|
This uses inspect.signature to dynamically detect accepted parameters, making it general for any model class.
|
|
Similar to how diffusers' @register_to_config decorator works.
|
|
|
|
Args:
|
|
tf_model_config: TransformerConfig instance containing model parameters
|
|
model_class: Optional model class to inspect for accepted __init__ parameters.
|
|
If None, all non-internal parameters are returned (backward compatibility).
|
|
|
|
Returns:
|
|
dict: Filtered dictionary of parameters suitable for transformer model initialization
|
|
"""
|
|
# Extract transformer config parameters, filtering out internal diffusers metadata
|
|
# TransformerConfig stores params in a 'params' dict, and we need to exclude
|
|
# internal keys like '_class_name' and '_diffusers_version'
|
|
tf_config_params = tf_model_config.to_dict()
|
|
|
|
# Filter out internal diffusers metadata keys that start with '_'
|
|
filtered_params = {k: v for k, v in tf_config_params.items() if not k.startswith("_")}
|
|
|
|
# If model_class is provided, use inspect.signature to get accepted parameters
|
|
if model_class is not None:
|
|
try:
|
|
# Get the signature of the model's __init__ method
|
|
sig = inspect.signature(model_class.__init__)
|
|
# Get all parameter names (excluding 'self' and special parameters)
|
|
accepted_params = {
|
|
name
|
|
for name, param in sig.parameters.items()
|
|
if name != "self" and param.kind != inspect.Parameter.VAR_KEYWORD # Exclude **kwargs
|
|
}
|
|
|
|
# Filter to only include parameters that are in the model's signature
|
|
filtered_params = {k: v for k, v in filtered_params.items() if k in accepted_params}
|
|
except (TypeError, AttributeError):
|
|
# If inspection fails, fall back to returning all non-internal params
|
|
# This maintains backward compatibility
|
|
pass
|
|
|
|
return filtered_params
|
|
|
|
|
|
def find_module_with_attr(model, attr_name="transformer"):
|
|
"""
|
|
This function searches for a module in the model that has the specified attribute.
|
|
If the model itself has the attribute, it returns the model.
|
|
If none of the modules have the attribute, it returns None.
|
|
"""
|
|
if hasattr(model, attr_name):
|
|
return model
|
|
|
|
for _, child in model.named_children():
|
|
if hasattr(child, attr_name):
|
|
return child
|
|
|
|
return None
|
|
|
|
|
|
def get_transformer_from_pipeline(pipeline: Any, transformer_name: str = "transformer"):
|
|
"""Return a named transformer from a pipeline or pipeline wrapper."""
|
|
pipe = find_module_with_attr(pipeline, attr_name=transformer_name)
|
|
|
|
if pipe is not None:
|
|
return getattr(pipe, transformer_name)
|
|
return None
|