Flex dm

class FlexDmConfig(PretrainedConfig):
    """Configuration for a converted Flex-DM MFP model.

    Args:
        dataset_name: Released dataset name.
        checkpoint_variant: Released checkpoint variant name.
        id2label: Public dataset-local label mapping.
        input_columns: Heterogeneous model column specs.
        attribute_groups: Model feature groups used for infilling.
        max_seq_length: Maximum document elements.
        latent_dim: Transformer hidden dimension.
        num_blocks: Number of DeepSVG-style transformer blocks.
        block_type: Released block type. Only ``deepsvg`` is implemented.
        masking_method: Released masking task selector.
        seq_type: Released sequence model type. ``default`` is the released path.
        arch_type: Released architecture type. ``oneshot`` is the released path.
        context: Optional reference context embedding mode.
        input_dtype: Released input ordering mode.
        use_elemwise_noise: Whether element-wise noise was enabled.
        dropout: Dropout probability.
        layer_norm_epsilon: LayerNorm epsilon matching Keras defaults.
        l2: Original L2 setting, stored for provenance.
        original_args: Raw reference ``args.json`` values.
        conversion_report: Checkpoint conversion diagnostics.
        kwargs: Extra ``PretrainedConfig`` fields.
    """

    model_type = "flex-dm"

    def __init__(
        self,
        dataset_name: str = "crello",
        checkpoint_variant: str = "ours-exp-ft",
        id2label: dict[int | str, str] | None = None,
        input_columns: dict[str, FlexDmColumnSpec] | None = None,
        attribute_groups: dict[str, tuple[str, ...] | list[str]] | None = None,
        max_seq_length: int = 50,
        latent_dim: int = 256,
        num_blocks: int = 4,
        block_type: str = "deepsvg",
        masking_method: str = "random",
        seq_type: str = "default",
        arch_type: str = "oneshot",
        context: str | None = None,
        input_dtype: str = "set",
        use_elemwise_noise: bool = False,
        dropout: float = 0.1,
        layer_norm_epsilon: float = 1e-3,
        l2: float | None = 1e-2,
        original_args: dict[str, FlexDmConfigValue] | None = None,
        conversion_report: dict[str, FlexDmConfigValue] | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize a Flex-DM config."""
        normalized_id2label = _normalize_id2label(id2label)
        kwargs.pop("label2id", None)
        super().__init__(
            id2label=normalized_id2label,
            label2id={label: idx for idx, label in normalized_id2label.items()},
            **kwargs,  # ty: ignore[invalid-argument-type]
        )
        self.dataset_name = dataset_name
        self.checkpoint_variant = checkpoint_variant
        self.input_columns = _normalize_columns(input_columns)
        if attribute_groups is None:
            from .data_specs import attribute_groups_for_dataset

            attribute_groups = dict(attribute_groups_for_dataset(dataset_name))
        groups = attribute_groups

        self.attribute_groups = {key: tuple(value) for key, value in groups.items()}
        self.max_seq_length = max_seq_length
        self.latent_dim = latent_dim
        self.num_blocks = num_blocks
        self.block_type = block_type
        self.masking_method = masking_method
        self.seq_type = seq_type
        self.arch_type = arch_type
        self.context = context
        self.input_dtype = input_dtype

        self.use_elemwise_noise = use_elemwise_noise
        self.dropout = dropout
        self.layer_norm_epsilon = layer_norm_epsilon
        self.l2 = l2
        self.original_args = original_args or {}
        self.conversion_report = conversion_report or {}

    @property
    def max_seq_length_with_length_lookup(self) -> int:
        """Return the max length used by the zero-based length lookup."""
        return self.max_seq_length

    @property
    def valid_sequence_keys(self) -> tuple[str, ...]:
        """Return non-demo sequence fields modeled by Flex-DM."""
        return tuple(
            key for key, column in self.input_columns.items() if column["is_sequence"]
        )

    @property
    def categorical_keys(self) -> tuple[str, ...]:
        """Return sequence fields with categorical heads."""
        return tuple(
            key
            for key, column in self.input_columns.items()
            if column["is_sequence"] and column["type"] == "categorical"
        )

    @property
    def numerical_keys(self) -> tuple[str, ...]:
        """Return sequence fields with numerical heads."""
        return tuple(
            key
            for key, column in self.input_columns.items()
            if column["is_sequence"] and column["type"] == "numerical"
        )

    @property
    def task_names(self) -> tuple[str, ...]:
        """Return task names in sampler order."""
        return ("random", "elem", *self.attribute_groups.keys())

    def mask_token_id_for(self, key: str) -> int:
        """Return the categorical mask token id for ``key``."""
        input_dim = self.input_columns[key]["input_dim"]
        if input_dim is None:
            raise ValueError(f"{key} is not categorical")

        return input_dim

    def unused_token_id_for(self, key: str) -> int:
        """Return the categorical unused token id for ``key``."""
        return self.mask_token_id_for(key) + 1

class FlexDmForMaskedDocumentModeling(FlexDmPreTrainedModel):
    """Flex-DM MFP model with a standard Transformers ``forward`` method."""

    def __init__(self, config: FlexDmConfig) -> None:
        """Initialize encoder, transformer blocks, and decoder."""
        super().__init__(config)
        if config.arch_type != "oneshot":
            raise ValueError("Only arch_type='oneshot' is supported")

        if config.block_type != "deepsvg":
            raise ValueError("Only block_type='deepsvg' is supported")

        self.encoder = FlexDmInputEncoder(config)
        self.blocks = nn.ModuleList(
            [FlexDmDeepSvgBlock(config) for _ in range(config.num_blocks)]
        )
        self.decoder = FlexDmDecoder(config)
        self.post_init()

    def forward(
        self,
        *,
        inputs: Mapping[str, Shaped[torch.Tensor, "..."]],
        masks: Mapping[str, Bool[torch.Tensor, "..."]] | None = None,
        labels: Mapping[str, Shaped[torch.Tensor, "..."]] | None = None,
        task_ids: Int[torch.Tensor, "batch"] | None = None,
        output_hidden_states: bool = False,
        return_dict: bool | None = None,
    ) -> (
        FlexDmModelOutput
        | tuple[dict[str, Shaped[torch.Tensor, "..."]], Float[torch.Tensor, ""] | None]
    ):
        """Run a Flex-DM forward pass.

        Args:
            inputs: Per-column model input tensors.
            masks: Optional hidden-field masks for diagnostics.
            labels: Optional per-column reconstruction targets.
            task_ids: Optional task ids.
            output_hidden_states: Whether to include final hidden states.
            return_dict: Whether to return a ``ModelOutput``.

        Returns:
            ``FlexDmModelOutput`` by default.
        """
        hidden_states, seq_mask = self.encoder(inputs, task_ids=task_ids)
        for block in self.blocks:
            hidden_states = block(hidden_states, seq_mask)
        logits = self.decoder(hidden_states)
        loss = self._compute_loss(logits, labels) if labels is not None else None
        output = FlexDmModelOutput(
            logits=logits,
            loss=loss,
            hidden_states=hidden_states if output_hidden_states else None,
            masks=dict(masks) if masks is not None else None,
        )
        if return_dict is False:
            return logits, loss
        return output

    def _compute_loss(
        self,
        logits: Mapping[str, Shaped[torch.Tensor, "..."]],
        labels: Mapping[str, Shaped[torch.Tensor, "..."]],
    ) -> Float[torch.Tensor, ""]:
        losses: list[Float[torch.Tensor, ""]] = []
        for key, target in labels.items():
            if key not in logits:
                continue
            column: FlexDmColumnSpec = self.config.input_columns[key]
            pred = logits[key]
            if column["type"] == "categorical":
                vocab = cast(int, column["input_dim"])
                losses.append(
                    F.cross_entropy(pred.view(-1, vocab), target.long().view(-1))
                )
            else:
                losses.append(F.mse_loss(pred, target.float()))
        if not losses:
            return torch.tensor(0.0, device=next(self.parameters()).device)
        return torch.stack(losses).sum()

@dataclass
class FlexDmModelOutput(ModelOutput):
    """Output of ``FlexDmForMaskedDocumentModeling``.

    Args:
        logits: Per-column logits or numerical predictions.
        loss: Optional summed reconstruction loss.
        hidden_states: Optional final hidden states.
        masks: Optional per-column hidden-field masks.
    """

    logits: dict[str, Shaped[torch.Tensor, "..."]]
    loss: Float[torch.Tensor, ""] | None = None
    hidden_states: Float[torch.Tensor, "batch seq channels"] | None = None
    masks: dict[str, Bool[torch.Tensor, "..."]] | None = None

    def __post_init__(self) -> None:
        """Keep the logits dictionary as one ModelOutput field."""
        if self.logits is not None:
            self["logits"] = self.logits
        if self.loss is not None:
            self["loss"] = self.loss
        if self.hidden_states is not None:
            self["hidden_states"] = self.hidden_states
        if self.masks is not None:
            self["masks"] = self.masks

class FlexDmPipeline(LayoutGenerationPipeline):
    """Run Flex-DM completion, refinement, and feature-level content infilling."""

    config_class: ClassVar[type[PretrainedConfig]] = FlexDmConfig
    component_specs: ClassVar = model_processor_component_specs(
        model_loader=_load_model_component,
        processor_loader=_load_processor_component,
    )

    config: FlexDmConfig
    model: FlexDmForMaskedDocumentModeling
    processor: FlexDmProcessor

    def __init__(
        self,
        model: FlexDmForMaskedDocumentModeling,
        processor: FlexDmProcessor | None = None,
        config: FlexDmConfig | None = None,
    ) -> None:
        """Initialize model and processor components."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor or FlexDmProcessor.from_config(self.config)

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, FlexDmPipelineComponent | None],
    ) -> "FlexDmPipeline":
        """Build a pipeline from loaded model and processor components."""
        return cls(
            config=cast(FlexDmConfig, config),
            model=cast(FlexDmForMaskedDocumentModeling, components["model"]),
            processor=cast(FlexDmProcessor, components["processor"]),
        )

    @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.completion,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "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: Literal["dataclass", "dict"] = "dataclass",
        return_intermediates: bool = False,
        attributes: Mapping[str, FlexDmValue] | None = None,
        content: Mapping[str, FlexDmValue] | None = None,
        feature_group: str | None = None,
        target_indices: Int[torch.Tensor, "..."] | None = None,
        **model_kwargs: FlexDmValue,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | Mapping[str, Shaped[torch.Tensor, "..."]]
            | None,
        ]
    ):
        """Infills masked Flex-DM document fields.

        Args:
            batch_size: Batch size used when synthetic empty inputs are created.
            seed: Common API compatibility argument. Flex-DM's public inference
                path is deterministic and does not currently consume randomness.
            generator: Common API compatibility argument. When supplied, it
                takes precedence over ``seed``; the deterministic Flex-DM
                inference path does not currently consume it.
            condition_type: Canonical condition or local task alias.
            labels: Public element labels.
            bbox: Public element boxes.
            mask: Public valid-element mask.
            num_elements: Optional element counts for synthetic inputs.
            box_format: Input box coordinate format.
            normalized: Whether input boxes are already normalized.
            canvas_size: Pixel canvas size when ``normalized=False``.
            num_inference_steps: Number of iterative decode steps.
            output_type: ``dataclass`` or ``dict``.
            return_intermediates: Whether to include logits and masks.
            attributes: Optional non-core document attributes.
            content: Optional Crello image/text embeddings.
            feature_group: Flex-DM task group such as ``pos`` or ``img``.
            target_indices: Optional element indexes for ``elem`` masking.
            model_kwargs: Reserved model keyword arguments.

        Returns:
            Common layout-generation output.

        Raises:
            NotImplementedError: If the requested canonical condition is not
                supported by released Flex-DM MFP checkpoints.
        """
        _ = model_kwargs
        model_device = next(self.model.parameters()).device
        if generator is None and seed is not None:
            generator = torch.Generator(device=model_device).manual_seed(seed)
        encoded = self.processor(
            condition_type=condition_type,
            labels=labels,
            bbox=bbox,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            attributes=attributes,
            content=content,
            feature_group=feature_group,
            target_indices=target_indices,
            batch_size=batch_size,
        )
        canonical = cast(ConditionType, encoded["condition_type"])
        feature = cast(str | None, encoded["feature_group"])
        self._validate_condition(canonical, feature)
        inputs = {
            key: value.to(model_device)
            for key, value in cast(
                dict[str, Shaped[torch.Tensor, "..."]], encoded["inputs"]
            ).items()
        }
        masks = {
            key: value.to(model_device)
            for key, value in cast(
                dict[str, Bool[torch.Tensor, "..."]], encoded["masks"]
            ).items()
        }
        masked_inputs = self._apply_masks(inputs, masks, generator=generator)
        was_training = self.model.training
        self.model.eval()
        try:
            if num_inference_steps is not None and num_inference_steps > 1:
                outputs = cast(
                    FlexDmModelOutput,
                    iterative_decode(
                        self.model,
                        inputs=masked_inputs,
                        masks=masks,
                        num_iter=num_inference_steps,
                        input_columns=self.config.input_columns,
                        source_inputs=inputs,
                    ),
                )
            else:
                outputs = self.model(
                    inputs=masked_inputs, masks=masks, return_dict=True
                )
        finally:
            self.model.train(was_training)
        return self.processor.post_process_document(
            outputs,
            original_inputs=inputs,
            masks=masks,
            output_type=output_type,
            return_intermediates=return_intermediates,
            refinement_input=inputs if canonical is ConditionType.refinement else None,
        )

    generate = __call__

    def _apply_masks(
        self,
        inputs: Mapping[str, Shaped[torch.Tensor, "..."]],
        masks: Mapping[str, Bool[torch.Tensor, "..."]],
        *,
        generator: torch.Generator | None,
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        modified = dict(inputs)
        for key, mask in masks.items():
            if key not in self.config.input_columns:
                continue
            column = self.config.input_columns[key]
            if column["is_sequence"] and mask.ndim == 2 and mask.any():
                modified[key] = apply_token(
                    modified[key],
                    column,
                    mask,
                    "masked",
                    generator=generator,
                )
        return modified

    def _validate_condition(
        self,
        condition_type: ConditionType,
        feature_group: str | None,
    ) -> None:
        if condition_type in {ConditionType.completion, ConditionType.refinement}:
            return
        if condition_type is ConditionType.content_image and feature_group in {
            "img",
            "txt",
        }:
            return
        if condition_type is ConditionType.label:
            raise NotImplementedError(
                "Flex-DM has no standalone label-conditioned mode; use "
                'condition_type="completion", feature_group="type".'
            )

        if condition_type is ConditionType.label_size:
            raise NotImplementedError("Flex-DM does not support label_size generation")

        if condition_type is ConditionType.unconditional:
            raise NotImplementedError(
                "Flex-DM released MFP checkpoints require an input document"
            )

        raise NotImplementedError(
            f"Flex-DM does not support condition_type={condition_type}"
        )

class FlexDmProcessor(ProcessorMixin):
    """Serialize vocabularies and convert public layouts to Flex-DM tensors.

    Flex-DM intentionally does not expose a ``PreTrainedTokenizer`` because the
    model consumes a dictionary of heterogeneous categorical and continuous
    fields rather than one discrete token stream.
    """

    attributes: list[str] = []

    def __init__(
        self,
        *,
        config: FlexDmConfig,
        vocabulary: dict[str, FlexDmValue] | None = None,
        discretizers: dict[str, FlexDmDiscretizerSpec] | None = None,
    ) -> None:
        """Initialize metadata-only processor state."""
        self.config = config
        self.vocabulary = vocabulary or {}
        self.discretizers = discretizers or {
            key: {"min": 0.0, "max": 1.0, "bins": 64} for key in GEOMETRY_KEYS
        }
        if "opacity" in config.input_columns:
            self.discretizers.setdefault("opacity", {"min": 0.0, "max": 1.0, "bins": 8})
        if "color" in config.input_columns:
            self.discretizers.setdefault(
                "color", {"min": 0.0, "max": 255.0, "bins": 16}
            )

    @classmethod
    def from_config(cls, config: FlexDmConfig) -> "FlexDmProcessor":
        """Create a processor from config metadata.

        Args:
            config: Flex-DM configuration.

        Returns:
            Processor with built-in discretizers.
        """
        return cls(config=config)

    def save_pretrained(
        self,
        save_directory: str | PathLike[str],
        push_to_hub: bool = False,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Save processor metadata next to a converted checkpoint."""
        _ = (push_to_hub, kwargs)
        root = Path(save_directory)
        root.mkdir(parents=True, exist_ok=True)
        (root / "processor_config.json").write_text(
            json.dumps(
                {
                    "processor_class": self.__class__.__name__,
                    "config": self.config.to_dict(),
                    "vocabulary": self.vocabulary,
                    "discretizers": self.discretizers,
                    "tokenizer_policy_deviation": (
                        "Flex-DM uses ProcessorMixin instead of PreTrainedTokenizer "
                        "because the model consumes heterogeneous dict tensors and "
                        "continuous image/text embeddings."
                    ),
                },
                indent=2,
                sort_keys=True,
            )
        )

    @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",
        *,
        subfolder: str | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> "FlexDmProcessor":
        """Load processor metadata from a local converted checkpoint."""
        _ = (cache_dir, force_download, local_files_only, token, revision, kwargs)
        root = Path(pretrained_model_name_or_path)
        if subfolder is not None:
            root = root / subfolder
        data = json.loads((root / "processor_config.json").read_text())
        return cls(
            config=FlexDmConfig.from_dict(data["config"]),
            vocabulary=cast(dict[str, FlexDmValue], data.get("vocabulary", {})),
            discretizers=cast(
                dict[str, FlexDmDiscretizerSpec], data.get("discretizers", {})
            ),
        )

    @classmethod
    def from_vocabulary(
        cls,
        *,
        dataset_name: str,
        vocabulary: dict[str, FlexDmValue],
        checkpoint_variant: str = "ours-exp-ft",
    ) -> "FlexDmProcessor":
        """Build config and processor metadata from vocabulary."""
        id2label = cast(
            dict[int | str, str],
            id2label_from_vocabulary(
                dataset_name, cast(dict[str, FlexDmVocabularyValue], vocabulary)
            ),
        )
        input_columns = build_column_specs(
            dataset_name=dataset_name,
            vocabulary=cast(dict[str, FlexDmVocabularyValue], vocabulary),
        )
        config = FlexDmConfig(
            dataset_name=dataset_name,
            checkpoint_variant=checkpoint_variant,
            id2label=id2label,
            input_columns=input_columns,
        )
        return cls(config=config, vocabulary=vocabulary)

    def __call__(
        self,
        *,
        condition_type: ConditionType | str = ConditionType.completion,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "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,
        attributes: Mapping[str, FlexDmValue] | None = None,
        content: Mapping[str, FlexDmValue] | None = None,
        feature_group: str | None = None,
        target_indices: Int[torch.Tensor, "..."] | None = None,
        batch_size: int = 1,
        return_tensors: Literal["pt"] = "pt",
    ) -> dict[
        str,
        dict[str, Shaped[torch.Tensor, "..."] | Bool[torch.Tensor, "..."]]
        | Shaped[torch.Tensor, "..."]
        | ConditionType
        | str
        | None,
    ]:
        """Convert public layout fields into Flex-DM model tensors."""
        if return_tensors != "pt":
            raise ValueError("FlexDmProcessor only supports return_tensors='pt'")

        if bbox is None or labels is None:
            count = self._num_elements_tensor(num_elements, batch_size)
            max_len = int(count.max().item()) if count.numel() else 0
            bbox_t = torch.zeros((batch_size, max_len, 4), dtype=torch.float32)
            labels_t = torch.zeros((batch_size, max_len), dtype=torch.long)
            mask_t = torch.arange(max_len).unsqueeze(0) < count.unsqueeze(1)
        else:
            bbox_t, labels_t, mask_t = prepare_layout_tensors(
                bbox=bbox,
                labels=labels,
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
                clamp_converted_normalized=True,
            )
        inputs = self._layout_to_inputs(
            bbox=bbox_t,
            labels=labels_t,
            mask=mask_t,
            attributes=attributes,
            content=content,
        )
        length = mask_t.long().sum(dim=1).clamp(min=1) - 1
        inputs["length"] = length.reshape(-1, 1).long()
        seq_mask = get_seq_mask(inputs["length"].reshape(-1), maxlen=bbox_t.size(1))
        filtered = filter_padding(inputs, self.config.input_columns, seq_mask)
        canonical, normalized_feature = self.normalize_condition_and_feature(
            condition_type,
            feature_group=feature_group,
        )
        masks = build_feature_masks(
            self.config.input_columns,
            seq_mask,
            condition_type=canonical,
            feature_group=normalized_feature,
            target_indices=target_indices,
        )
        return {
            "inputs": filtered,
            "masks": masks,
            "bbox": bbox_t,
            "labels": labels_t,
            "mask": mask_t,
            "condition_type": canonical,
            "feature_group": normalized_feature,
        }

    def normalize_condition_and_feature(
        self,
        condition_type: ConditionType | str,
        *,
        feature_group: str | None = None,
    ) -> tuple[ConditionType, str | None]:
        """Normalize canonical conditions plus local Flex-DM task aliases."""
        aliases = {"random", "elem", "type", "pos", "attr", "img", "txt"}
        if isinstance(condition_type, str) and condition_type in aliases:
            return ConditionType.completion, condition_type
        canonical = normalize_condition_type(condition_type)
        if canonical is ConditionType.content_image and feature_group is None:
            return canonical, "img"
        return canonical, feature_group

    def _num_elements_tensor(
        self,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None,
        batch_size: int,
    ) -> Int[torch.Tensor, "batch"]:
        if num_elements is None:
            return torch.full(
                (batch_size,), min(1, self.config.max_seq_length), dtype=torch.long
            )
        tensor = torch.as_tensor(num_elements, dtype=torch.long)
        if tensor.ndim == 0:
            tensor = tensor.repeat(batch_size)
        return tensor

    def _layout_to_inputs(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"],
        attributes: Mapping[str, FlexDmValue] | None,
        content: Mapping[str, FlexDmValue] | None,
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        ltwh = xywh_to_ltwh(bbox).clamp(0.0, 1.0)
        inputs: dict[str, Shaped[torch.Tensor, "..."]] = {}
        for idx, key in enumerate(GEOMETRY_KEYS):
            inputs[key] = self._discretize(key, ltwh[..., idx : idx + 1]).long()
        inputs["type"] = labels.unsqueeze(-1).long()
        attrs = attributes or {}
        cnt = content or {}
        for key, column in self.config.input_columns.items():
            if key in inputs or key == "length":
                continue
            if not column["is_sequence"]:
                inputs[key] = torch.zeros((bbox.size(0), 1), dtype=torch.long)
            else:
                source = cnt.get(key, attrs.get(key))
                inputs[key] = self._coerce_column_value(key, column, source, mask)
        return inputs

    def _coerce_column_value(
        self,
        key: str,
        column: FlexDmColumnSpec,
        value: FlexDmValue,
        mask: Bool[torch.Tensor, "batch elements"],
    ) -> Shaped[torch.Tensor, "batch elements channels"]:
        batch, seq_len = mask.shape
        shape = (batch, seq_len, int(column["shape"][-1]))
        if value is None:
            dtype = torch.float32 if column["type"] == "numerical" else torch.long
            return torch.zeros(shape, dtype=dtype)
        tensor = torch.as_tensor(value)
        if tensor.ndim == 2:
            tensor = tensor.unsqueeze(-1)
        if key in self.discretizers and tensor.dtype.is_floating_point:
            tensor = self._discretize(key, tensor.float())
        return tensor.reshape(shape).to(
            dtype=torch.float32 if column["type"] == "numerical" else torch.long
        )

    def _discretize(
        self, key: str, value: Float[torch.Tensor, "..."]
    ) -> Float[torch.Tensor, "..."]:
        spec = self.discretizers[key]
        scaled = (value - spec["min"]) / (spec["max"] - spec["min"])
        return torch.clamp((scaled * spec["bins"]).floor(), 0, spec["bins"] - 1)

    def _continuize(
        self, key: str, value: Shaped[torch.Tensor, "..."]
    ) -> Float[torch.Tensor, "..."]:
        spec = self.discretizers[key]
        scale = (spec["max"] - spec["min"]) / spec["bins"]
        return value.float() * scale + spec["min"]

    def post_process_document(
        self,
        outputs: FlexDmModelOutput,
        *,
        original_inputs: Mapping[str, Shaped[torch.Tensor, "..."]],
        masks: Mapping[str, Bool[torch.Tensor, "..."]],
        output_type: Literal["dataclass", "dict"] = "dataclass",
        return_intermediates: bool = False,
        refinement_input: Mapping[str, Shaped[torch.Tensor, "..."]] | None = None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Shaped[torch.Tensor, "..."]
            | dict[int, str]
            | Mapping[str, Shaped[torch.Tensor, "..."]]
            | None,
        ]
    ):
        """Decode Flex-DM model outputs to the common layout schema."""
        decoded = self._decode_logits(outputs.logits, original_inputs, masks)
        ltwh = torch.cat([decoded[key].float() for key in GEOMETRY_KEYS], dim=-1)
        bbox = ltwh_to_xywh(ltwh).clamp(0.0, 1.0).detach().cpu()
        labels = decoded["type"].squeeze(-1).long().detach().cpu()
        valid_mask = get_seq_mask(
            original_inputs["length"].reshape(-1), maxlen=labels.size(1)
        )
        intermediates = None
        if return_intermediates:
            intermediates = {
                "attributes": {
                    key: value.detach().cpu()
                    for key, value in decoded.items()
                    if key not in (*GEOMETRY_KEYS, "type")
                    and self.config.input_columns[key]["is_sequence"]
                },
                "masks": {key: value.detach().cpu() for key, value in masks.items()},
                "logits": {
                    key: value.detach().cpu() for key, value in outputs.logits.items()
                },
            }
            if refinement_input is not None:
                intermediates["refinement_input"] = {
                    key: value.detach().cpu() for key, value in refinement_input.items()
                }
        result = LayoutGenerationOutput(
            bbox=bbox,
            labels=labels,
            mask=valid_mask.detach().cpu(),
            id2label=cast(dict[int, str], self.config.id2label),
            intermediates=intermediates,
        )
        if output_type == "dict":
            return dict(result)
        if output_type != "dataclass":
            raise ValueError(f"Unsupported output_type: {output_type}")

        return result

    def _decode_logits(
        self,
        logits: Mapping[str, Shaped[torch.Tensor, "..."]],
        original_inputs: Mapping[str, Shaped[torch.Tensor, "..."]],
        masks: Mapping[str, Bool[torch.Tensor, "..."]],
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        decoded: dict[str, Shaped[torch.Tensor, "..."]] = {}
        for key, column in self.config.input_columns.items():
            if not column["is_sequence"]:
                continue
            if key in logits:
                pred = (
                    logits[key].argmax(dim=-1)
                    if column["type"] == "categorical"
                    else logits[key]
                )
            else:
                pred = original_inputs[key]
            mask = masks.get(key)
            if mask is not None:
                pred = torch.where(
                    mask.unsqueeze(-1).to(pred.device),
                    pred,
                    original_inputs[key].to(pred.device),
                )
            if key in self.discretizers and column["type"] == "categorical":
                pred = self._continuize(key, pred)
            decoded[key] = pred
        return decoded