Model Card for CanvasVAE¶

class CanvasVAEConfig(PretrainedConfig):
    """Configuration for a CanvasVAE model trained on RICO.

    A layout has a document-level element count and eight categorical
    per-element fields. Box coordinates and sizes are discretized into
    ``num_bins`` bins over ``[0, 1]``; ``component``, ``icon``, and
    ``text_button`` use the lookup tables in ``vocabularies``, whose index 0 is
    the out-of-vocabulary token.

    Args:
        vocabularies: Lookup tables for ``component``, ``icon``, and
            ``text_button``, each starting with the out-of-vocabulary token.
            Required for a usable model; the default only exists so that
            Transformers can instantiate an empty config while serializing.
        max_length: Maximum number of elements per layout.
        num_bins: Number of discretization bins for box coordinates and sizes.
        latent_dim: Hidden and latent dimension.
        num_blocks: Number of transformer blocks in the encoder and decoder.
        num_heads: Number of attention heads.
        dropout: Dropout rate of position embeddings and blocks.
        kl_weight: Weight of the KL divergence term.
        l2_weight: Weight of the squared-norm penalty on dense and embedding
            weights.
        layer_norm_epsilon: Layer-normalization epsilon.
        batch_norm_epsilon: Batch-normalization epsilon.
        batch_norm_momentum: Batch-normalization moving-average momentum.
        **kwargs: Extra ``PretrainedConfig`` fields.

    Examples:
        >>> config = CanvasVAEConfig(
        ...     vocabularies={
        ...         "component": ["[UNK]", "", "Text"],
        ...         "icon": ["[UNK]", ""],
        ...         "text_button": ["[UNK]", ""],
        ...     },
        ...     latent_dim=16,
        ...     num_heads=2,
        ... )
        >>> config.field_sizes["component"], config.primary_label_id
        (3, 1)
    """

    model_type = "canvas-vae"

    def __init__(
        self,
        vocabularies: Mapping[str, Sequence[str]] | None = None,
        max_length: int = 50,
        num_bins: int = 64,
        latent_dim: int = 256,
        num_blocks: int = 1,
        num_heads: int = 8,
        dropout: float = 0.1,
        kl_weight: float = 16.0,
        l2_weight: float = 1e-6,
        layer_norm_epsilon: float = 1e-3,
        batch_norm_epsilon: float = 1e-3,
        batch_norm_momentum: float = 0.99,
        **kwargs: str | float | bool | None,
    ) -> None:
        """Initialize the config; see the class docstring for arguments."""
        tables = {
            str(key): [str(token) for token in tokens]
            for key, tokens in (vocabularies or {}).items()
        }
        components = tables.get(CanvasVAEField.component, [])
        kwargs.pop("id2label", None)
        kwargs.pop("label2id", None)
        super().__init__(
            id2label=dict(enumerate(components)),
            label2id={token: index for index, token in enumerate(components)},
            **kwargs,  # ty: ignore[invalid-argument-type]
        )
        self.vocabularies = tables
        self.max_length = max_length
        self.num_bins = num_bins
        self.latent_dim = latent_dim
        self.num_blocks = num_blocks
        self.num_heads = num_heads
        self.dropout = dropout
        self.kl_weight = kl_weight
        self.l2_weight = l2_weight
        self.layer_norm_epsilon = layer_norm_epsilon
        self.batch_norm_epsilon = batch_norm_epsilon
        self.batch_norm_momentum = batch_norm_momentum

    @property
    def field_sizes(self) -> dict[str, int]:
        """Return the number of classes of every per-element field.

        Raises:
            ValueError: If a required vocabulary is missing.
        """
        missing = [key for key in VOCABULARY_FIELDS if key not in self.vocabularies]
        if missing:
            raise ValueError(f"CanvasVAEConfig is missing vocabularies: {missing}")

        sizes = {str(key): self.num_bins for key in GEOMETRY_FIELDS}
        sizes[CanvasVAEField.clickable] = NUM_CLICKABLE_CLASSES
        for key in VOCABULARY_FIELDS:
            sizes[key] = len(self.vocabularies[key])

        return sizes

    @property
    def primary_label_id(self) -> int:
        """Return the ``component`` id of the empty label used as background."""
        return self.vocabularies[CanvasVAEField.component].index(PRIMARY_LABEL_TOKEN)

