MiniMaxAI

/

MiniMax-H3

like 4.02k

Follow

MiniMax 9.42k

Thanks for opening source such great model ! When testing the video VAE, I found that the result obtained by encoding and then decoding the image showed color differences compared to the original image. Here is my testing code

import numpy as np
from PIL import Image
import torch
from diffusers import ModularPipeline, ComponentsManager

if __name__ == '__main__':
    manager = ComponentsManager()
    manager.enable_auto_cpu_offload(device="cuda")

    pipe = ModularPipeline.from_pretrained("/mnt/data/0/all_users/llx/minimax_h3/MiniMax-H3", workflow="t2va", components_manager=manager)
    pipe.load_components(dtype=torch.bfloat16)

    with torch.no_grad():
        pixel_mean = torch.tensor([0.485, 0.456, 0.406]).reshape(1,3,1,1,1).cuda()
        pixel_std = torch.tensor([0.229, 0.224, 0.225]).reshape(1,3,1,1,1).cuda()
        img = np.array(Image.open('test.png'))
        img_tensor = (torch.from_numpy(img/255.).float().permute(2,0,1).unsqueeze(0).unsqueeze(2).cuda() - pixel_mean) / pixel_std
        img_latent = pipe.vae.encode(img_tensor).latent_dist.sample()
        recon_img = pipe.vae._decode_clip(img_latent)
        recon_img = ((recon_img * pixel_std + pixel_mean).clamp(0,1).squeeze().permute(1,2,3,0).cpu().numpy() * 255).astype(np.uint8)
        Image.fromarray(recon_img[-1:]).save('recon.png')

The input image and the reconstructed image as follows