import torch
from diffusers import AutoencoderKL

class Flux1Autoencoder(AutoencoderKL):
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
        kwargs['torch_dtype']=torch.float32
        return super().from_pretrained(pretrained_model_name_or_path, **kwargs)