Skip to content

vllm_omni.diffusion.cache.teacache.coefficient_estimator

BagelAdapter

Bases: DefaultAdapter

Adapter for Bagel model.

model_class_name class-attribute instance-attribute 

model_class_name = 'BagelPipeline'

uses_tf_config class-attribute instance-attribute 

uses_tf_config = False

get_transformer staticmethod 

get_transformer(pipeline: Any) -> tuple[Any, str]

install_hook staticmethod 

install_hook(
    transformer: Any, hook: DataCollectionHook
) -> None

DataCollectionHook

Bases: ModelHook

Hook to collect modulated inputs and model outputs for TeaCache coefficient estimation.

current_trajectory instance-attribute 

current_trajectory: list[tuple[ndarray, ndarray]] = []

extractor_fn instance-attribute 

extractor_fn = None

transformer_type instance-attribute 

transformer_type = transformer_type

initialize_hook 

initialize_hook(module: Module) -> Module

new_forward 

new_forward(
    module: Module, *args: Any, **kwargs: Any
) -> Any

start_collection 

start_collection()

stop_collection 

stop_collection() -> list[tuple[ndarray, ndarray]]

DefaultAdapter

Default adapter for standard diffusers pipelines.

model_class_name class-attribute instance-attribute 

model_class_name = None

uses_tf_config class-attribute instance-attribute 

uses_tf_config = True

get_transformer staticmethod 

get_transformer(pipeline: Any) -> tuple[Any, str]

install_hook staticmethod 

install_hook(
    transformer: Any, hook: DataCollectionHook
) -> None

load_pipeline classmethod 

load_pipeline(
    model_path: str, device: str, dtype: dtype
) -> Any

Flux2Adapter

Bases: DefaultAdapter

Adapter for Flux2 model coefficient estimation.

model_class_name class-attribute instance-attribute 

model_class_name = 'Flux2Pipeline'

LongCatAdapter

Bases: DefaultAdapter

Adapter for LongCat Image - NOTE: currently this model needs the vLLM context to be correctly configured to actually run the estimation, since it uses vLLM norm layers etc.

model_class_name class-attribute instance-attribute 

model_class_name = 'LongCatImagePipeline'

StableAudioAdapter

Bases: DefaultAdapter

Adapter for Stable Audio Open 1.0 coefficient estimation.

model_class_name class-attribute instance-attribute 

model_class_name = 'StableAudioPipeline'

TeaCacheCoefficientEstimator

Model-agnostic helper class to collect data and estimate TeaCache coefficients.

collected_data instance-attribute 

collected_data: list[list[tuple[ndarray, ndarray]]] = []

hook instance-attribute 

hook = DataCollectionHook(transformer_type)

pipeline instance-attribute 

pipeline = load_pipeline(model_path, device, dtype)

collect_from_prompt 

collect_from_prompt(prompt: str, **generate_kwargs)

estimate 

estimate(poly_order: int = 4) -> list[float]

Estimate polynomial coefficients from collected data.

Parameters:

Name Type Description Default
poly_order int

Order of polynomial fit (default: 4)

 4

Returns:

Type Description
list[float]

List of polynomial coefficients [a_n, a_{n-1}, ..., a_1, a_0]

Raises:

Type Description
RuntimeError

If no data has been collected

calculate_relative_l1 

calculate_relative_l1(
    tensor_current: ndarray, tensor_next: ndarray
) -> float

Calculate relative L1 distance (Eq. 4 from TeaCache paper).

estimate_teacache_coefficients 

estimate_teacache_coefficients(
    collected_data: list[list[tuple[ndarray, ndarray]]],
    poly_order: int = 4,
) -> list[float]

Estimate polynomial coefficients for TeaCache using np.polyfit.