class CanvasVAEField(StrEnum):
    """Per-element fields of a CanvasVAE layout, in model column order."""

    left = auto()
    top = auto()
    width = auto()
    height = auto()
    clickable = auto()
    component = auto()
    icon = auto()
    text_button = auto()

class CanvasVAEModel(CanvasVAEPreTrainedModel):
    """CanvasVAE variational autoencoder over layout element sequences.

    Given layout ids, the model encodes them, takes a posterior sample in
    training mode (or when ``posterior_noise`` is given) and the posterior mean
    otherwise, and decodes with the ground-truth element mask in training mode
    or the predicted element count in evaluation mode. Given only
    ``latents``, it decodes them with the predicted element count.

    Args:
        config: Model config with vocabularies.

    Examples:
        >>> config = CanvasVAEConfig(
        ...     vocabularies={
        ...         "component": ["[UNK]", "", "Text"],
        ...         "icon": ["[UNK]", ""],
        ...         "text_button": ["[UNK]", ""],
        ...     },
        ...     max_length=4,
        ...     num_bins=8,
        ...     latent_dim=16,
        ...     num_heads=2,
        ... )
        >>> model = CanvasVAEModel(config).eval()
        >>> output = model(latents=torch.zeros(2, 16))
        >>> output.length_logits.shape
        torch.Size([2, 4])
    """

    def __init__(self, config: CanvasVAEConfig) -> None:
        """Create the encoder and decoder."""
        super().__init__(config)
        self.encoder = CanvasVAEEncoder(config)
        self.decoder = CanvasVAEDecoder(config)
        self.post_init()

    def forward(
        self,
        num_elements: Int[torch.Tensor, "batch"] | None = None,
        element_ids: Int[torch.Tensor, "batch elements fields"] | None = None,
        *,
        latents: Float[torch.Tensor, "batch dim"] | None = None,
        posterior_noise: Float[torch.Tensor, "batch dim"] | None = None,
    ) -> CanvasVAEModelOutput:
        """Encode and decode a layout batch, or decode latent codes.

        Args:
            num_elements: Element count of every layout.
            element_ids: Field ids in :class:`CanvasVAEField` order, padded to
                the longest layout in the batch.
            latents: Latent codes to decode instead of encoding a layout.
            posterior_noise: Standard-normal noise for the posterior sample.
                Drawn from the global generator in training mode when omitted.

        Returns:
            Model outputs; ``loss`` is set only in training mode.

        Raises:
            ValueError: If neither a layout nor latent codes are given.
        """
        if latents is not None:
            length_logits, logits, mask = self.decoder(latents)
            return CanvasVAEModelOutput(
                length_logits=length_logits,
                element_logits=logits,
                mask=mask,
                latents=latents,
            )

        if num_elements is None or element_ids is None:
            raise ValueError("pass num_elements and element_ids, or latents")

        length_ids = num_elements - 1
        z_mean, z_log_var = self.encoder(length_ids, element_ids)
        kl_divergence = -0.5 * torch.mean(
            1 + z_log_var - z_mean.square() - torch.exp(z_log_var)
        )
        if posterior_noise is None and self.training:
            posterior_noise = randn(
                z_mean.shape, device=z_mean.device, dtype=z_mean.dtype
            )

        latents = z_mean
        if posterior_noise is not None:
            latents = z_mean + torch.exp(0.5 * z_log_var) * posterior_noise

        mask = length_mask(length_ids, element_ids.shape[1]) if self.training else None
        length_logits, logits, mask = self.decoder(latents, mask)
        output = CanvasVAEModelOutput(
            length_logits=length_logits,
            element_logits=logits,
            mask=mask,
            latents=latents,
            z_mean=z_mean,
            z_log_var=z_log_var,
            kl_divergence=kl_divergence,
        )
        if not self.training:
            return output

        losses = {
            "length": nn.functional.cross_entropy(length_logits, length_ids),
            **{
                key: _sequence_cross_entropy(logits[key], element_ids[..., index], mask)
                for index, key in enumerate(CanvasVAEField)
            },
        }
        reconstruction = sum(losses.values(), torch.zeros((), device=z_mean.device))
        output.reconstruction_losses = losses
        output.loss = reconstruction + self.config.kl_weight * kl_divergence
        return output

