Requesting "stats/vae_trial1.pkl" checkpoint

[snoop2head]

Dear Willis,

Thank you for your awesome work!

If possible, can you share "stats/vae_trial1.pkl" checkpoint at

diffuse_nnx/networks/encoders/sd_vae.py

Lines 55 to 58 in 023afd2

	def initialize(self):
	ckpt_path = os.path.join(Path(__file__).parent, self.pretrained_path)
	if not os.path.exists(ckpt_path):
	utils.download_blob('will-data', 'stats/vae_trial1.pkl', ckpt_path)

Thank you so much!

Best regards,

Young Jin

[willisma]

Hi Young,

Thanks! We will release all checkpoints soon!

[snoop2head]

Dear Willis,

Thank you for your reply! Nice work on SiT by the way. I've learned a lot, and is current replicating it.

Would replacing sd-vae with FlaxAutoencoderKL and using pretrained checkpoint do the job?

I implemented this way referring https://github.com/kvfrans/jax-flow :

"""Stable Diffusion VAE wrapper for deterministic latent pipelines."""

from __future__ import annotations

from functools import partial
from typing import Optional

import jax
import jax.numpy as jnp
from chex import Array, PRNGKey
from diffusers import FlaxAutoencoderKL

class StableVAE:
    """Thin wrapper around `FlaxAutoencoderKL` with deterministic encode helpers."""

    def __init__(self, checkpoint_id: str = "pcuenq/sd-vae-ft-mse-flax"):
        module, params = FlaxAutoencoderKL.from_pretrained(checkpoint_id)
        self.module = module
        self.params = jax.device_get(params)
        self.scaling_factor = float(module.config.scaling_factor)

    @staticmethod
    def _to_nchw(x: Array) -> Array:
        return jnp.transpose(x, (0, 3, 1, 2))

    @staticmethod
    def _to_nhwc(x: Array) -> Array:
        return jnp.transpose(x, (0, 2, 3, 1))

    # Keep `self` and the deterministic flag static for compilation caching.
    @partial(jax.jit, static_argnums=(0,), static_argnames=("deterministic",))
    def encode(
        self,
        images: Array,
        key: Optional[PRNGKey] = None,
        *,
        deterministic: bool = True,
    ) -> Array:
        """Encode RGB images in [-1, 1] range to scaled latents (NHWC).

        Args:
            images: Batch of images shaped (B, H, W, 3) normalized to [-1, 1].
            key: RNG key (only used when `deterministic=False`).
            deterministic: When True, return posterior mean instead of sampling.
        """
        pixels = self._to_nchw(images)
        posterior = self.module.apply(
            {"params": self.params},
            pixels,
            method=self.module.encode,
        ).latent_dist

        if deterministic:
            latents = posterior.mean
        else:
            if key is None:
                raise ValueError("PRNGKey is required when deterministic=False")
            latents = posterior.sample(key)

        # Diffusers already returns NHWC latents; avoid extra transposes
        latents = latents * self.scaling_factor
        return latents

    @partial(jax.jit, static_argnums=(0,))
    def decode(self, latents: Array) -> Array:
        """Decode NHWC latents (scaled) back to RGB images in [-1, 1]."""
        latents_unscaled = latents / self.scaling_factor
        images = self.module.apply(
            {"params": self.params},
            latents_unscaled,
            method=self.module.decode,
        ).sample
        return self._to_nhwc(images)


def create_stable_vae(pretrained_path: Optional[str] = None) -> StableVAE:
    """Factory used across training / evaluation entry-points."""
    checkpoint_id = pretrained_path or "pcuenq/sd-vae-ft-mse-flax"
    return StableVAE(checkpoint_id)

Thank you.

Best regards,
Young Jin

[willisma]

Hi Jin,

Thanks for your kind words! Yes, I think it will work! Before you launch training, for a sanity check maybe enabling visualize_reconstruction to see if the reconstruction results make sense.

[snoop2head]

@willisma

I think it works. Mind if I create a pull request regarding this?
I will close the issue after the pull request

[willisma]

That would be great! Thanks @snoop2head!