Coarse to fine

class ConditionType(StrEnum):
    """Canonical condition names used by layout generation interfaces."""

    unconditional = auto()
    label = auto()
    label_size = auto()
    completion = auto()
    refinement = auto()
    text = auto()
    content_image = auto()
    relation = auto()
    hierarchical = auto()
    retrieval = auto()

@dataclass
class LayoutGenerationOutput(ModelOutput):
    """Canonical layout-generation output for Transformers-style APIs.

    Attributes:
        bbox: Normalized center ``xywh`` boxes with shape
            ``(batch, elements, 4)``.
        labels: Dataset-local integer labels with shape ``(batch, elements)``.
        mask: Boolean valid-element mask with shape ``(batch, elements)``.
        id2label: Mapping from integer label ids to display names.
        sequences: Optional raw token sequences.
        scores: Optional per-token or per-element scores.
        trajectory: Optional sampling trajectory.
        intermediates: Optional model-specific debug or auxiliary data.

    Examples:
        >>> import numpy as np
        >>> output = LayoutGenerationOutput(
        ...     bbox=np.zeros((1, 1, 4), dtype=np.float32),
        ...     labels=np.zeros((1, 1), dtype=np.int64),
        ...     mask=np.ones((1, 1), dtype=bool),
        ...     id2label={0: "text"},
        ... )
        >>> output["bbox"].shape
        (1, 1, 4)
    """

    bbox: (
        Float[np.ndarray, "batch elements 4"] | Float[torch.Tensor, "batch elements 4"]
    )
    labels: Int[np.ndarray, "batch elements"] | Int[torch.Tensor, "batch elements"] = (
        cast(
            'Int[np.ndarray, "batch elements"] | Int[torch.Tensor, "batch elements"]',
            None,
        )
    )
    mask: Bool[np.ndarray, "batch elements"] | Bool[torch.Tensor, "batch elements"] = (
        cast(
            'Bool[np.ndarray, "batch elements"] | Bool[torch.Tensor, "batch elements"]',
            None,
        )
    )
    id2label: dict[int, str] = cast(dict[int, str], None)
    sequences: object | None = None
    scores: object | None = None
    trajectory: object | None = None
    intermediates: object | None = None

