Skip to content

moss_image: the FLUX.2 image decoder as a native Metal forward #1504

Open

Description

@BobClawblaw

MOSS-Image's image decoder is FLUX.2's small VAE decoder (AutoencoderKLFlux2, the checkpoint's vae/: about 28M decoder parameters in float32). The pipeline undoes the latent batch norm with the saved statistics, unpacks the 2x2 patches, and decodes in float32: post_quant_conv, conv_in, a middle block (two residual blocks around one single-head self-attention), four upsampling stages of GroupNorm, SiLU and 3x3 convolutions, then conv_out and a clamp.

This issue tracks that decoder as a native float32 Metal forward inside the moss_image family (#1456, #1457, #1459), so that tf-moss-image writes the picture itself instead of handing latents to torch:

  • the checkpoint's float32 weights as published, run in float32 as the pipeline does;
  • the family's own kernels (3x3 and 1x1 convolution, GroupNorm with SiLU, float32 attention) registered with the kernel contract;
  • the pictures checked against the publisher's float32 decode of the same latents.

No shared code changes. The starting noise (torch's CPU generator) stays out of scope.

Activity

  1. BobClawblaw commented on Oct 27, 2026

    @BobClawblaw
    ContributorAuthor

    /atc claim

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions