25 lines
748 B
Python
25 lines
748 B
Python
import mlx.core as mx
|
|
from mlx import nn
|
|
|
|
from mflux.models.vae.decoder.decoder import Decoder
|
|
from mflux.models.vae.encoder.encoder import Encoder
|
|
|
|
|
|
class VAE(nn.Module):
|
|
scaling_factor: int = 0.3611
|
|
shift_factor: int = 0.1159
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.decoder = Decoder()
|
|
self.encoder = Encoder()
|
|
|
|
def decode(self, latents: mx.array) -> mx.array:
|
|
scaled_latents = (latents / self.scaling_factor) + self.shift_factor
|
|
return self.decoder.decode(scaled_latents)
|
|
|
|
def encode(self, latents: mx.array) -> mx.array:
|
|
latents = self.encoder.encode(latents)
|
|
mean, _ = mx.split(latents, 2, axis=1)
|
|
return (mean - self.shift_factor) * self.scaling_factor
|