class CoarseToFineConfig(PretrainedConfig):
    """Stores architecture, discretization, and label metadata.

    Args:
        dataset: Dataset name for the converted checkpoint.
        num_labels: Optional label vocabulary size. Defaults to the dataset
            vocabulary length.
        id2label: Optional label id to display-label mapping.
        label2id: Optional display-label to id mapping.
        max_num_elements: Maximum flat element count used by the checkpoint model.
        discrete_x_grid: Number of x-axis bins.
        discrete_y_grid: Number of y-axis bins.
        d_model: Transformer hidden dimension.
        d_z: VAE latent dimension.
        n_layers: Number of encoder layers.
        n_layers_decoder: Number of decoder layers.
        n_heads: Number of attention heads.
        dim_feedforward: Transformer feed-forward dimension.
        dropout: Dropout probability.
        internal_box_format: Reference internal box format.
        public_box_format: Public output box format.
        vendor_label_offset: Offset from public label ids to internal ids.
        **kwargs: Additional ``PretrainedConfig`` fields.

    Examples:
        >>> CoarseToFineConfig(dataset="publaynet").num_labels
        5
    """

    model_type = "coarse_to_fine"

    def __init__(
        self,
        dataset: DatasetName | str = DatasetName.rico25,
        num_labels: int | None = None,
        id2label: Mapping[int | str, str] | None = None,
        label2id: Mapping[str, int] | None = None,
        max_num_elements: int = 20,
        discrete_x_grid: int = 128,
        discrete_y_grid: int = 128,
        d_model: int = 512,
        d_z: int = 512,
        n_layers: int = 4,
        n_layers_decoder: int = 4,
        n_heads: int = 8,
        dim_feedforward: int = 2048,
        dropout: float = 0.1,
        internal_box_format: BoxFormat | str = BoxFormat.ltwh,
        public_box_format: BoxFormat | str = BoxFormat.xywh,
        vendor_label_offset: int = 1,
        eval_batch_size: int | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize Coarse-to-Fine config values."""
        normalized_dataset = normalize_dataset_name(dataset)
        if normalized_dataset not in SUPPORTED_DATASETS:
            raise ValueError(f"Unsupported Coarse-to-Fine dataset: {dataset}")

        resolved_id2label = (
            {int(key): str(value) for key, value in id2label.items()}
            if id2label is not None
            else id2label_for_dataset(normalized_dataset)
        )
        resolved_label2id = (
            {str(key): int(value) for key, value in label2id.items()}
            if label2id is not None
            else label2id_for_dataset(normalized_dataset)
        )
        resolved_num_labels = num_labels or len(resolved_id2label)
        super().__init__(id2label=resolved_id2label, label2id=resolved_label2id)
        for key, value in kwargs.items():
            setattr(self, key, value)

        self.dataset = str(normalized_dataset)
        self.num_labels = resolved_num_labels
        self.max_num_elements = max_num_elements
        self.discrete_x_grid = discrete_x_grid
        self.discrete_y_grid = discrete_y_grid
        self.d_model = d_model
        self.d_z = d_z
        self.n_layers = n_layers
        self.n_layers_decoder = n_layers_decoder
        self.n_heads = n_heads
        self.dim_feedforward = dim_feedforward
        self.dropout = dropout

        self.internal_box_format = str(normalize_box_format(internal_box_format))
        self.public_box_format = str(normalize_box_format(public_box_format))
        self.vendor_label_offset = vendor_label_offset
        self.eval_batch_size = eval_batch_size or max_num_elements

        self.element_sos_id = resolved_num_labels + 1
        self.element_eos_id = resolved_num_labels + 2
        self.group_sos_index = 0
        self.group_eos_index = resolved_num_labels + 1
        self.group_label_size = resolved_num_labels + 2
        self.element_label_size = resolved_num_labels + 3
        self.bbox_vocab_size = max(discrete_x_grid, discrete_y_grid)
        self.architectures = ["CoarseToFineForLayoutGeneration"]

@dataclass
class CoarseToFineHierarchy:
    """Decoded Coarse-to-Fine hierarchy returned in ``intermediates``."""

    group_bbox: Float[torch.Tensor, "batch groups 4"]
    group_mask: Bool[torch.Tensor, "batch groups"]
    label_histogram: Float[torch.Tensor, "batch groups labels"]
    element_group_index: Int[torch.Tensor, "batch elements"]
    relative_bbox: Float[torch.Tensor, "batch groups elements 4"]
    relative_mask: Bool[torch.Tensor, "batch groups elements"]
    discrete_group_bbox: Int[torch.Tensor, "batch groups 4"] | None = None
    discrete_relative_bbox: Int[torch.Tensor, "batch groups elements 4"] | None = None

@dataclass
class CoarseToFineHierarchyEncoding:
    """Training/reference hierarchy tensors before padding."""

    group_bounding_box: Int[torch.Tensor, "groups 4"]
    label_in_one_group: Float[torch.Tensor, "groups labels"]
    grouped_labels: list[Int[torch.Tensor, "elements"]]
    grouped_bbox: list[Int[torch.Tensor, "elements 4"]]

class CoarseToFineForLayoutGeneration(PreTrainedModel):
    """Transformers ``PreTrainedModel`` with checkpoint-compatible module names."""

    config_class = CoarseToFineConfig
    base_model_prefix = "coarse_to_fine"
    main_input_name = "labels"
    _tied_weights_keys = {
        "encoder.embedding.label_embed.weight": "layout_embd.label_embed.weight",
        "group_decoder.layout_embd.label_embed.weight": "layout_embd.label_embed.weight",
        "ele_decoder.layout_embd.label_embed.weight": "layout_embd.label_embed.weight",
        "encoder.embedding.bbox_embed.weight": "layout_embd.bbox_embed.weight",
        "group_decoder.layout_embd.bbox_embed.weight": "layout_embd.bbox_embed.weight",
        "ele_decoder.layout_embd.bbox_embed.weight": "layout_embd.bbox_embed.weight",
        "encoder.embedding.proj_cat.weight": "layout_embd.proj_cat.weight",
        "group_decoder.layout_embd.proj_cat.weight": "layout_embd.proj_cat.weight",
        "ele_decoder.layout_embd.proj_cat.weight": "layout_embd.proj_cat.weight",
        "encoder.embedding.proj_cat.bias": "layout_embd.proj_cat.bias",
        "group_decoder.layout_embd.proj_cat.bias": "layout_embd.proj_cat.bias",
        "ele_decoder.layout_embd.proj_cat.bias": "layout_embd.proj_cat.bias",
        "encoder.embedding.group_label_embed.weight": (
            "layout_embd.group_label_embed.weight"
        ),
        "group_decoder.layout_embd.group_label_embed.weight": (
            "layout_embd.group_label_embed.weight"
        ),
        "ele_decoder.layout_embd.group_label_embed.weight": (
            "layout_embd.group_label_embed.weight"
        ),
        "encoder.embedding.group_label_embed.bias": (
            "layout_embd.group_label_embed.bias"
        ),
        "group_decoder.layout_embd.group_label_embed.bias": (
            "layout_embd.group_label_embed.bias"
        ),
        "ele_decoder.layout_embd.group_label_embed.bias": (
            "layout_embd.group_label_embed.bias"
        ),
    }

    def __init__(self, config: CoarseToFineConfig) -> None:
        """Initialize checkpoint-compatible modules."""
        super().__init__(config)
        self.layout_embd = LayoutEmbedding(config)
        self.encoder = Encoder(config, self.layout_embd)
        self.vae = VAE(config)
        self.group_decoder = GroupDecoder(config, self.layout_embd)
        self.ele_decoder = ElementDecoder(config, self.layout_embd)
        self.all_tied_weights_keys = dict(self._tied_weights_keys)

    @property
    def device(self) -> torch.device:
        """Return the current parameter device."""
        return next(self.parameters()).device

    def _sample_latent(
        self,
        *,
        batch_size: int,
        generator: torch.Generator | None,
        device: torch.device,
    ) -> Float[torch.Tensor, "1 batch latent"]:
        sample_device = generator.device if generator is not None else device
        return cast(
            torch.FloatTensor,
            torch.randn(
                (1, batch_size, self.config.d_z),
                generator=generator,
                device=sample_device,
            ).to(device),
        )

    def forward(
        self,
        labels: Int[torch.Tensor, "batch elements"],
        bbox: Int[torch.Tensor, "batch elements 4"],
        mask: Bool[torch.Tensor, "batch elements"],
        group_bounding_box: Int[torch.Tensor, "batch seq 4"],
        label_in_one_group: Float[torch.Tensor, "batch seq vocab"],
        group_mask: Bool[torch.Tensor, "batch seq"],
        grouped_bbox: Int[torch.Tensor, "batch seq elements 4"],
        grouped_labels: Int[torch.Tensor, "batch seq elements"],
        grouped_mask: Bool[torch.Tensor, "batch seq elements"],
        latent_z: Float[torch.Tensor, "1 batch latent"] | None = None,
        use_teacher_forcing: bool = True,
        return_dict: bool | None = None,
    ) -> (
        dict[str, Shaped[torch.Tensor, "..."]] | tuple[Shaped[torch.Tensor, "..."], ...]
    ):
        """Run teacher-forced or greedy hierarchical decoding.

        Args:
            labels: Batch-first internal label ids.
            bbox: Batch-first discrete ``ltwh`` ids.
            mask: Batch-first valid element mask.
            group_bounding_box: Batch-first group discrete ``ltwh`` ids.
            label_in_one_group: Batch-first group label histograms.
            group_mask: Batch-first valid group mask.
            grouped_bbox: Batch-first group-relative discrete ``ltwh`` ids.
            grouped_labels: Batch-first group-relative internal labels.
            grouped_mask: Batch-first valid grouped-element mask.
            latent_z: Optional latent tensor to bypass stochastic sampling.
            use_teacher_forcing: Whether to use provided hierarchy tensors.
            return_dict: Return a dictionary when true.

        Returns:
            Dictionary of raw logits and latent tensors.
        """
        batch_size, groups, seq, dims = grouped_bbox.shape
        seq_bbox = make_seq_first(bbox)
        seq_group_bbox = make_seq_first(group_bounding_box)
        grouped_box = make_seq_first(
            grouped_bbox.reshape(batch_size, groups * seq, dims)
        )
        grouped_box = make_seq_first(
            grouped_box.reshape(groups, seq, batch_size * dims)
        ).reshape(seq, groups, batch_size, dims)
        seq_labels = make_seq_first(labels.unsqueeze(2)).squeeze(2)
        seq_group_labels = make_seq_first(label_in_one_group)
        grouped_label = make_seq_first(
            grouped_labels.reshape(batch_size, groups * seq, 1)
        )
        grouped_label = make_seq_first(
            grouped_label.reshape(groups, seq, batch_size)
        ).unsqueeze(3)
        memory = self.encoder(seq_labels, seq_bbox, mask)
        if latent_z is None:
            z, mu, logvar = self.vae(memory)
        else:
            z = (
                latent_z
                if latent_z.dim() == 3 and latent_z.size(0) == 1
                else make_seq_first(latent_z)
            )
            mu = torch.zeros_like(z)
            logvar = torch.zeros_like(z)
        if use_teacher_forcing:
            group_embd, rec_group_bbox, rec_group_label = self.group_decoder(
                seq_group_labels, seq_group_bbox, z, group_mask
            )
            rec_box, rec_label = self.ele_decoder(
                grouped_label,
                grouped_box,
                group_embd,
                z,
                grouped_mask,
            )
        else:
            group_embd, rec_group_bbox, rec_group_label = self.group_decoder.inference(
                z, max_group_num=seq_group_labels.shape[0], device=self.device
            )
            rec_box, rec_label = self.ele_decoder.inference(
                group_embd,
                z,
                max_num_elements=grouped_label.shape[0],
                device=self.device,
            )
        rec_box = make_group_first(rec_box)
        rec_label = make_group_first(rec_label)
        flat_box = make_batch_first(rec_box.reshape(groups * seq, batch_size, -1))
        flat_label = make_batch_first(rec_label.reshape(groups * seq, batch_size, -1))
        rec_group_bbox = make_batch_first(rec_group_bbox)
        rec_group_label = make_batch_first(rec_group_label)
        output = {
            "group_bounding_box_logits": rec_group_bbox.reshape(
                batch_size, rec_group_bbox.size(1), 4, self.config.bbox_vocab_size
            ),
            "label_in_one_group_logits": rec_group_label,
            "grouped_bbox_logits": flat_box.reshape(
                batch_size, groups, seq, 4, self.config.bbox_vocab_size
            ),
            "grouped_label_logits": flat_label.reshape(batch_size, groups, seq, -1),
            "mu": make_batch_first(mu),
            "logvar": make_batch_first(logvar),
            "latent_z": make_batch_first(z),
        }
        if return_dict is False:
            return tuple(output.values())
        return output

    @torch.no_grad()
    def _decode_hierarchy(
        self, latent_z: Float[torch.Tensor, "1 batch latent"]
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        """Decode raw hierarchy logits from a seq-first latent tensor."""
        z = latent_z.to(self.device)
        group_embd, group_bbox, group_label = self.group_decoder.inference(
            z,
            max_group_num=self.config.max_num_elements + 2,
            device=self.device,
        )
        elem_bbox, elem_label = self.ele_decoder.inference(
            group_embd,
            z,
            max_num_elements=self.config.max_num_elements + 2,
            device=self.device,
        )
        elem_bbox = make_group_first(elem_bbox)
        elem_label = make_group_first(elem_label)
        groups, seq, batch_size, _ = elem_bbox.shape
        return {
            "group_bounding_box_logits": make_batch_first(group_bbox).reshape(
                batch_size,
                group_bbox.size(0),
                4,
                self.config.bbox_vocab_size,
            ),
            "label_in_one_group_logits": make_batch_first(group_label),
            "grouped_bbox_logits": make_batch_first(
                elem_bbox.reshape(groups * seq, batch_size, -1)
            ).reshape(batch_size, groups, seq, 4, self.config.bbox_vocab_size),
            "grouped_label_logits": make_batch_first(
                elem_label.reshape(groups * seq, batch_size, -1)
            ).reshape(batch_size, groups, seq, -1),
        }

class CoarseToFinePipeline(Pipeline):
    """Orchestrate Coarse-to-Fine layout generation."""

    model: CoarseToFineForLayoutGeneration
    processor: CoarseToFineProcessor

    def __init__(
        self,
        model: CoarseToFineForLayoutGeneration,
        processor: CoarseToFineProcessor,
    ) -> None:
        """Initialize the pipeline with a model and processor."""
        super().__init__(model=model, tokenizer=None)
        self.processor = processor

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

    def preprocess(
        self,
        input_: CoarseToFinePipelineParameter,
        **preprocess_parameters: dict[str, CoarseToFinePipelineParameter],
    ) -> dict[str, GenericTensor]:
        """Satisfy the abstract pipeline API; direct calls bypass this path."""
        _ = (input_, preprocess_parameters)
        return {}

    def _forward(
        self,
        input_tensors: dict[str, GenericTensor],
        **forward_parameters: dict[str, CoarseToFinePipelineParameter],
    ) -> ModelOutput:
        _ = (input_tensors, forward_parameters)
        raise NotImplementedError("Use CoarseToFinePipeline.__call__ directly")

    def postprocess(
        self,
        model_outputs: LayoutGenerationOutput | CoarseToFineOutputDict,
        **kwargs: dict[str, CoarseToFinePipelineParameter],
    ) -> LayoutGenerationOutput | CoarseToFineOutputDict:
        """Return model outputs without additional formatting."""
        _ = kwargs
        return model_outputs

    def __call__(
        self,
        *,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.unconditional,
        labels: Int[torch.Tensor, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Sequence[ArrayLikeInput]
        | 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,
        latent_z: Float[torch.Tensor, "1 batch latent"] | None = None,
    ) -> LayoutGenerationOutput | CoarseToFineOutputDict:
        """Generate layouts through model decode and processor post-processing."""
        del labels, bbox, mask, num_elements
        del box_format, normalized, canvas_size, num_inference_steps
        condition = normalize_condition_type(condition_type)
        if condition is not ConditionType.unconditional:
            raise NotImplementedError(
                "Coarse-to-Fine released checkpoints support only unconditional generation"
            )

        if generator is None and seed is not None:
            generator = torch.Generator(device=self.model.device).manual_seed(seed)
        if latent_z is None:
            sampled_z = self.model._sample_latent(
                batch_size=batch_size,
                generator=generator,
                device=self.model.device,
            )
        else:
            sampled_z = cast(torch.FloatTensor, latent_z.to(self.model.device))
        raw = self.model._decode_hierarchy(sampled_z)
        hierarchy = decode_hierarchy_from_logits(
            group_bbox_logits=raw["group_bounding_box_logits"],
            group_label_logits=raw["label_in_one_group_logits"],
            grouped_bbox_logits=raw["grouped_bbox_logits"],
            grouped_label_logits=raw["grouped_label_logits"],
            num_labels=self.model.config.num_labels,
            group_eos_index=self.model.config.group_eos_index,
            element_eos_id=self.model.config.element_eos_id,
            discrete_x_grid=self.model.config.discrete_x_grid,
            discrete_y_grid=self.model.config.discrete_y_grid,
        )
        output = cast(
            LayoutGenerationOutput,
            self.processor.post_process_hierarchy(
                hierarchy,
                output_type=OutputType.dataclass,
                return_intermediates=return_intermediates,
            ),
        )
        output.sequences = cast(torch.Tensor, hierarchy.discrete_relative_bbox)
        output.scores = None
        output.trajectory = raw if return_intermediates else None
        if normalize_output_type(output_type) is OutputType.dict:
            return dict(output)
        return output

class CoarseToFineProcessor(ProcessorMixin):
    """Convert public layouts to Coarse-to-Fine hierarchy tensors."""

    attributes: list[str] = []
    processor_class = "CoarseToFineProcessor"

    def __init__(
        self,
        dataset: DatasetName | str = DatasetName.rico25,
        *,
        x_grid: int = 128,
        y_grid: int = 128,
        max_num_elements: int = 20,
        id2label: dict[int, str] | None = None,
    ) -> None:
        """Initialize label maps and discretization settings.

        Args:
            dataset: Canonical layout dataset.
            x_grid: Number of x/width bins.
            y_grid: Number of y/height bins.
            max_num_elements: Padded flat sequence length.
            id2label: Optional public label map.

        Examples:
            >>> CoarseToFineProcessor(dataset="publaynet").id2label[0]
            'text'
        """
        normalized_dataset = normalize_dataset_name(dataset)
        self.dataset = str(normalized_dataset)
        self.x_grid = x_grid
        self.y_grid = y_grid
        self.max_num_elements = max_num_elements
        self.id2label = id2label or id2label_for_dataset(normalized_dataset)
        self.label2id = label2id_for_dataset(normalized_dataset)

    @classmethod
    def from_config(
        cls,
        dataset: DatasetName | str = DatasetName.rico25,
        *,
        x_grid: int = 128,
        y_grid: int = 128,
        max_num_elements: int = 20,
        id2label: dict[int, str] | None = None,
    ) -> "CoarseToFineProcessor":
        """Construct a processor without external files."""
        return cls(
            dataset=dataset,
            x_grid=x_grid,
            y_grid=y_grid,
            max_num_elements=max_num_elements,
            id2label=id2label,
        )

    def to_dict(self) -> dict[str, str | int | dict[str, str]]:
        """Serialize processor state to a JSON-compatible dictionary."""
        return {
            "processor_class": self.processor_class,
            "dataset": self.dataset,
            "x_grid": self.x_grid,
            "y_grid": self.y_grid,
            "max_num_elements": self.max_num_elements,
            "id2label": {str(key): value for key, value in self.id2label.items()},
        }

    def save_pretrained(
        self,
        save_directory: str | PathLike[str],
        push_to_hub: bool = False,
        **kwargs: str | int | float | bool | None,
    ) -> list[str]:
        """Save processor metadata next to a converted checkpoint."""
        _ = (push_to_hub, kwargs)
        path = Path(save_directory)
        path.mkdir(parents=True, exist_ok=True)
        out_file = path / "preprocessor_config.json"
        out_file.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n")
        return [str(out_file)]

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | PathLike[str],
        cache_dir: str | PathLike[str] | None = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: str | bool | None = None,
        revision: str = "main",
        **kwargs: str | int | float | bool | None,
    ) -> "CoarseToFineProcessor":
        """Load processor metadata from a local ``save_pretrained`` directory."""
        _ = (cache_dir, force_download, local_files_only, token, revision, kwargs)
        path = Path(pretrained_model_name_or_path) / "preprocessor_config.json"
        data = json.loads(path.read_text())
        return cls(
            dataset=str(data["dataset"]),
            x_grid=int(data["x_grid"]),
            y_grid=int(data["y_grid"]),
            max_num_elements=int(data["max_num_elements"]),
            id2label={int(key): str(value) for key, value in data["id2label"].items()},
        )

    def _labels_to_vendor(
        self, labels: Int[torch.Tensor, "..."]
    ) -> Int[torch.Tensor, "..."]:
        return cast(torch.LongTensor, labels.long() + 1)

    def __call__(
        self,
        labels: list[list[int | str]]
        | Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | None = None,
        bbox: list[list[list[float]]]
        | Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Build discrete flat tensors from public layout inputs.

        Args:
            labels: Public zero-based labels or label strings.
            bbox: Public normalized boxes.
            mask: Optional valid-element mask.
            box_format: Input box format.
            normalized: Must be true; pixel boxes are dataset-loader work.
            return_tensors: Only ``"pt"`` is supported.

        Returns:
            BatchEncoding with internal labels, discrete boxes, and masks.

        Raises:
            ValueError: If labels or boxes are missing.
        """
        if return_tensors != "pt":
            raise ValueError("Only return_tensors='pt' is supported")

        if not normalized:
            raise ValueError("CoarseToFineProcessor expects normalized boxes")

        if labels is None or bbox is None:
            raise ValueError("labels and bbox are required for processor encoding")

        label_tensor = self._coerce_labels(labels)
        bbox_tensor = torch.as_tensor(bbox, dtype=torch.float32)
        encoded_mask = (
            cast(torch.BoolTensor, torch.ones(label_tensor.shape, dtype=torch.bool))
            if mask is None
            else cast(torch.BoolTensor, torch.as_tensor(mask, dtype=torch.bool))
        )
        ltwh = public_to_ltwh(bbox_tensor, box_format=box_format)
        discrete = discretize_ltwh(ltwh, num_x_grid=self.x_grid, num_y_grid=self.y_grid)
        vendor_labels = self._labels_to_vendor(label_tensor)
        return BatchEncoding(
            {"labels": vendor_labels, "bbox": discrete, "mask": encoded_mask}
        )

    def _coerce_labels(
        self,
        labels: list[list[int | str]]
        | Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"],
    ) -> Int[torch.Tensor, "batch elements"]:
        if isinstance(labels, torch.Tensor | np.ndarray):
            return cast(torch.LongTensor, torch.as_tensor(labels, dtype=torch.long))
        rows: list[list[int]] = []
        for row in labels:
            values: list[int] = []
            for label in row:
                if isinstance(label, int):
                    values.append(label)
                else:
                    values.append(self.label2id[label])
            rows.append(values)
        max_len = max((len(row) for row in rows), default=1)
        padded = [row + [0] * (max_len - len(row)) for row in rows]
        return cast(torch.LongTensor, torch.tensor(padded, dtype=torch.long))

    def build_hierarchy_batch(
        self,
        labels: Int[torch.Tensor, "batch elements"],
        bbox: Float[torch.Tensor, "batch elements 4"],
        mask: Bool[torch.Tensor, "batch elements"],
    ) -> BatchEncoding:
        """Build padded hierarchy tensors for training/reference batches."""
        batch = labels.size(0)
        flat_enc = self(labels=labels, bbox=bbox, mask=mask, return_tensors="pt")
        encodings = []
        for idx in range(batch):
            valid = mask[idx].bool()
            encodings.append(
                build_cut_hierarchy(
                    public_to_ltwh(bbox[idx, valid]),
                    self._labels_to_vendor(labels[idx, valid]),
                    num_labels=len(self.id2label),
                    discrete_x_grid=self.x_grid,
                    discrete_y_grid=self.y_grid,
                )
            )
        max_groups = max(enc.group_bounding_box.size(0) for enc in encodings) + 2
        max_decode_groups = max_groups - 2
        max_group_elems = (
            max(max(group.size(0) for group in enc.grouped_labels) for enc in encodings)
            + 2
        )
        group_boxes = []
        group_labels = []
        group_masks = []
        grouped_boxes = []
        grouped_labels = []
        grouped_masks = []
        for enc in encodings:
            group_count = enc.group_bounding_box.size(0)
            sos_group_box = torch.zeros((1, 4), dtype=torch.long)
            eos_group_box = torch.zeros((1, 4), dtype=torch.long)
            padded_group_box = torch.cat(
                (sos_group_box, enc.group_bounding_box.long(), eos_group_box), dim=0
            )
            padded_group_box = F.pad(
                padded_group_box, (0, 0, 0, max_groups - padded_group_box.size(0))
            )
            hist = torch.zeros(
                (group_count + 2, len(self.id2label) + 2), dtype=torch.float32
            )
            hist[0, 0] = 1.0
            hist[1 : group_count + 1, 1:-1] = enc.label_in_one_group.float()
            hist[group_count + 1, -1] = 1.0
            hist = F.pad(hist, (0, 0, 0, max_groups - hist.size(0)))
            group_mask = torch.arange(max_groups) < group_count + 2
            per_group_boxes = []
            per_group_labels = []
            per_group_masks = []
            for group_idx in range(group_count):
                box = enc.grouped_bbox[group_idx].long()
                label = enc.grouped_labels[group_idx].long()
                label = torch.cat(
                    (
                        torch.tensor([len(self.id2label) + 1]),
                        label,
                        torch.tensor([len(self.id2label) + 2]),
                    )
                )
                box = torch.cat(
                    (
                        torch.zeros((1, 4), dtype=torch.long),
                        box,
                        torch.zeros((1, 4), dtype=torch.long),
                    )
                )
                valid_count = min(label.size(0), max_group_elems)
                per_group_boxes.append(
                    F.pad(
                        box[:max_group_elems], (0, 0, 0, max_group_elems - valid_count)
                    )
                )
                per_group_labels.append(
                    F.pad(label[:max_group_elems], (0, max_group_elems - valid_count))
                )
                per_group_masks.append(torch.arange(max_group_elems) < valid_count)
            while len(per_group_boxes) < max_decode_groups:
                per_group_boxes.append(
                    torch.zeros((max_group_elems, 4), dtype=torch.long)
                )
                per_group_labels.append(
                    torch.zeros((max_group_elems,), dtype=torch.long)
                )
                per_group_masks.append(
                    torch.zeros((max_group_elems,), dtype=torch.bool)
                )
            group_boxes.append(padded_group_box)
            group_labels.append(hist)
            group_masks.append(group_mask)
            grouped_boxes.append(torch.stack(per_group_boxes[:max_decode_groups]))
            grouped_labels.append(torch.stack(per_group_labels[:max_decode_groups]))
            grouped_masks.append(torch.stack(per_group_masks[:max_decode_groups]))
        flat_enc.update(
            {
                "group_bounding_box": torch.stack(group_boxes),
                "label_in_one_group": torch.stack(group_labels),
                "group_mask": torch.stack(group_masks),
                "grouped_bbox": torch.stack(grouped_boxes),
                "grouped_labels": torch.stack(grouped_labels),
                "grouped_mask": torch.stack(grouped_masks),
            }
        )
        return flat_enc

    def post_process_hierarchy(
        self,
        hierarchy: CoarseToFineHierarchy,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | None,
        ]
    ):
        """Convert decoded hierarchy tensors to the shared output schema."""
        output = flatten_hierarchy(
            hierarchy,
            id2label=dict(self.id2label),
            max_num_elements=self.max_num_elements,
        )
        if not return_intermediates:
            output.intermediates = None
        if normalize_output_type(output_type) is OutputType.dict:
            return dict(output)
        return output