Layoutvae

class LayoutVAEConfig(PretrainedConfig):
    """Configuration for LayoutVAE.

    Args:
        dataset_name: Dataset key. The first release supports PubLayNet.
        num_labels: Public label vocabulary size.
        internal_num_labels: Internal label-set size including the empty label.
        max_position_embeddings: Maximum number of layout elements.
        count_latent_dim: Latent dimension for the count module.
        bbox_latent_dim: Latent dimension for the box module.
        bbox_format: Internal box format.
        bbox_normalized: Whether internal boxes are normalized.
        id2label: Optional public ID-to-label mapping.
        label2id: Optional public label-to-ID mapping.
        **kwargs: Extra `PretrainedConfig` keyword arguments.

    Raises:
        ValueError: If `dataset_name` is not supported.

    Examples:
        >>> LayoutVAEConfig().model_type
        'layoutvae'
    """

    model_type = "layoutvae"

    def __init__(
        self,
        dataset_name: DatasetName | str = DatasetName.publaynet,
        num_labels: int = 5,
        internal_num_labels: int = 6,
        max_position_embeddings: int | None = None,
        count_latent_dim: int = 32,
        bbox_latent_dim: int = 32,
        bbox_format: BoxFormat | str = BoxFormat.ltwh,
        bbox_normalized: bool = True,
        id2label: Id2LabelMapping | None = None,
        label2id: dict[str, int] | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize a LayoutVAE config.

        Args:
            dataset_name: Dataset key. The first release supports PubLayNet.
            num_labels: Public label vocabulary size.
            internal_num_labels: Internal label-set size including the empty label.
            max_position_embeddings: Maximum number of generated elements.
            count_latent_dim: Latent dimension for the count module.
            bbox_latent_dim: Latent dimension for the box module.
            bbox_format: Internal box format.
            bbox_normalized: Whether internal boxes are normalized.
            id2label: Optional public ID-to-label mapping.
            label2id: Optional public label-to-ID mapping.
            **kwargs: Extra `PretrainedConfig` keyword arguments.

        Raises:
            ValueError: If `dataset_name` is not PubLayNet.

        Examples:
            >>> LayoutVAEConfig(dataset_name="publaynet").num_labels
            5
        """
        canonical_dataset = normalize_dataset_name(dataset_name)
        if canonical_dataset is not DatasetName.publaynet:
            raise ValueError("LayoutVAE v1 supports only dataset_name='publaynet'")

        raw_id2label = id2label or id2label_for_dataset(canonical_dataset)
        normalized_id2label = {int(k): v for k, v in raw_id2label.items()}
        normalized_label2id = label2id or label2id_for_dataset(canonical_dataset)
        super().__init__(id2label=normalized_id2label, label2id=normalized_label2id)
        for key, value in kwargs.items():
            setattr(self, key, value)
        self.dataset_name = str(canonical_dataset)
        self.num_labels = num_labels
        self.internal_num_labels = internal_num_labels
        self.max_position_embeddings = (
            max_position_embeddings or max_elements_for_dataset(canonical_dataset)
        )
        self.count_latent_dim = count_latent_dim
        self.bbox_latent_dim = bbox_latent_dim
        self.bbox_format = str(normalize_box_format(bbox_format))
        self.bbox_normalized = bbox_normalized
        self.architectures = ["LayoutVAEModel"]

class LayoutVAEModel(PreTrainedModel):
    """Transformers-compatible LayoutVAE model.

    Args:
        config: LayoutVAE configuration.

    Examples:
        >>> model = LayoutVAEModel(LayoutVAEConfig())
        >>> model.config.model_type
        'layoutvae'
    """

    config_class = LayoutVAEConfig
    base_model_prefix = "layoutvae"
    supports_gradient_checkpointing = False

    def __init__(self, config: LayoutVAEConfig) -> None:
        """Initialize LayoutVAE submodules."""
        super().__init__(config)
        self.countvae = CountVAEModel(
            config.internal_num_labels, config.count_latent_dim
        )
        self.bboxvae = BboxVAEModel(
            config.internal_num_labels,
            4,
            config.max_position_embeddings,
            config.bbox_latent_dim,
        )
        self.post_init()

    def forward(
        self,
        label_set: Float[torch.Tensor, "batch internal_labels"],
        *,
        count_latents: Float[torch.Tensor, "batch internal_labels latent"]
        | None = None,
        bbox_latents: Float[torch.Tensor, "batch elements latent"] | None = None,
        bbox_noise: Float[torch.Tensor, "batch elements 4"] | None = None,
        class_counts: Float[torch.Tensor, "batch internal_labels"] | None = None,
        count_samples: Float[torch.Tensor, "batch internal_labels"] | None = None,
        generator: torch.Generator | None = None,
        return_dict: bool = True,
    ) -> (
        LayoutVAEModelOutput
        | tuple[
            Float[torch.Tensor, "batch elements 4"],
            Float[torch.Tensor, "batch elements 4"],
            Int[torch.Tensor, "batch elements"],
            Bool[torch.Tensor, "batch elements"],
            Float[torch.Tensor, "batch internal_labels"],
            Int[torch.Tensor, "batch elements"],
        ]
    ):
        """Run label-conditioned layout generation.

        Args:
            label_set: Six-way label-set tensor.
            count_latents: Optional fixed count latents.
            bbox_latents: Optional fixed box latents.
            bbox_noise: Optional fixed output noise.
            class_counts: Optional fixed six-way class counts.
            count_samples: Optional fixed count samples.
            generator: Optional PyTorch random generator.
            return_dict: Whether to return a dataclass.

        Returns:
            Model output dataclass or tuple.

        Raises:
            ValueError: If shapes are invalid.

        Examples:
            >>> model = LayoutVAEModel(LayoutVAEConfig())
            >>> label_set = torch.tensor([[0, 1, 0, 0, 0, 1]], dtype=torch.float32)
            >>> out = model(label_set, class_counts=torch.tensor([[7, 1, 0, 0, 0, 1.]]))
            >>> tuple(out.bbox.shape)
            (1, 9, 4)
        """
        label_set = label_set.to(device=self.device, dtype=self.dtype)
        if label_set.ndim != 2 or label_set.shape[1] != self.config.internal_num_labels:
            raise ValueError(
                "label_set must have shape (batch, config.internal_num_labels)"
            )

        if class_counts is None:
            class_counts = self.countvae(
                label_set,
                latents=count_latents,
                count_samples=count_samples,
                generator=generator,
            )
            class_counts = self._normalize_counts(class_counts)
        else:
            class_counts = class_counts.to(
                device=label_set.device, dtype=label_set.dtype
            )
            if class_counts.shape != label_set.shape:
                raise ValueError("class_counts must match label_set shape")

        class_labels = self._labels_from_counts(class_counts)
        raw_ltwh = self.bboxvae(
            class_counts,
            class_labels,
            latents=bbox_latents,
            output_noise=bbox_noise,
            generator=generator,
        )
        internal_ids = torch.argmax(class_labels, dim=2).to(dtype=torch.long)
        labels = torch.clamp(internal_ids - 1, min=0)
        mask = internal_ids != INTERNAL_EMPTY_LABEL_ID
        public_bbox = clamp_boxes(ltwh_to_xywh(raw_ltwh))
        raw_ltwh = torch.where(mask.unsqueeze(-1), raw_ltwh, torch.zeros_like(raw_ltwh))
        public_bbox = torch.where(
            mask.unsqueeze(-1), public_bbox, torch.zeros_like(public_bbox)
        )
        if not return_dict:
            return raw_ltwh, public_bbox, labels, mask, class_counts, internal_ids
        return LayoutVAEModelOutput(
            raw_ltwh=raw_ltwh,
            bbox=public_bbox,
            labels=labels,
            mask=mask,
            class_counts=class_counts,
            label_set=label_set,
            internal_labels=internal_ids,
        )

    def _normalize_counts(
        self,
        class_counts: Float[torch.Tensor, "batch internal_labels"],
    ) -> Float[torch.Tensor, "batch internal_labels"]:
        counts = class_counts.clamp_min(0)
        denom = counts.sum(dim=1, keepdim=True).clamp_min(1)
        counts = torch.floor(self.config.max_position_embeddings * (counts / denom))
        totals = counts.sum(dim=1)
        shortfall = self.config.max_position_embeddings - totals
        counts[:, INTERNAL_EMPTY_LABEL_ID] = counts[
            :, INTERNAL_EMPTY_LABEL_ID
        ] + torch.clamp(shortfall, min=0)
        return counts

    def _labels_from_counts(
        self,
        class_counts: Float[torch.Tensor, "batch internal_labels"],
    ) -> Float[torch.Tensor, "batch elements internal_labels"]:
        labels = torch.zeros(
            (
                class_counts.shape[0],
                self.config.max_position_embeddings,
                self.config.internal_num_labels,
            ),
            device=class_counts.device,
            dtype=class_counts.dtype,
        )
        for batch_index, counts in enumerate(class_counts):
            position = 0
            for class_index in reversed(range(self.config.internal_num_labels)):
                count = int(counts[class_index].item())
                for _ in range(count):
                    if position >= self.config.max_position_embeddings:
                        break
                    labels[batch_index, position, class_index] = 1.0
                    position += 1
        return labels

@dataclass
class LayoutVAEModelOutput(ModelOutput):
    """Raw LayoutVAE model output.

    Args:
        raw_ltwh: Internal normalized left-top-width-height boxes.
        bbox: Public normalized center `xywh` boxes.
        labels: Public label IDs.
        mask: Valid-element mask.
        class_counts: Six-way class-count tensor.
        label_set: Optional label-set input.
        internal_labels: Optional six-way generated label IDs.

    Examples:
        >>> output = LayoutVAEModelOutput(
        ...     raw_ltwh=torch.zeros(1, 1, 4),
        ...     bbox=torch.zeros(1, 1, 4),
        ...     labels=torch.zeros(1, 1, dtype=torch.long),
        ...     mask=torch.ones(1, 1, dtype=torch.bool),
        ...     class_counts=torch.ones(1, 6),
        ... )
        >>> tuple(output.bbox.shape)
        (1, 1, 4)
    """

    raw_ltwh: Float[torch.Tensor, "batch elements 4"]
    bbox: Float[torch.Tensor, "batch elements 4"] = cast(
        Float[torch.Tensor, "batch elements 4"], None
    )
    labels: Int[torch.Tensor, "batch elements"] = cast(
        Int[torch.Tensor, "batch elements"], None
    )
    mask: Bool[torch.Tensor, "batch elements"] = cast(
        Bool[torch.Tensor, "batch elements"], None
    )
    class_counts: Float[torch.Tensor, "batch internal_labels"] = cast(
        Float[torch.Tensor, "batch internal_labels"], None
    )
    label_set: Float[torch.Tensor, "batch internal_labels"] | None = None
    internal_labels: Int[torch.Tensor, "batch elements"] | None = None

class LayoutVAEPipeline(LayoutGenerationPipeline):
    """Transformers pipeline for LayoutVAE label-conditioned generation.

    Args:
        model: LayoutVAE model instance.
        processor: Optional processor for label-set encoding.
        config: Optional root pipeline config. Defaults to `model.config`.
        device: Optional torch device.
        binary_output: Reserved compatibility flag.

    Examples:
        >>> model = LayoutVAEModel(LayoutVAEConfig())
        >>> pipe = LayoutVAEPipeline(model=model)
        >>> pipe.config.model_type
        'layoutvae'
    """

    config_class: ClassVar[type[PretrainedConfig]] = LayoutVAEConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = (
        model_processor_component_specs(
            model_loader=_load_model_component,
            processor_loader=_load_processor_component,
        )
    )

    config: LayoutVAEConfig
    model: LayoutVAEModel
    processor: LayoutVAEProcessor

    def __init__(
        self,
        model: LayoutVAEModel,
        processor: LayoutVAEProcessor | None = None,
        config: LayoutVAEConfig | None = None,
        device: int | torch.device | None = None,
        binary_output: bool = False,
    ) -> None:
        """Initialize a LayoutVAE pipeline."""
        _ = binary_output
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor or LayoutVAEProcessor(
            dataset_name=model.config.dataset_name,
            id2label=model.config.id2label,
        )
        if device is not None:
            self.to(_resolve_device(device))

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> LayoutVAEPipeline:
        """Build a pipeline from loaded root components."""
        return cls(
            config=cast(LayoutVAEConfig, config),
            model=cast(LayoutVAEModel, components["model"]),
            processor=cast(LayoutVAEProcessor, components["processor"]),
        )

    def _sanitize_parameters(
        self, **kwargs: LayoutVAEParam
    ) -> tuple[
        dict[str, LayoutVAEParam], dict[str, LayoutVAEParam], dict[str, LayoutVAEParam]
    ]:
        return {}, kwargs, {}

    def preprocess(
        self,
        input_: list[list[str | int]]
        | list[str | int]
        | Int[torch.Tensor, ...]
        | None = None,
        **preprocess_parameters: LayoutVAEParam,
    ) -> BatchEncoding:
        """Encode labels into model inputs."""
        labels = preprocess_parameters.pop("labels", input_)
        if labels is None:
            raise ValueError("labels are required for LayoutVAEPipeline")

        encoded = self.processor(
            cast(
                list[list[str | int]] | list[str | int] | Int[torch.Tensor, "..."],
                labels,
            )
        )
        encoded.update(preprocess_parameters)
        return encoded

    def _forward(
        self, model_inputs: dict[str, LayoutVAEParam], **forward_params: LayoutVAEParam
    ) -> LayoutGenerationOutput | LayoutVAEOutputDict:
        del forward_params
        label_set = torch.as_tensor(model_inputs.pop("label_set"), dtype=torch.float32)
        condition_type = cast(
            ConditionType | str, model_inputs.pop("condition_type", ConditionType.label)
        )
        options = _pop_generation_options(model_inputs)
        count_latents = cast(
            Float[torch.Tensor, "batch internal_labels latent"] | None,
            model_inputs.pop("count_latents", None),
        )
        bbox_latents = cast(
            Float[torch.Tensor, "batch elements latent"] | None,
            model_inputs.pop("bbox_latents", None),
        )
        bbox_noise = cast(
            Float[torch.Tensor, "batch elements 4"] | None,
            model_inputs.pop("bbox_noise", None),
        )
        class_counts = cast(
            Float[torch.Tensor, "batch internal_labels"] | None,
            model_inputs.pop("class_counts", None),
        )
        if model_inputs:
            unknown = ", ".join(sorted(model_inputs))
            raise ValueError(f"Unsupported generation kwargs: {unknown}")

        return self._generate(
            label_set=label_set,
            condition_type=condition_type,
            options=options,
            count_latents=count_latents,
            bbox_latents=bbox_latents,
            bbox_noise=bbox_noise,
            class_counts=class_counts,
        )

    def postprocess(
        self,
        model_outputs: LayoutGenerationOutput | LayoutVAEOutputDict,
        **kwargs: str | int | float | bool | None,
    ) -> LayoutGenerationOutput | LayoutVAEOutputDict:
        """Return generated layouts unchanged."""
        del kwargs
        return model_outputs

    @torch.no_grad()
    def __call__(
        self,
        labels: list[list[str | int]]
        | list[str | int]
        | Int[torch.Tensor, ...]
        | None = None,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.label,
        bbox: Float[torch.Tensor, "batch elements 4"] | None = None,
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = BoxFormat.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,
        count_latents: Float[torch.Tensor, "batch internal_labels latent"]
        | None = None,
        bbox_latents: Float[torch.Tensor, "batch elements latent"] | None = None,
        bbox_noise: Float[torch.Tensor, "batch elements 4"] | None = None,
        class_counts: Float[torch.Tensor, "batch internal_labels"] | None = None,
    ) -> LayoutGenerationOutput | LayoutVAEOutputDict:  # ty: ignore[invalid-method-override]
        """Generate PubLayNet layouts from label conditions.

        Args:
            labels: Public PubLayNet label strings or IDs.
            batch_size: Used when `labels` is omitted.
            seed: Optional random seed used when `generator` is absent.
            generator: Optional PyTorch random generator. Takes precedence.
            condition_type: Condition type or alias. Only `label` is supported.
            bbox: Reserved compatibility argument.
            mask: Reserved compatibility argument.
            num_elements: Reserved compatibility argument.
            box_format: Reserved compatibility argument.
            normalized: Reserved compatibility argument.
            canvas_size: Reserved compatibility argument.
            num_inference_steps: Reserved compatibility argument.
            output_type: Return format.
            return_intermediates: Whether to include raw generation tensors.
            count_latents: Optional fixed count latents.
            bbox_latents: Optional fixed box latents.
            bbox_noise: Optional fixed box output noise.
            class_counts: Optional fixed class counts.

        Returns:
            Layout generation output.

        Raises:
            ValueError: If labels are missing or condition options are unsupported.

        Examples:
            >>> pipe = LayoutVAEPipeline(LayoutVAEModel(LayoutVAEConfig()))
            >>> out = pipe(labels=["text"], class_counts=torch.tensor([[8, 1, 0, 0, 0, 0.]]))
            >>> tuple(out.bbox.shape)
            (1, 9, 4)
        """
        if labels is None:
            if batch_size < 1:
                raise ValueError("batch_size must be positive")

            labels = [["text"] for _ in range(batch_size)]
        encoded = self.processor(labels)

        option_kwargs: GenerationOptionsKwargs = {}
        option_kwargs["bbox"] = bbox
        option_kwargs["mask"] = mask
        option_kwargs["num_elements"] = num_elements
        option_kwargs["box_format"] = box_format
        option_kwargs["normalized"] = normalized
        option_kwargs["canvas_size"] = canvas_size
        option_kwargs["seed"] = seed
        option_kwargs["generator"] = generator
        option_kwargs["num_inference_steps"] = num_inference_steps
        option_kwargs["output_type"] = output_type
        option_kwargs["return_intermediates"] = return_intermediates

        options = _make_generation_options(option_kwargs)

        return self._generate(
            label_set=cast(
                Float[torch.Tensor, "batch internal_labels"], encoded["label_set"]
            ),
            condition_type=condition_type,
            options=options,
            count_latents=count_latents,
            bbox_latents=bbox_latents,
            bbox_noise=bbox_noise,
            class_counts=class_counts,
        )

    def _generate(
        self,
        *,
        label_set: Float[torch.Tensor, "batch internal_labels"],
        condition_type: ConditionType | str,
        options: GenerationOptions,
        count_latents: Float[torch.Tensor, "batch internal_labels latent"] | None,
        bbox_latents: Float[torch.Tensor, "batch elements latent"] | None,
        bbox_noise: Float[torch.Tensor, "batch elements 4"] | None,
        class_counts: Float[torch.Tensor, "batch internal_labels"] | None,
    ) -> LayoutGenerationOutput | LayoutVAEOutputDict:
        _ = (
            options.bbox,
            options.labels,
            options.mask,
            options.num_elements,
            options.normalized,
            options.canvas_size,
            options.num_inference_steps,
        )
        normalize_box_format(options.box_format)
        canonical = normalize_condition_type(condition_type)
        if canonical is not ConditionType.label:
            raise ValueError(f"Unsupported condition_type for layoutvae: {canonical}")

        device = next(self.model.parameters()).device
        prepared_generator = self.prepare_generator(
            generator=options.generator,
            seed=options.seed,
            device=device,
        )
        out = self.model(
            label_set.to(device=device),
            count_latents=count_latents,
            bbox_latents=bbox_latents,
            bbox_noise=bbox_noise,
            class_counts=class_counts,
            generator=prepared_generator,
            return_dict=True,
        )
        assert isinstance(out, LayoutVAEModelOutput)
        intermediates = None
        if options.return_intermediates:
            intermediates = {
                "condition_type": canonical,
                "raw_ltwh": out.raw_ltwh.detach().cpu(),
                "internal_labels": out.internal_labels.detach().cpu()
                if out.internal_labels is not None
                else None,
                "class_counts": out.class_counts.detach().cpu(),
            }
        layout = LayoutGenerationOutput(
            bbox=out.bbox.detach().cpu(),
            labels=out.labels.detach().cpu(),
            mask=out.mask.detach().cpu(),
            id2label={
                int(k): v for k, v in cast(Id2Label, self.config.id2label).items()
            },
            intermediates=intermediates,
        )
        resolved_output_type = normalize_output_type(options.output_type)
        if resolved_output_type is OutputType.dict:
            return dict(layout)
        if resolved_output_type is OutputType.dataclass:
            return layout
        assert_never(resolved_output_type)

class LayoutVAEProcessor(ProcessorMixin):
    """Encode PubLayNet labels into LayoutVAE label sets.

    Args:
        dataset_name: Dataset key. The first release supports PubLayNet.
        id2label: Optional public ID-to-label mapping.

    Raises:
        ValueError: If the dataset is unsupported.

    Examples:
        >>> processor = LayoutVAEProcessor()
        >>> processor.label2id["text"]
        0
    """

    config_name = "preprocessor_config.json"

    def __init__(
        self,
        dataset_name: DatasetName | str = DatasetName.publaynet,
        id2label: Id2LabelMapping | None = None,
    ) -> None:
        """Initialize the processor.

        Args:
            dataset_name: Dataset key. The first release supports PubLayNet.
            id2label: Optional public ID-to-label mapping.

        Raises:
            ValueError: If the dataset is unsupported.

        Examples:
            >>> LayoutVAEProcessor("publaynet").id2label[4]
            'figure'
        """
        self.chat_template = None
        canonical_dataset = normalize_dataset_name(dataset_name)
        if canonical_dataset is not DatasetName.publaynet:
            raise ValueError("LayoutVAEProcessor supports only PubLayNet")

        self.dataset_name = str(canonical_dataset)
        raw_id2label = id2label or id2label_for_dataset(canonical_dataset)
        self.id2label = {int(k): v for k, v in raw_id2label.items()}
        self.label2id = label2id_for_dataset(canonical_dataset)
        self.internal_id2label = {
            INTERNAL_EMPTY_LABEL_ID: "None",
            **{index + 1: label for index, label in self.id2label.items()},
        }

    def __call__(
        self,
        labels: list[list[str | int]] | list[str | int] | Int[torch.Tensor, "..."],
        *,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Encode public labels as a six-way label-set tensor.

        Args:
            labels: Public label names or IDs. A flat list is treated as one row.
            return_tensors: Tensor framework. Only `pt` is supported.

        Returns:
            Batch encoding with `label_set`.

        Raises:
            ValueError: If labels are empty, unknown, or tensors are unsupported.

        Examples:
            >>> encoded = LayoutVAEProcessor()(["text", "figure"])
            >>> encoded["label_set"].tolist()
            [[0.0, 1.0, 0.0, 0.0, 0.0, 1.0]]
        """
        if return_tensors != "pt":
            raise ValueError("LayoutVAEProcessor only supports return_tensors='pt'")

        rows = self._normalize_rows(labels)
        label_set = torch.zeros(
            len(rows), len(self.internal_id2label), dtype=torch.float32
        )
        for row_index, row in enumerate(rows):
            for label in row:
                public_id = self._label_to_id(label)
                label_set[row_index, public_id + 1] = 1.0
        return BatchEncoding({"label_set": label_set})

    def public_from_internal(
        self,
        internal_labels: Int[torch.Tensor, "batch elements"],
    ) -> tuple[
        Int[torch.Tensor, "batch elements"], Bool[torch.Tensor, "batch elements"]
    ]:
        """Map six-way labels to public labels and validity masks.

        Args:
            internal_labels: Internal label IDs where zero marks empty slots.

        Returns:
            Public label IDs and mask tensors.

        Examples:
            >>> processor = LayoutVAEProcessor()
            >>> labels, mask = processor.public_from_internal(torch.tensor([[0, 1, 5]]))
            >>> labels.tolist(), mask.tolist()
            ([[0, 0, 4]], [[False, True, True]])
        """
        labels = torch.clamp(internal_labels.to(dtype=torch.long) - 1, min=0)
        mask = internal_labels.to(dtype=torch.long) != INTERNAL_EMPTY_LABEL_ID
        return labels, mask

    def batch_decode(
        self,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
    ) -> list[list[DecodedLayoutRecord]]:
        """Decode layout tensors into records.

        Args:
            bbox: Public normalized center `xywh` boxes.
            labels: Public label IDs.
            mask: Optional valid-element mask.

        Returns:
            Nested records with label text, label ID, and box coordinates.

        Raises:
            KeyError: If a public label ID is unknown.

        Examples:
            >>> records = LayoutVAEProcessor().batch_decode(
            ...     torch.zeros(1, 1, 4), torch.tensor([[0]])
            ... )
            >>> records[0][0]["label"]
            'text'
        """
        bbox_t = torch.as_tensor(bbox, dtype=torch.float32)
        labels_t = torch.as_tensor(labels, dtype=torch.long)
        labels_t, bbox_t = self._ensure_batched(labels_t, bbox_t)
        mask_t = self._prepare_mask(mask, labels_t.shape)
        records: list[list[DecodedLayoutRecord]] = []
        for boxes, ids, valid in zip(bbox_t, labels_t, mask_t, strict=True):
            row: list[DecodedLayoutRecord] = []
            for box, label_id in zip(boxes[valid], ids[valid], strict=True):
                idx = int(label_id.item())
                row.append(
                    {
                        "label": self.id2label[idx],
                        "label_id": idx,
                        "bbox": box.tolist(),
                    }
                )
            records.append(row)
        return records

    def _ensure_batched(
        self, labels: Int[torch.Tensor, "..."], bbox: Float[torch.Tensor, "... 4"]
    ) -> tuple[
        Int[torch.Tensor, "batch elements"], Float[torch.Tensor, "batch elements 4"]
    ]:
        if labels.ndim != 1:
            return labels, bbox
        return labels.unsqueeze(0), bbox.unsqueeze(0)

    def _prepare_mask(
        self,
        mask: Bool[torch.Tensor, "batch elements"] | None,
        shape: torch.Size,
    ) -> Bool[torch.Tensor, "batch elements"]:
        if mask is None:
            return torch.ones(shape, dtype=torch.bool)
        mask_t = torch.as_tensor(mask, dtype=torch.bool)
        return mask_t.unsqueeze(0) if mask_t.ndim == 1 else mask_t

    def _normalize_rows(
        self,
        labels: list[list[str | int]] | list[str | int] | Int[torch.Tensor, "..."],
    ) -> list[list[str | int]]:
        if isinstance(labels, torch.Tensor):
            if labels.ndim == 0 or labels.ndim > 2:
                raise ValueError("labels tensor must have one or two dimensions")

            if labels.ndim == 1:
                return [[int(value) for value in labels.tolist()]]
            return [[int(value) for value in row] for row in labels.tolist()]
        if not labels:
            raise ValueError("labels must not be empty")

        contains_rows = [isinstance(item, list) for item in labels]
        if any(contains_rows) and not all(contains_rows):
            raise ValueError("labels must be a flat list or list of rows")

        if all(contains_rows):
            return [list(row) for row in cast(list[list[str | int]], labels)]
        return [list(cast(list[str | int], labels))]

    def _label_to_id(self, label: str | int) -> int:
        if not isinstance(label, int):
            if label in self.label2id:
                return self.label2id[label]
            raise ValueError(f"Unknown label: {label}")

        if 0 <= label < len(self.id2label):
            return label
        raise ValueError(f"Unknown label id: {label}")