def optimize_feature(
    sample, flows, occs, correlation_matrix=[], intra_weight=1e2, iters=20, unet_chunk_size=2, optimize_temporal=True
):
    """
    FRESO-guided latent feature optimization
    * optimize spatial correspondence (match correlation_matrix)
    * optimize temporal correspondence (match warped_image)
    """
    if (flows is None or occs is None or (not optimize_temporal)) and (
        intra_weight == 0 or len(correlation_matrix) == 0
    ):
        return sample
    # flows=[fwd_flows, bwd_flows]: (N-1)*2*H1*W1
    # occs=[fwd_occs, bwd_occs]: (N-1)*H1*W1
    # sample: 2N*C*H*W
    torch.cuda.empty_cache()
    video_length = sample.shape[0] // unet_chunk_size
    latent = rearrange(sample.to(torch.float32), "(b f) c h w -> b f c h w", f=video_length)

    cs = torch.nn.Parameter((latent.detach().clone()))
    optimizer = torch.optim.Adam([cs], lr=0.2)

    # unify resolution
    if flows is not None and occs is not None:
        scale = sample.shape[2] * 1.0 / flows[0].shape[2]
        kernel = int(1 / scale)
        bwd_flow_ = F.interpolate(flows[1] * scale, scale_factor=scale, mode="bilinear").repeat(
            unet_chunk_size, 1, 1, 1
        )
        bwd_occ_ = F.max_pool2d(occs[1].unsqueeze(1), kernel_size=kernel).repeat(
            unet_chunk_size, 1, 1, 1
        )  # 2(N-1)*1*H1*W1
        fwd_flow_ = F.interpolate(flows[0] * scale, scale_factor=scale, mode="bilinear").repeat(
            unet_chunk_size, 1, 1, 1
        )
        fwd_occ_ = F.max_pool2d(occs[0].unsqueeze(1), kernel_size=kernel).repeat(
            unet_chunk_size, 1, 1, 1
        )  # 2(N-1)*1*H1*W1
        # match frame 0,1,2,3 and frame 1,2,3,0
        reshuffle_list = list(range(1, video_length)) + [0]

    # attention_probs is the GRAM matrix of the normalized feature
    attention_probs = None
    for tmp in correlation_matrix:
        if sample.shape[2] * sample.shape[3] == tmp.shape[1]:
            attention_probs = tmp  # 2N*HW*HW
            break

    n_iter = [0]
    while n_iter[0] < iters:

        def closure():
            optimizer.zero_grad()

            loss = 0

            # temporal consistency loss
            if optimize_temporal and flows is not None and occs is not None:
                c1 = rearrange(cs[:, :], "b f c h w -> (b f) c h w")
                c2 = rearrange(cs[:, reshuffle_list], "b f c h w -> (b f) c h w")
                warped_image1 = flow_warp(c1, bwd_flow_)
                warped_image2 = flow_warp(c2, fwd_flow_)
                loss = (
                    abs((c2 - warped_image1) * (1 - bwd_occ_)) + abs((c1 - warped_image2) * (1 - fwd_occ_))
                ).mean() * 2

            # spatial consistency loss
            if attention_probs is not None and intra_weight > 0:
                cs_vector = rearrange(cs, "b f c h w -> (b f) (h w) c")
                # attention_scores = torch.bmm(cs_vector, cs_vector.transpose(-1, -2))
                # cs_attention_probs = attention_scores.softmax(dim=-1)
                cs_vector = cs_vector / ((cs_vector**2).sum(dim=2, keepdims=True) ** 0.5)
                cs_attention_probs = torch.bmm(cs_vector, cs_vector.transpose(-1, -2))
                tmp = F.l1_loss(cs_attention_probs, attention_probs) * intra_weight
                loss = tmp + loss

            loss.backward()
            n_iter[0] += 1

            return loss

        optimizer.step(closure)

    torch.cuda.empty_cache()
    return adaptive_instance_normalization(rearrange(cs.data.to(sample.dtype), "b f c h w -> (b f) c h w"), sample)

    FRESO-guided latent feature optimization
    * optimize spatial correspondence (match correlation_matrix)
    * optimize temporal correspondence (match warped_image)
    
examples/community/fresco_v2v.py