import torch

def vae(data_batch, pipe, key):

    # Video
    raw_state = torch.stack([val for key, val in data_batch['data']['rgb'].items() if key[1] == 0], 2).bfloat16()
    data_batch[key] = raw_state = 2.0 * raw_state - 1.0
    latent_state = pipe.encode(raw_state).contiguous()
    data_batch[f'{key}_latent'] = latent_state

    # Text
    t5_text_embedding = pipe.encode_prompt(data_batch['prompt']).to(dtype=pipe.torch_dtype)
    data_batch['t5_text_embeddings'] = t5_text_embedding

    # Return updated batch
    return data_batch