Spaces:
Runtime error
Runtime error
File size: 645 Bytes
0c8d55e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from diffusers.models.autoencoders import autoencoder_kl, vae
import torch
class CompiledAutoencoderKL(autoencoder_kl.AutoencoderKL):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@torch.compile
def encode(self, *args, **kwargs):
return super().encode(*args, **kwargs)
autoencoder_kl.AutoencoderKL = CompiledAutoencoderKL
class CompiledEncoder(vae.Encoder):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@torch.compile
def forward(self, *args, **kwargs):
return super().forward(*args, **kwargs)
vae.Encoder = CompiledEncoder |