diff --git a/assets/prompt_0.txt b/assets/prompt_0.txt new file mode 100644 index 000000000..ad0d62213 --- /dev/null +++ b/assets/prompt_0.txt @@ -0,0 +1 @@ +Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting. diff --git a/assets/prompt_1.txt b/assets/prompt_1.txt new file mode 100644 index 000000000..c56a97618 --- /dev/null +++ b/assets/prompt_1.txt @@ -0,0 +1 @@ +A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature. diff --git a/assets/prompt_2.txt b/assets/prompt_2.txt new file mode 100644 index 000000000..5b4663776 --- /dev/null +++ b/assets/prompt_2.txt @@ -0,0 +1 @@ +A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere. diff --git a/assets/prompt_3.txt b/assets/prompt_3.txt new file mode 100644 index 000000000..f317af2f2 --- /dev/null +++ b/assets/prompt_3.txt @@ -0,0 +1 @@ +A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones. diff --git a/assets/prompt_4.txt b/assets/prompt_4.txt new file mode 100644 index 000000000..26edb8b65 --- /dev/null +++ b/assets/prompt_4.txt @@ -0,0 +1 @@ +A superintelligent humanoid robot waking up. The robot has a sleek metallic body with futuristic design features. Its glowing red eyes are the focal point, emanating a sharp, intense light as it powers on. The scene is set in a dimly lit, high-tech laboratory filled with glowing control panels, robotic arms, and holographic screens. The setting emphasizes advanced technology and an atmosphere of mystery. The ambiance is eerie and dramatic, highlighting the moment of awakening and the robots immense intelligence. Photorealistic style with a cinematic, dark sci-fi aesthetic. Aspect ratio: 16:9 --v 6.1 diff --git a/assets/prompt_5.txt b/assets/prompt_5.txt new file mode 100644 index 000000000..67d4f90d8 --- /dev/null +++ b/assets/prompt_5.txt @@ -0,0 +1 @@ +fox in the forest close-up quickly turned its head to the left diff --git a/assets/prompt_6.txt b/assets/prompt_6.txt new file mode 100644 index 000000000..86751da5d --- /dev/null +++ b/assets/prompt_6.txt @@ -0,0 +1 @@ +Man walking his dog in the woods on a hot sunny day diff --git a/assets/prompt_7.txt b/assets/prompt_7.txt new file mode 100644 index 000000000..11e03e476 --- /dev/null +++ b/assets/prompt_7.txt @@ -0,0 +1 @@ +A majestic lion strides across the golden savanna, its powerful frame glistening under the warm afternoon sun. The tall grass ripples gently in the breeze, enhancing the lion's commanding presence. The tone is vibrant, embodying the raw energy of the wild. Low angle, steady tracking shot, cinematic. diff --git a/examples/inference/sta_mask_search/inference_wan_sta.sh b/examples/inference/sta_mask_search/inference_wan_sta.sh index ba448f5cc..4bd5d87ca 100644 --- a/examples/inference/sta_mask_search/inference_wan_sta.sh +++ b/examples/inference/sta_mask_search/inference_wan_sta.sh @@ -5,7 +5,7 @@ export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers base_port=29503 -num_gpu=1 +num_gpu=8 gpu_ids=$(seq 0 $((num_gpu-1))) skip_time_steps=12 diff --git a/fastvideo/attention/backends/STA_configuration.py b/fastvideo/attention/backends/STA_configuration.py index 9111dcab3..d97e48791 100644 --- a/fastvideo/attention/backends/STA_configuration.py +++ b/fastvideo/attention/backends/STA_configuration.py @@ -350,7 +350,7 @@ def select_best_mask_strategy( def save_mask_search_results( - mask_search_final_result: list[dict[str, list[float]]], + mask_search_final_result: list[Any], prompt: str, mask_strategies: list[str], output_dir: str = 'output/mask_search_result/') -> str | None: @@ -358,8 +358,9 @@ def save_mask_search_results( print("No mask search results to save") return None - # Create result dictionary with defaultdict for nested lists - mask_search_dict: dict[str, dict[str, list[list[float]]]] = { + # Create result dictionary with nested lists: + # [timesteps][layers][heads]. + mask_search_dict: dict[str, dict[str, list[list[list[float]]]]] = { "L2_loss": defaultdict(list), "L1_loss": defaultdict(list) } @@ -371,23 +372,79 @@ def save_mask_search_results( masks_list = [int(x) for x in mask.split(',')] selected_masks.append(masks_list) + def _to_float_list(loss_values: Any, loss_name: str) -> list[float]: + if isinstance(loss_values, np.ndarray): + loss_values = loss_values.tolist() + if not isinstance(loss_values, list | tuple): + raise ValueError( + f"{loss_name} must be a sequence of numeric values") + + float_values = [] + for loss in loss_values: + try: + float_values.append(float(loss)) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Invalid {loss_name} value {loss!r}: expected a number" + ) from exc + return float_values + + def _extract_timestep_layer_losses(step_data: Any, loss_name: str, + strategy_idx: int) -> list[list[float]]: + if isinstance(step_data, dict): + layer_data_list = [step_data] + elif isinstance(step_data, list): + layer_data_list = step_data + else: + return [] + + timestep_layer_losses: list[list[float]] = [] + for layer_data in layer_data_list: + if not isinstance(layer_data, dict) or loss_name not in layer_data: + continue + + raw_losses = layer_data[loss_name] + if isinstance(raw_losses, np.ndarray): + raw_losses = raw_losses.tolist() + if not isinstance(raw_losses, list | tuple): + raise ValueError(f"{loss_name} must be a list or tuple") + raw_losses = list(raw_losses) + + if raw_losses and isinstance(raw_losses[0], list | tuple + | np.ndarray): + if strategy_idx >= len(raw_losses): + raise ValueError( + f"Missing strategy index {strategy_idx} in {loss_name}") + strategy_losses = raw_losses[strategy_idx] + else: + if strategy_idx > 0: + continue + strategy_losses = raw_losses + + timestep_layer_losses.append( + _to_float_list(strategy_losses, loss_name)) + + return timestep_layer_losses + # Process each mask strategy for i, mask_strategy in enumerate(selected_masks): mask_strategy_str = str(mask_strategy) - # Process L2 loss - step_results: list[list[float]] = [] - for step_data in mask_search_final_result: - if isinstance(step_data, dict) and "L2_loss" in step_data: - layer_losses = [float(loss) for loss in step_data["L2_loss"]] - step_results.append(layer_losses) - mask_search_dict["L2_loss"][mask_strategy_str] = step_results + l2_step_results: list[list[list[float]]] = [] + l1_step_results: list[list[list[float]]] = [] - step_results = [] for step_data in mask_search_final_result: - if isinstance(step_data, dict) and "L1_loss" in step_data: - layer_losses = [float(loss) for loss in step_data["L1_loss"]] - step_results.append(layer_losses) - mask_search_dict["L1_loss"][mask_strategy_str] = step_results + l2_layer_losses = _extract_timestep_layer_losses( + step_data, "L2_loss", i) + if l2_layer_losses: + l2_step_results.append(l2_layer_losses) + + l1_layer_losses = _extract_timestep_layer_losses( + step_data, "L1_loss", i) + if l1_layer_losses: + l1_step_results.append(l1_layer_losses) + + mask_search_dict["L2_loss"][mask_strategy_str] = l2_step_results + mask_search_dict["L1_loss"][mask_strategy_str] = l1_step_results # Create the output directory if it doesn't exist os.makedirs(output_dir, exist_ok=True) diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 0630fdc5b..e67b985a2 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -722,20 +722,14 @@ class DenoisingStage(PipelineStage): from fastvideo.attention.backends.STA_configuration import save_mask_search_results if batch.mask_search_final_result_pos is not None and batch.prompt is not None: save_mask_search_results( - [ - dict(layer_data) - for layer_data in batch.mask_search_final_result_pos - ], + batch.mask_search_final_result_pos, prompt=str(batch.prompt), mask_strategies=sparse_mask_candidates_searching, output_dir=f'output/mask_search_result_pos_{size[0]}x{size[1]}/' ) if batch.mask_search_final_result_neg is not None and batch.prompt is not None: save_mask_search_results( - [ - dict(layer_data) - for layer_data in batch.mask_search_final_result_neg - ], + batch.mask_search_final_result_neg, prompt=str(batch.prompt), mask_strategies=sparse_mask_candidates_searching, output_dir=f'output/mask_search_result_neg_{size[0]}x{size[1]}/'