Model Card for creative-graphic-design/pixel-vae-crello¶

class PixelVAEConfig(PretrainedConfig):
    """Configuration for the 256 × 256 RGBA MobileNetV2 VAE.

    Args:
        latent_dim: Dimension of each posterior vector.
        quantize_factor: Number of low bits removed from each target channel.
        kl_weight: Multiplier on the mean KL divergence.
        l2_weight: L2 coefficient for the two variational-head dense layers.
        batch_norm_epsilon: Variance epsilon used by batch normalization.
        backbone_batch_norm_momentum: Moving-average momentum in MobileNetV2.
        decoder_batch_norm_momentum: Moving-average momentum in the decoder.
        **kwargs: Extra Transformers configuration fields.

    Examples:
        >>> PixelVAEConfig().latent_dim
        256
    """

    model_type = "pixel-vae"

    def __init__(
        self,
        *,
        latent_dim: int = 256,
        quantize_factor: int = 4,
        kl_weight: float = 100.0,
        l2_weight: float = 1e-6,
        batch_norm_epsilon: float = 1e-3,
        backbone_batch_norm_momentum: float = 0.999,
        decoder_batch_norm_momentum: float = 0.99,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize the configuration with the production model defaults."""
        if latent_dim <= 0:
            raise ValueError("latent_dim must be positive")

        if quantize_factor < 0 or quantize_factor > 8:
            raise ValueError("quantize_factor must be between 0 and 8")

        if kl_weight < 0 or l2_weight < 0:
            raise ValueError("loss weights must be non-negative")

        if batch_norm_epsilon <= 0:
            raise ValueError("batch_norm_epsilon must be positive")

        if not 0 <= backbone_batch_norm_momentum < 1:
            raise ValueError("backbone_batch_norm_momentum must be in [0, 1)")

        if not 0 <= decoder_batch_norm_momentum < 1:
            raise ValueError("decoder_batch_norm_momentum must be in [0, 1)")

        super().__init__(**kwargs)  # ty: ignore[invalid-argument-type]
        self.latent_dim = latent_dim
        self.quantize_factor = quantize_factor
        self.kl_weight = kl_weight
        self.l2_weight = l2_weight
        self.batch_norm_epsilon = batch_norm_epsilon
        self.backbone_batch_norm_momentum = backbone_batch_norm_momentum
        self.decoder_batch_norm_momentum = decoder_batch_norm_momentum

class PixelVAEImageProcessor(BaseImageProcessor):
    """Decode PNG bytes to unmodified 256 × 256 RGBA model inputs.

    Examples:
        >>> from io import BytesIO
        >>> from PIL import Image
        >>> stream = BytesIO()
        >>> Image.new("RGBA", (256, 256)).save(stream, format="PNG")
        >>> batch = PixelVAEImageProcessor()(stream.getvalue())
        >>> tuple(batch.pixel_values.shape)
        (1, 4, 256, 256)
    """

    model_input_names = ["pixel_values"]

    def __init__(self, **kwargs: str | int | float | bool | None) -> None:
        """Initialize the processor without adding transform settings."""
        super().__init__(**kwargs)

    def preprocess(
        self,
        images: bytes | Sequence[bytes],
        *,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchFeature:
        """Decode PNG bytes and scale the exact RGBA pixels to ``[0, 1]``.

        Args:
            images: One PNG byte string or a sequence of PNG byte strings.
            return_tensors: Tensor framework. Only ``pt`` is supported.

        Returns:
            A batch feature with float32 ``pixel_values`` shaped ``[B, 4, 256, 256]``.

        Raises:
            ValueError: If a PNG is malformed, is not 256 × 256, or another
                tensor framework is requested.
        """
        if return_tensors != "pt":
            raise ValueError("PixelVAEImageProcessor only supports return_tensors='pt'")

        image_batch = [images] if isinstance(images, bytes) else list(images)
        if not image_batch:
            raise ValueError("at least one PNG image is required")

        decoded = np.stack([decode_pixelvae_png(image) for image in image_batch])
        pixel_values: Float[torch.Tensor, "batch 4 256 256"] = (
            torch.from_numpy(decoded)
            .permute(0, 3, 1, 2)
            .to(torch.float32)
            .mul(1.0 / 255.0)
        )
        return BatchFeature({"pixel_values": pixel_values})

class PixelVAEEncoder(PreTrainedModel):
    """MobileNetV2 encoder that exports posterior-mean embeddings."""

    config_class = PixelVAEConfig
    base_model_prefix = "pixel_vae_encoder"
    main_input_name = "pixel_values"

    def __init__(self, config: PixelVAEConfig) -> None:
        """Create the feature extractor and posterior heads."""
        super().__init__(config)
        self.backbone = PixelVAEMobileNetV2(config)
        self.z_mean = nn.Linear(self.backbone.output_channels, config.latent_dim)
        self.z_log_variance = nn.Linear(
            self.backbone.output_channels, config.latent_dim
        )
        self.post_init()

    def forward(
        self, pixel_values: Float[torch.Tensor, "batch 4 256 256"]
    ) -> PixelVAEEncoderOutput:
        """Return posterior mean and log variance without sampling."""
        _validate_pixel_values(pixel_values)
        features = self.backbone(pixel_values)
        return PixelVAEEncoderOutput(
            posterior_mean=self.z_mean(features),
            posterior_log_variance=self.z_log_variance(features),
        )

    def encode(
        self, pixel_values: Float[torch.Tensor, "batch 4 256 256"]
    ) -> Float[torch.Tensor, "batch latent"]:
        """Return the float32 posterior mean used by Crello documents."""
        result = self(pixel_values)
        assert result.posterior_mean is not None
        return result.posterior_mean

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Conv2d):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, PixelVAEBatchNorm2d):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

@dataclass
class PixelVAEEncoderOutput(ModelOutput):
    """Posterior statistics emitted by :class:`PixelVAEEncoder`."""

    posterior_mean: Float[torch.Tensor, "batch latent"] | None = None
    posterior_log_variance: Float[torch.Tensor, "batch latent"] | None = None

class PixelVAEModel(PreTrainedModel):
    """Transformers model implementing the RGBA PixelVAE objective."""

    config_class = PixelVAEConfig
    base_model_prefix = "pixel_vae"
    main_input_name = "pixel_values"

    def __init__(self, config: PixelVAEConfig) -> None:
        """Create the posterior encoder and autoregressive image decoder."""
        super().__init__(config)
        self.encoder = PixelVAEEncoder(config)
        self.decoder = PixelVAEDecoder(config)
        self.post_init()

    def encode(
        self, pixel_values: Float[torch.Tensor, "batch 4 256 256"]
    ) -> Float[torch.Tensor, "batch latent"]:
        """Export the encoder posterior mean without sampling."""
        return self.encoder.encode(pixel_values)

    def forward(
        self,
        pixel_values: Float[torch.Tensor, "batch 4 256 256"],
        *,
        sample: bool | None = None,
        epsilon: Float[torch.Tensor, "batch latent"] | None = None,
    ) -> PixelVAEOutput:
        """Compute logits and the source reconstruction, KL, and L2 losses.

        Args:
            pixel_values: Float32 RGBA images scaled to ``[0, 1]``.
            sample: Whether to sample the posterior; defaults to module mode.
            epsilon: Optional fixed standard-normal noise for deterministic runs.

        Returns:
            Decoder logits, posterior statistics, latent vectors, and losses.

        Raises:
            ValueError: If image dimensions or sampling-noise shape are invalid.
            TypeError: If inputs are not float32 tensors.

        Examples:
            >>> model = PixelVAEModel(PixelVAEConfig(latent_dim=8)).eval()
            >>> image = torch.zeros(1, 4, 256, 256)
            >>> model(image).posterior_mean.shape
            torch.Size([1, 8])
        """
        encoded = self.encoder(pixel_values)
        assert encoded.posterior_mean is not None
        assert encoded.posterior_log_variance is not None
        z_mean, z_log_variance = encoded.posterior_mean, encoded.posterior_log_variance
        should_sample = self.training if sample is None else sample
        if should_sample:
            if epsilon is None:
                epsilon = randn(
                    z_mean.shape,
                    device=z_mean.device,
                    dtype=z_mean.dtype,
                )

            if epsilon.shape != z_mean.shape:
                raise ValueError("epsilon must have the same shape as the posterior")

            latents = z_mean + z_log_variance.mul(0.5).exp() * epsilon
        else:
            if epsilon is not None:
                raise ValueError("epsilon is only valid when sample=True")

            latents = z_mean

        logits = self.decoder(latents)
        reconstruction_loss = pixelvae_reconstruction_loss(
            pixel_values, logits, self.config.quantize_factor
        )
        kl_divergence = (
            -0.5 * (1 + z_log_variance - z_mean.square() - z_log_variance.exp()).mean()
        )
        l2_loss = self.config.l2_weight * (
            self.encoder.z_mean.weight.square().sum()
            + self.encoder.z_mean.bias.square().sum()
            + self.encoder.z_log_variance.weight.square().sum()
            + self.encoder.z_log_variance.bias.square().sum()
        )
        loss = reconstruction_loss + self.config.kl_weight * kl_divergence + l2_loss
        return PixelVAEOutput(
            loss=loss,
            reconstruction_loss=reconstruction_loss,
            kl_divergence=kl_divergence,
            logits=logits,
            posterior_mean=z_mean,
            posterior_log_variance=z_log_variance,
            latents=latents,
        )

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Conv2d | nn.ConvTranspose2d):
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, PixelVAEBatchNorm2d):
            nn.init.ones_(module.weight)
            nn.init.zeros_(module.bias)

@dataclass
class PixelVAEOutput(ModelOutput):
    """Reconstruction, posterior statistics, and losses from PixelVAE."""

    loss: Float[torch.Tensor, ""] | None = None
    reconstruction_loss: Float[torch.Tensor, ""] | None = None
    kl_divergence: Float[torch.Tensor, ""] | None = None
    logits: Float[torch.Tensor, "batch height width channels levels"] | None = None
    posterior_mean: Float[torch.Tensor, "batch latent"] | None = None
    posterior_log_variance: Float[torch.Tensor, "batch latent"] | None = None
    latents: Float[torch.Tensor, "batch latent"] | None = None