@dataclass
class CanvasVAEModelOutput(ModelOutput):
    """Outputs of :class:`CanvasVAEModel`.

    Attributes:
        loss: Reconstruction plus weighted KL loss; present when layout inputs
            are given and the model is in training mode.
        length_logits: Logits over the zero-based element count.
        element_logits: Logits of every per-element field, keyed by field name.
        mask: Valid-element mask used by the decoder.
        latents: Latent codes fed to the decoder.
        z_mean: Posterior mean; present when layout inputs are given.
        z_log_var: Posterior log-variance; present when layout inputs are given.
        kl_divergence: Unweighted KL divergence to the standard normal prior.
        reconstruction_losses: Per-field reconstruction losses, including
            ``length``; present together with ``loss``.
    """

    loss: Float[torch.Tensor, ""] | None = None
    length_logits: Float[torch.Tensor, "batch lengths"] | None = None
    element_logits: dict[str, Float[torch.Tensor, "batch elements classes"]] | None = (
        None
    )
    mask: Bool[torch.Tensor, "batch elements"] | None = None
    latents: Float[torch.Tensor, "batch latent"] | None = None
    z_mean: Float[torch.Tensor, "batch latent"] | None = None
    z_log_var: Float[torch.Tensor, "batch latent"] | None = None
    kl_divergence: Float[torch.Tensor, ""] | None = None
    reconstruction_losses: dict[str, Float[torch.Tensor, ""]] | None = None

