def prepare_mask_coef_by_score(
    video_shape: list,
    cond_frame_idx: list,
    sim_range: list = [0.2, 1.0],
    statistic: list = [1, 100],
    coef_max: int = 0.98,
    score: Optional[torch.Tensor] = None,
):
    """
    the shape of video_data is (b f c h w)
    cond_frame_idx is a list, with length of batch_size
    the shape of statistic  is (f 2)
    the shape of score      is (b f)
    the shape of coef       is (b f)
    """
    assert (
        len(video_shape) == 2
    ), f"the shape of video_shape should be (b f c h w), but now get {len(video_shape.shape)} channels"

    batch_size, frame_num = video_shape[0], video_shape[1]

    score = score.permute(0, 2, 1).squeeze(0)

    # list -> b 1
    cond_fram_mat = torch.tensor(cond_frame_idx).unsqueeze(-1)

    statistic = torch.tensor(statistic)
    # (f 2) -> (b f 2)
    statistic = statistic.repeat(batch_size, 1, 1)

    # shape of order (b f), shape of cond_mat (b f)
    order = torch.arange(0, frame_num, 1)
    order = order.repeat(batch_size, 1)
    cond_mat = torch.ones((batch_size, frame_num)) * cond_fram_mat
    order = abs(order - cond_mat)

    statistic = statistic[:, order.to(torch.long)][0, :, :, :]

    # score (b f) max_s (b f 1)
    max_stats = torch.max(statistic, dim=2).values.to(dtype=score.dtype)
    min_stats = torch.min(statistic, dim=2).values.to(dtype=score.dtype)

    score[score > max_stats] = max_stats[score > max_stats] * 0.95
    score[score < min_stats] = min_stats[score < min_stats]

    eps = 1e-10
    coef = 1 - abs((score / (max_stats + eps)) * (max(sim_range) - min(sim_range)))

    indices = torch.arange(coef.shape[0]).unsqueeze(1)
    coef[indices, cond_fram_mat] = 1.0

    return coef

    the shape of video_data is (b f c h w)
    cond_frame_idx is a list, with length of batch_size
    the shape of statistic  is (f 2)
    the shape of score      is (b f)
    the shape of coef       is (b f)
    
animatediff/utils/util.py