Remote VAE allow unpacked latents

[vladmandic]

Did you like the remote VAE solution?

Did I like remote VAE solution? Absolutely! Everything about it, its a great service from HF.

What can be improved about the current solution?

Currently FLUX.1 endpoint accepts packed latents and performs unpack itself.
Ask is to allow unpacked latents as well - this would basically be "skip unpacking if already unpacked".
Reason? HiDream uses same VAE as FLUX.1, but it works with unpacked latents and user should not need to pack them just to be able to send them.
This should not really require any additional params as determination if latent is packed or unpacked can easily be inferred based on latent shape.

What other VAEs you would like to see if the pilot goes well?

N/A

Notify the members of the team

@hlky @sayakpaul

[hlky]

Hi @vladmandic. Unpacked latents can be used with Flux endpoint already.

import torch
from diffusers.utils.remote_utils import remote_decode
from diffusers.utils.constants import DECODE_ENDPOINT_FLUX

image = remote_decode(
    endpoint=DECODE_ENDPOINT_FLUX,
    tensor=torch.randn([1, 16, 128, 128], dtype=torch.float16),
    scaling_factor=0.3611,
    shift_factor=0.1159,
)

If v1 or XL latents happened to be packed for some reason, that works too.

import torch
from diffusers.utils.remote_utils import remote_decode
from diffusers.utils.constants import DECODE_ENDPOINT_SD_V1


def _pack_latents(
    latents: torch.Tensor, batch_size, num_channels_latents, height, width
):
    latents = latents.view(
        batch_size, num_channels_latents, height // 2, 2, width // 2, 2
    )
    latents = latents.permute(0, 2, 4, 1, 3, 5)
    latents = latents.reshape(
        batch_size, (height // 2) * (width // 2), num_channels_latents * 4
    )

    return latents


packed = _pack_latents(
    torch.randn([1, 4, 64, 64], dtype=torch.float16),
    batch_size=1,
    num_channels_latents=4,
    height=64,
    width=64,
)

image = remote_decode(
    endpoint=DECODE_ENDPOINT_SD_V1,
    tensor=packed,
    scaling_factor=0.18215,
    height=512,
    width=512,
)

[vladmandic]

ah, my bad, error message i was getting was pointing me in a wrong direction:
(i use url endpoints directly, don't rely on diffusers.utils.remote_utils)

url="https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud"
args={'input_tensor_type': 'binary', 'shape': [16, 128, 128], 'dtype': 'bfloat16', 'partial_postprocess': False, 'output_type': 'pt', 'output_tensor_type': 'binary', 'width': 512, 'height': 512, 'scaling_factor': 0.3611, 'shift_factor': 0.1159}
response={'detail': "shape '[16, 32, 32, 32, 2, 2]' is invalid for input of size 262144"}

what was missing is batch size at the front, input tensor for h1 should be [1,16,128,128]
remote vae for hidream using f1 endpoint works!