class CanvasVAEPipeline(LayoutGenerationPipeline):
    """Sample RICO layouts from the CanvasVAE prior.

    Each call draws latent codes from the standard normal prior, decodes the
    element count and every element field with ``argmax``, and returns
    normalized center ``xywh`` boxes clamped to ``[0, 1]`` with ``component``
    labels. ``intermediates`` holds the ``clickable``, ``icon``, and
    ``text_button`` ids and the latent codes.

    Args:
        model: CanvasVAE model.

    Examples:
        >>> from canvas_vae import CanvasVAEConfig
        >>> config = CanvasVAEConfig(
        ...     vocabularies={
        ...         "component": ["[UNK]", "", "Text"],
        ...         "icon": ["[UNK]", ""],
        ...         "text_button": ["[UNK]", ""],
        ...     },
        ...     max_length=4,
        ...     num_bins=8,
        ...     latent_dim=16,
        ...     num_heads=2,
        ... )
        >>> pipe = CanvasVAEPipeline(model=CanvasVAEModel(config))
        >>> out = pipe(batch_size=2, seed=0)
        >>> out.bbox.shape[0], out.bbox.shape[-1]
        (2, 4)
    """

    config_class: ClassVar[type[CanvasVAEConfig]] = CanvasVAEConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = {
        "model": PipelineComponentSpec(attribute_name="model", loader=_load_model)
    }

    def __init__(self, model: CanvasVAEModel) -> None:
        """Wrap a model; the pipeline config is the model config."""
        super().__init__(model.config)
        self.model = model

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> CanvasVAEPipeline:
        """Build the pipeline from a loaded model."""
        del config
        return cls(model=cast(CanvasVAEModel, components["model"]))

    @torch.no_grad()
    def __call__(  # ty: ignore[invalid-method-override]
        self,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.unconditional,
        labels: Shaped[torch.Tensor, "..."]
        | Shaped[np.ndarray, "..."]
        | list[list[int]]
        | None = None,
        bbox: Shaped[torch.Tensor, "..."]
        | Shaped[np.ndarray, "..."]
        | list[list[list[float]]]
        | None = None,
        mask: Shaped[torch.Tensor, "..."]
        | Shaped[np.ndarray, "..."]
        | list[list[bool]]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: Literal["xywh", "ltwh", "ltrb"] = "xywh",
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
    ) -> LayoutGenerationOutput | dict[str, Shaped[torch.Tensor, "..."]]:
        """Generate layouts from prior samples.

        Args:
            batch_size: Number of layouts.
            seed: Seed for a local generator when ``generator`` is absent.
            generator: Generator for the prior samples; takes precedence over
                ``seed``.
            condition_type: Only ``unconditional`` is supported.
            labels: Unsupported; must be ``None``.
            bbox: Unsupported; must be ``None``.
            mask: Unsupported; must be ``None``.
            num_elements: Unsupported; the model predicts the element count.
            box_format: Only ``xywh`` output is supported.
            normalized: Only normalized output is supported.
            canvas_size: Unsupported; must be ``None``.
            num_inference_steps: Unsupported; decoding is a single pass.
            output_type: ``dataclass`` or ``dict``.
            return_intermediates: Whether to fill ``intermediates``.

        Returns:
            Generated layouts in the shared layout schema.

        Raises:
            NotImplementedError: If a condition or layout input other than
                unconditional generation is requested.
        """
        condition = normalize_condition_type(condition_type)
        if condition is not ConditionType.unconditional:
            raise NotImplementedError(
                f"CanvasVAE supports only unconditional generation, got {condition}"
            )

        given = {
            "labels": labels,
            "bbox": bbox,
            "mask": mask,
            "num_elements": num_elements,
            "canvas_size": canvas_size,
            "num_inference_steps": num_inference_steps,
        }
        unsupported = [name for name, value in given.items() if value is not None]
        if unsupported or BoxFormat(box_format) is not BoxFormat.xywh or not normalized:
            raise NotImplementedError(
                "CanvasVAE unconditional generation returns normalized xywh boxes "
                f"and takes no layout inputs; got {unsupported or box_format}"
            )

        generator = self.prepare_generator(generator=generator, seed=seed)
        device = self.model.device
        latents = randn(
            (batch_size, self.model.config.latent_dim),
            generator=generator,
            device=device,
            dtype=self.model.dtype,
        )
        was_training = self.model.training
        self.model.eval()
        try:
            output = self.model(latents=latents)
        finally:
            self.model.train(was_training)

        ids = {
            key: logits.argmax(dim=-1) for key, logits in output.element_logits.items()
        }
        geometry = torch.stack(
            [
                ids[CanvasVAEField.left],
                ids[CanvasVAEField.top],
                ids[CanvasVAEField.width],
                ids[CanvasVAEField.height],
            ],
            dim=-1,
        )
        boxes = clamp_boxes(
            ltwh_to_xywh(bins_to_ltwh(geometry, self.model.config.num_bins))
        )
        valid = cast(Bool[torch.Tensor, "batch elements"], output.mask)
        result = LayoutGenerationOutput(
            bbox=boxes * valid.unsqueeze(-1),
            labels=ids[CanvasVAEField.component] * valid,
            mask=valid,
            id2label=dict(self.model.config.id2label),
            intermediates=(
                {
                    "clickable": ids[CanvasVAEField.clickable],
                    "icon": ids[CanvasVAEField.icon],
                    "text_button": ids[CanvasVAEField.text_button],
                    "latents": latents,
                }
                if return_intermediates
                else None
            ),
        )
        if OutputType(output_type) is OutputType.dict:
            return dict(result)

        return result

