fix_sta

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.

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.

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.

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.

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

assets/prompt_5.txt

@@ -0,0 +1 @@
+fox in the forest close-up quickly turned its head to the left

assets/prompt_6.txt

@@ -0,0 +1 @@
+Man walking his dog in the woods on a hot sunny day

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.

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
 

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)

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]}/'