Function drag_gan

draggan/deprecated/api.py:103–155

(
    g_ema,
    latent: torch.Tensor,
    noise,
    F,
    handle_points,
    target_points,
    mask,
    max_iters=1000,
    r1=3,
    r2=12,
    lam=20,
    d=2,
    lr=2e-3,
)

def drag_gan(
    g_ema,
    latent: torch.Tensor,
    noise,
    F,
    handle_points,
    target_points,
    mask,
    max_iters=1000,
    r1=3,
    r2=12,
    lam=20,
    d=2,
    lr=2e-3,
):
    handle_points0 = copy.deepcopy(handle_points)
    handle_points = torch.stack(handle_points)
    handle_points0 = torch.stack(handle_points0)
    target_points = torch.stack(target_points)

    F0 = F.detach().clone()
    device = latent.device

    latent_trainable = latent[:, :6, :].detach().clone().requires_grad_(True)
    latent_untrainable = latent[:, 6:, :].detach().clone().requires_grad_(False)
    optimizer = torch.optim.Adam([latent_trainable], lr=lr)
    for _ in range(max_iters):
        if torch.allclose(handle_points, target_points, atol=d):
            break

        optimizer.zero_grad()
        latent = torch.cat([latent_trainable, latent_untrainable], dim=1)
        sample2, F2 = g_ema.generate(latent, noise)

        # motion supervision
        loss = motion_supervison(handle_points, target_points, F2, r1, device)

        if mask is not None:
            loss += ((F2 - F0) * (1 - mask)).abs().mean() * lam

        loss.backward()
        optimizer.step()

        with torch.no_grad():
            latent = torch.cat([latent_trainable, latent_untrainable], dim=1)
            sample2, F2 = g_ema.generate(latent, noise)
            handle_points = point_tracking(F2, F0, handle_points, handle_points0, r2, device)

        F = F2.detach().clone()
        # if iter % 1 == 0:
        #     print(iter, loss.item(), handle_points, target_points)

        yield sample2, latent, F2, handle_points


def motion_supervison(handle_points, target_points, F2, r1, device):