class CanvasVAEProcessor(ProcessorMixin):
    """Encode flattened layouts into CanvasVAE model ids.

    Coordinates and sizes in ``[0, 1]`` map to ``num_bins`` equal-width bins
    whose float32 boundaries are ``linspace(0, 1, num_bins)[1:]``; values at or
    above a boundary fall in the next bin. Strings map through the lookup
    tables with unknown values at id 0. Layouts are padded to the longest one
    in the batch with zero geometry, ``clickable=0``, and the id of the empty
    string.

    Args:
        vocabularies: Lookup tables for ``component``, ``icon``, and
            ``text_button``.
        num_bins: Number of geometry bins.

    Examples:
        >>> processor = CanvasVAEProcessor(
        ...     vocabularies={
        ...         "component": ["[UNK]", "", "Text"],
        ...         "icon": ["[UNK]", ""],
        ...         "text_button": ["[UNK]", ""],
        ...     },
        ...     num_bins=4,
        ... )
        >>> element = {"left": 0.5, "top": 0.0, "width": 1.0, "height": 0.2,
        ...            "clickable": 1, "component": "Text", "icon": "", "text_button": ""}
        >>> processor([[element]])["element_ids"].tolist()
        [[[1, 0, 3, 0, 1, 2, 1, 1]]]
    """

    def __init__(
        self, vocabularies: Mapping[str, Sequence[str]], num_bins: int = 64
    ) -> None:
        """Store lookup tables and the bin count."""
        self.chat_template = None
        self.vocabularies = {key: list(tokens) for key, tokens in vocabularies.items()}
        self.token_ids = {
            key: {token: index for index, token in reversed(list(enumerate(tokens)))}
            for key, tokens in self.vocabularies.items()
        }
        self.num_bins = num_bins

    @property
    def bin_boundaries(self) -> Float[np.ndarray, "boundaries"]:
        """Return the float32 geometry bin boundaries."""
        return np.linspace(0.0, 1.0, self.num_bins)[1:].astype(np.float32)

    def discretize(
        self, values: Float[np.ndarray, "*shape"]
    ) -> Int[np.ndarray, "*shape"]:
        """Map normalized geometry values to bin ids.

        Args:
            values: Normalized coordinates or sizes.

        Returns:
            Bin ids in ``[0, num_bins)``.
        """
        return np.searchsorted(
            self.bin_boundaries, values.astype(np.float32), side="right"
        )

    def lookup(self, field: str, value: str) -> int:
        """Return the id of ``value`` in the ``field`` table, or 0 when unknown.

        Args:
            field: Vocabulary field name.
            value: String value.

        Returns:
            Lookup id.
        """
        return self.token_ids[field].get(value, 0)

    def __call__(self, layouts: Sequence[Sequence[RicoElement]]) -> BatchFeature:
        """Encode a batch of flattened layouts.

        Args:
            layouts: Element lists, each non-empty.

        Returns:
            ``num_elements`` with shape ``(batch,)`` and ``element_ids`` with
            shape ``(batch, elements, 8)`` in :class:`CanvasVAEField` order.

        Raises:
            ValueError: If a layout has no elements.
        """
        if any(not layout for layout in layouts):
            raise ValueError("every layout needs at least one element")

        width = max(len(layout) for layout in layouts)
        geometry = np.zeros((len(layouts), width, len(GEOMETRY_FIELDS)), np.float32)
        categorical = np.zeros((len(layouts), width, 4), np.int64)
        categorical[..., 1:] = [
            self.lookup(field, PRIMARY_LABEL_TOKEN) for field in VOCABULARY_FIELDS
        ]
        for row, layout in enumerate(layouts):
            for column, element in enumerate(layout):
                geometry[row, column] = [element[key] for key in GEOMETRY_FIELDS]  # ty: ignore[invalid-key]
                categorical[row, column] = [
                    element["clickable"],
                    *(self.lookup(key, element[key]) for key in VOCABULARY_FIELDS),  # ty: ignore[invalid-key]
                ]

        element_ids = np.concatenate([self.discretize(geometry), categorical], axis=-1)
        return BatchFeature(
            {
                "num_elements": torch.tensor([len(layout) for layout in layouts]),
                "element_ids": torch.from_numpy(element_ids),
            }
        )