Birch-san/slerp.py

from torch import FloatTensor
import torch

device=torch.device('cuda')
dtype=torch.float16

# spherical linear midpoint between [1,0] and [0,1] is [sine(pi/4), sine(pi/4)]
start: FloatTensor = tensor([1,0], dtype=dtype, device=device)
end: FloatTensor = tensor([0,1], dtype=dtype, device=device)
slerp(start, end, 0.5)
# tensor([0.7070, 0.7070], device='cuda:0', dtype=torch.float16)

# many-dimensional data supported; interpolation is performed row-wise
start: FloatTensor = tensor([[1,0], [2,0]], dtype=dtype, device=device)
end: FloatTensor = tensor([[0,1], [0,2]], dtype=dtype, device=device)
slerp(start, end, 0.5)
# tensor([[0.7070, 0.7070],
#         [1.4141, 1.4141]], device='cuda:0', dtype=torch.float16)

# any row where either vector has all zeroes, is computed via fallback to lerp():
start: FloatTensor = tensor([[0,0], [1,0]], dtype=dtype, device=device)
end: FloatTensor = tensor([[1,1], [0,1]], dtype=dtype, device=device)
slerp(start, end, time)
# tensor([[0.5000, 0.5000],
#         [0.7070, 0.7070]], device='cuda:0', dtype=torch.float16)

# any row where both vectors are approx. colinear, is computed via fallback to lerp():
start: FloatTensor = tensor([1,1], dtype=dtype, device=device)
end: FloatTensor = tensor([2,2], dtype=dtype, device=device)
slerp(start, end, 0.5)
# tensor([1.5000, 1.5000], device='cuda:0', dtype=torch.float16)

# we can compute slerp at a batch of timepoints simultaneously,
# by providing a time tensor that is more-dimensional than start/end
start: FloatTensor = tensor([1,0], dtype=dtype, device=device)
end: FloatTensor = tensor([0,1], dtype=dtype, device=device)
time: FloatTensor = tensor([[0.25], [0.5]], dtype=torch.float16, device=device)
slerp(start, end, time)
# tensor([[0.9238, 0.3826],
#         [0.7070, 0.7070]], device='cuda:0', dtype=torch.float16)

import torch
from diffusers import AutoencoderKL
from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
from PIL import Image
from torchvision import transforms
from torchvision.transforms.functional import pil_to_tensor


smiling = <path_img1>
not_smiling = <path_img2>


vae: AutoencoderKL = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae"
)
with torch.no_grad():
    # vae:AutoencoderKL = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix")

    vae.to(dtype=torch.float32)  # otherwise it produces NaNs, even madebyollin's VAE
    vae.to(device="cuda")

    assert vae.device == torch.device("cuda:0")
    assert vae.dtype == torch.float32

    ## 1) smile image
    # make image as tensor
    img = Image.open(smiling).resize((1024, 1024)).convert("RGB")
    img_tensor = pil_to_tensor(img).unsqueeze(0) / 255.0
    img_tensor = img_tensor.to(vae.device)
    img_tensor = img_tensor.to(vae.dtype)

    # get the inferred latent distribution
    latent_dist: DiagonalGaussianDistribution = vae.encode(
        img_tensor, return_dict=False
    )[0]
    print(
        f"{latent_dist.mean.shape=} {latent_dist.std.shape=} {latent_dist.mean.mean()=} {latent_dist.std.mean()=}"
    )

    assert not latent_dist.mean.isnan().any()
    assert not latent_dist.std.isnan().any()
    assert latent_dist.deterministic is False

    latent_smiling = (
        latent_dist.mode()
    )  # .squeeze(0)  # sample from the latent distribution

    ## 2) no smile image
    # make image as tensor
    img = Image.open(not_smiling).resize((1024, 1024)).convert("RGB")
    img_tensor = pil_to_tensor(img).detach().cpu().unsqueeze(0) / 255.0
    img_tensor = img_tensor.to(vae.device)
    img_tensor = img_tensor.to(vae.dtype)

    # get the inferred latent distribution
    latent_dist: DiagonalGaussianDistribution = vae.encode(
        img_tensor, return_dict=False
    )[0]
    print(
        f"{latent_dist.mean.shape=} {latent_dist.std.shape=} {latent_dist.mean.mean()=} {latent_dist.std.mean()=}"
    )

    assert not latent_dist.mean.isnan().any()
    assert not latent_dist.std.isnan().any()
    assert latent_dist.deterministic is False

    latent_not_smiling = (
        latent_dist.mode()
    )  # .squeeze(0)  # sample from the latent distribution

    assert (latent_not_smiling != latent_smiling).any()

    # avg then decode
    t = 0.5
    avg_latent = (
        t * latent_smiling + (1 - t) * latent_not_smiling
    )  # linear interpolation
    # avg_latent = slerp(latent_smiling, latent_not_smiling, t)
    print(
        f"{avg_latent.shape=} {avg_latent.mean()=} {avg_latent.std()=} {avg_latent.min()=} {avg_latent.max()=}"
    )
    image = vae.decode(avg_latent).sample.squeeze(0).cpu().detach()
    image = transforms.ToPILImage()(image)
    image.save("interpolated_.png")