fastvideo.v1.STA_configuration

Module Contents

Functions

average_head_losses	Average losses across all prompts for each mask strategy.
configure_sta	Configure Sliding Tile Attention (STA) parameters based on the specified mode.
read_specific_json_files	Read and parse JSON files containing mask search results.
save_mask_search_results
select_best_mask_strategy	Select the best mask strategy for each head based on loss minimization.

API

fastvideo.v1.STA_configuration.average_head_losses(results: list[dict[str, Any]], selected_masks: list[list[int]]) → dict[str, dict[str, numpy.ndarray]]

Average losses across all prompts for each mask strategy.

fastvideo.v1.STA_configuration.configure_sta(mode: str = 'STA_searching', layer_num: int = 40, time_step_num: int = 50, head_num: int = 40, **kwargs) → list[list[list[Any]]]

Configure Sliding Tile Attention (STA) parameters based on the specified mode.

Parameters:

mode : str The STA mode to use. Options are: - ‘STA_searching’: Generate a set of mask candidates for initial search - ‘STA_tuning’: Select best mask strategy based on previously saved results - ‘STA_inference’: Load and use a previously tuned mask strategy layer_num: int, number of layers time_step_num: int, number of timesteps head_num: int, number of heads

**kwargs : dict Mode-specific parameters:

For 'STA_searching':
- mask_candidates: list of str, optional, mask candidates to use
- mask_selected: list of int, optional, indices of selected masks

For 'STA_tuning':
- mask_search_files_path: str, required, path to mask search results
- mask_candidates: list of str, optional, mask candidates to use
- mask_selected: list of int, optional, indices of selected masks
- skip_time_steps: int, optional, number of time steps to use full attention (default 12)
- save_dir: str, optional, directory to save mask strategy (default "mask_candidates")

For 'STA_inference':
- load_path: str, optional, path to load mask strategy (default "mask_candidates/mask_strategy.json")

fastvideo.v1.STA_configuration.read_specific_json_files(folder_path: str) → list[dict[str, Any]]

Read and parse JSON files containing mask search results.

fastvideo.v1.STA_configuration.save_mask_search_results(mask_search_final_result: list[dict[str, list[float]]], prompt: str, mask_strategies: list[str], output_dir: str = 'output/mask_search_result/') → str | None

fastvideo.v1.STA_configuration.select_best_mask_strategy(averaged_results: dict[str, dict[str, numpy.ndarray]], selected_masks: list[list[int]], skip_time_steps: int = 12, timesteps: int = 50, head_num: int = 40) → tuple[dict[str, list[int]], float, dict[str, int]]

Select the best mask strategy for each head based on loss minimization.