Ltnet

class LTNetConfig(PretrainedConfig):
    """Architecture and processor metadata for LT-Net checkpoints.

    Args:
        dataset_name: Dataset slug for the converted checkpoint.
        vocab_size: Mixed special/object/predicate token vocabulary size.
        obj_classes_size: Object-id embedding/classifier vocabulary size.
        hidden_size: Transformer hidden dimension.
        num_hidden_layers: Number of relation encoder layers.
        num_attention_heads: Number of relation encoder attention heads.
        dropout: Dropout used by embeddings, encoder, and bbox heads.
        enable_noise: Whether the original config enabled relation noise.
        noise_size: Original relation-noise size.
        decoder_head_type: ``GMM`` or ``Linear`` bbox head.
        decoder_box_loss: ``PDF`` or ``Reg`` objective family.
        decoder_schedule_sample: Whether training used scheduled sampling.
        decoder_two_path: Whether the original config enabled two-path decoding.
        decoder_global_feature: Whether to concatenate max-pooled global features.
        decoder_greedy: Whether inference uses GMM means instead of sampling.
        xy_temperature: GMM mixture temperature for center coordinates.
        wh_temperature: GMM mixture temperature for box size.
        refine: Whether a refinement head is present.
        refine_head_type: Refinement bbox head type.
        refine_box_loss: Refinement objective family.
        refine_x_softmax: Whether the decoder records XY PDF scores for refine.
        max_sequence_length: Processor/model sequence length.
        id2label: Public object label mapping.
        relation_id2label: Public relation label mapping.
        model_type: Ignored compatibility field from serialized configs.
        transformers_version: Ignored compatibility field.
        kwargs: Additional ``PretrainedConfig`` fields.

    Examples:
        >>> config = LTNetConfig(hidden_size=32, num_attention_heads=4)
        >>> config.model_type
        'ltnet'
    """

    model_type = "ltnet"

    def __init__(
        self,
        *,
        dataset_name: str = "coco",
        vocab_size: int = 206,
        obj_classes_size: int = 155,
        hidden_size: int = 256,
        num_hidden_layers: int = 4,
        num_attention_heads: int = 4,
        dropout: float = 0.1,
        enable_noise: bool = False,
        noise_size: int = 64,
        decoder_head_type: DecoderHeadType | str = DecoderHeadType.gmm,
        decoder_box_loss: BoxLossType | str = BoxLossType.pdf,
        decoder_schedule_sample: bool = False,
        decoder_two_path: bool = False,
        decoder_global_feature: bool = True,
        decoder_greedy: bool = True,
        xy_temperature: float = 1.0,
        wh_temperature: float = 1.0,
        refine: bool = False,
        refine_head_type: DecoderHeadType | str = DecoderHeadType.linear,
        refine_box_loss: BoxLossType | str = BoxLossType.reg,
        refine_x_softmax: bool = True,
        max_sequence_length: int = 128,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        relation_id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        pad_token_id: int = 0,
        mask_token_id: int = 3,
        model_type: str | None = None,
        transformers_version: str | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize LT-Net architecture and metadata fields."""
        _ = (model_type, transformers_version)

        self.dataset_name = dataset_name
        self.vocab_size = vocab_size
        self.obj_classes_size = obj_classes_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.dropout = dropout
        self.enable_noise = enable_noise
        self.noise_size = noise_size

        self.decoder_head_type = str(_normalize_head_type(decoder_head_type))
        self.decoder_box_loss = str(_normalize_box_loss(decoder_box_loss))
        self.decoder_schedule_sample = decoder_schedule_sample
        self.decoder_two_path = decoder_two_path
        self.decoder_global_feature = decoder_global_feature
        self.decoder_greedy = decoder_greedy
        self.xy_temperature = xy_temperature
        self.wh_temperature = wh_temperature
        self.refine = refine

        self.refine_head_type = str(_normalize_head_type(refine_head_type))
        self.refine_box_loss = str(_normalize_box_loss(refine_box_loss))
        self.refine_x_softmax = refine_x_softmax
        self.max_sequence_length = max_sequence_length
        self.mask_token_id = mask_token_id
        _ = kwargs

        super().__init__()
        normalized_id2label = id2label or DEFAULT_ID2LABEL
        normalized_relation_id2label = relation_id2label or DEFAULT_RELATION_ID2LABEL
        self.id2label = {
            int(key): str(value) for key, value in normalized_id2label.items()
        }
        self.relation_id2label = {
            int(key): str(value) for key, value in normalized_relation_id2label.items()
        }
        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
        self.pad_token_id = pad_token_id
        self.label2id = {value: key for key, value in self.id2label.items()}

class LTNetForLayoutGeneration(PreTrainedModel):
    """Transformers ``PreTrainedModel`` for LT-Net relation-to-layout inference."""

    config_class = LTNetConfig
    base_model_prefix = "ltnet"
    main_input_name = "input_token"
    _tied_weights_keys: dict[str, str] = {}

    def __init__(self, config: LTNetConfig) -> None:
        """Initialize relation encoder and bbox head."""
        super().__init__(config)
        self.encoder = RelEncoder(config)
        self.bbox_head = BBoxHead(config)
        self.all_tied_weights_keys = dict(self._tied_weights_keys)

    def forward(
        self,
        input_token: Int[torch.Tensor, "batch sequence"],
        input_obj_id: Int[torch.Tensor, "batch sequence"],
        segment_label: Int[torch.Tensor, "batch sequence"],
        token_type: Int[torch.Tensor, "batch sequence"],
        src_mask: Bool[torch.Tensor, "batch 1 sequence"] | None = None,
        global_mask: Bool[torch.Tensor, "batch sequence"] | None = None,
        bbox: Float[torch.Tensor, "batch sequence 4"] | None = None,
        bbox_mask: Bool[torch.Tensor, "batch sequence"] | None = None,
        inference: bool = False,
        generator: torch.Generator | None = None,
        output_hidden_states: bool = False,
        return_dict: bool = True,
    ) -> (
        LTNetModelOutput
        | tuple[Float[torch.Tensor, "batch sequence feature"] | None, ...]
    ):
        """Run LT-Net relation encoding and bbox prediction.

        Args:
            input_token: Mixed object/predicate token ids.
            input_obj_id: Stable object ids for object-token positions.
            segment_label: Relation segment ids.
            token_type: Token type ids ``0/1/2/3``.
            src_mask: Valid-token mask shaped ``(batch, 1, sequence)``.
            global_mask: Optional reference global-feature mask.
            bbox: Optional teacher-forced boxes for training/parity paths.
            bbox_mask: Optional box validity mask, reserved for parity paths.
            inference: Compatibility flag; greedy decoding is pipeline-owned.
            generator: Optional PyTorch generator for stochastic GMM sampling.
            output_hidden_states: Whether to include encoder hidden states.
            return_dict: Whether to return a ``ModelOutput``.

        Returns:
            Raw LT-Net output dataclass or tuple.

        Raises:
            ValueError: If required tensor shapes are invalid.
        """
        _ = bbox_mask
        if input_token.shape != input_obj_id.shape:
            raise ValueError("input_token and input_obj_id must have the same shape")

        effective_src_mask = (
            src_mask
            if src_mask is not None
            else input_token.ne(self.config.pad_token_id).unsqueeze(1)
        )
        if effective_src_mask.ndim == 2:
            effective_src_mask = effective_src_mask.unsqueeze(1)
        encoder_outputs = self.encoder(
            input_token,
            input_obj_id,
            segment_label,
            token_type,
            effective_src_mask,
        )
        (
            hidden_states,
            vocab_logits,
            obj_id_logits,
            token_type_logits,
            src,
            class_embeds,
        ) = encoder_outputs
        effective_global_mask = (
            global_mask if global_mask is not None else input_token.ge(2)
        )
        if inference:
            coarse_box, coarse_gmm, refine_box, refine_gmm = self.bbox_head.inference(
                hidden_states,
                effective_src_mask,
                src,
                class_embeds,
                effective_global_mask,
                generator=generator,
            )
        else:
            if bbox is None:
                bbox = input_token.new_full(
                    (input_token.size(0), input_token.size(1) - 1, 4), 2.0
                ).float()
            elif bbox.size(1) == input_token.size(1):
                bbox = bbox[:, :-1, :]
            trg_mask = effective_src_mask.new_ones([1, 1, 1])
            coarse_box, coarse_gmm, refine_box, refine_gmm = self.bbox_head(
                0,
                hidden_states,
                effective_src_mask,
                src,
                class_embeds,
                bbox,
                trg_mask,
                effective_global_mask,
                generator=generator,
            )
        output = LTNetModelOutput(
            vocab_logits=vocab_logits,
            obj_id_logits=obj_id_logits,
            token_type_logits=token_type_logits,
            coarse_box=coarse_box,
            coarse_gmm=coarse_gmm,
            refine_box=refine_box,
            refine_gmm=refine_gmm,
            hidden_states=hidden_states if output_hidden_states else None,
        )
        if return_dict:
            return output
        return output.to_tuple()

    @torch.no_grad()
    def _generate_boxes(
        self,
        input_token: Int[torch.Tensor, "batch sequence"],
        input_obj_id: Int[torch.Tensor, "batch sequence"],
        segment_label: Int[torch.Tensor, "batch sequence"],
        token_type: Int[torch.Tensor, "batch sequence"],
        src_mask: Bool[torch.Tensor, "batch 1 sequence"] | None = None,
        global_mask: Bool[torch.Tensor, "batch sequence"] | None = None,
        generator: torch.Generator | None = None,
    ) -> LTNetModelOutput:
        """Private pipeline helper for layout-level generation."""
        output = self(
            input_token=input_token,
            input_obj_id=input_obj_id,
            segment_label=segment_label,
            token_type=token_type,
            src_mask=src_mask,
            global_mask=global_mask,
            inference=True,
            generator=generator,
            return_dict=True,
        )
        return cast(LTNetModelOutput, output)

@dataclass
class LTNetModelOutput(ModelOutput):
    """Raw LT-Net model outputs.

    Attributes:
        vocab_logits: Mixed token vocabulary logits.
        obj_id_logits: Object-id classifier logits.
        token_type_logits: Token-type classifier logits.
        coarse_box: Coarse normalized center ``xywh`` boxes.
        coarse_gmm: Optional coarse GMM parameters.
        refine_box: Optional refined normalized center ``xywh`` boxes.
        refine_gmm: Optional refinement GMM parameters.
        hidden_states: Optional encoder hidden states.

    Examples:
        >>> import torch
        >>> out = LTNetModelOutput(
        ...     vocab_logits=torch.zeros(1, 2, 3),
        ...     obj_id_logits=torch.zeros(1, 2, 4),
        ...     token_type_logits=torch.zeros(1, 2, 4),
        ...     coarse_box=torch.zeros(1, 2, 4),
        ... )
        >>> out.coarse_box.shape
        torch.Size([1, 2, 4])
    """

    vocab_logits: Float[torch.Tensor, "batch sequence vocab"] | None = None
    obj_id_logits: Float[torch.Tensor, "batch sequence object_classes"] | None = None
    token_type_logits: Float[torch.Tensor, "batch sequence token_types"] | None = None
    coarse_box: Float[torch.Tensor, "batch sequence 4"] | None = None
    coarse_gmm: Float[torch.Tensor, "batch sequence gmm_params"] | None = None
    refine_box: Float[torch.Tensor, "batch sequence 4"] | None = None
    refine_gmm: Float[torch.Tensor, "batch sequence gmm_params"] | None = None
    hidden_states: Float[torch.Tensor, "batch sequence hidden"] | None = None

class LTNetPipeline(LayoutGenerationPipeline):
    """Compose an LT-Net model and processor for scene-graph layout inference.

    Args:
        model: Converted LT-Net model.
        processor: Matching scene-graph processor/tokenizer.
        config: Optional root pipeline config. Defaults to ``model.config``.

    Examples:
        >>> processor = LTNetProcessor.from_config()
        >>> config = LTNetConfig(
        ...     vocab_size=processor.tokenizer.vocab_size,
        ...     hidden_size=32,
        ...     num_hidden_layers=1,
        ...     num_attention_heads=4,
        ... )
        >>> pipe = LTNetPipeline(
        ...     model=LTNetForLayoutGeneration(config),
        ...     processor=processor,
        ... )
        >>> pipe.config.model_type
        'ltnet'
    """

    config_class: ClassVar[type[PretrainedConfig]] = LTNetConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = {
        "model": PipelineComponentSpec(
            attribute_name="model",
            loader=_load_model_component,
            marker_file="config.json",
        ),
        "processor": PipelineComponentSpec(
            attribute_name="processor",
            loader=_load_processor_component,
            marker_file="preprocessor_config.json",
            save_with_is_main_process=False,
        ),
    }

    config: LTNetConfig
    model: LTNetForLayoutGeneration
    processor: LTNetProcessor

    def __init__(
        self,
        model: LTNetForLayoutGeneration,
        processor: LTNetProcessor,
        config: LTNetConfig | None = None,
    ) -> None:
        """Initialize the pipeline with model and processor components."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor

    @classmethod
    def _from_pretrained_components(  # ty: ignore[invalid-method-override]
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, LTNetForLayoutGeneration | LTNetProcessor | None],
    ) -> "LTNetPipeline":
        """Build a pipeline from loaded root components."""
        return cls(
            config=cast(LTNetConfig, config),
            model=cast(LTNetForLayoutGeneration, components["model"]),
            processor=cast(LTNetProcessor, 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.relation,
        labels: Int[torch.Tensor, "batch elements"] | Sequence[int] | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Sequence[Sequence[float]]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"] | Sequence[bool] | 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 = "dataclass",
        return_intermediates: bool = False,
        scene_graph: SceneGraphInput | SceneGraphMapping | None = None,
        objects: Sequence[LayoutObject] | None = None,
        relations: Sequence[LayoutRelation] | None = None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | dict[str, Float[torch.Tensor, "..."] | None],
        ]
    ):
        """Generate layouts from a public scene graph.

        Args:
            batch_size: Number of layouts to generate from the same graph.
            seed: Convenience seed used only when ``generator`` is absent.
            generator: Optional PyTorch generator; takes precedence over seed.
            condition_type: Must normalize to ``relation``.
            labels: Reserved v1 interface input; LT-Net uses ``scene_graph``.
            bbox: Reserved v1 box constraint input.
            mask: Reserved v1 validity-mask input.
            num_elements: Reserved v1 element-count input.
            box_format: Output box format; LT-Net returns normalized ``xywh``.
            normalized: Whether output boxes should be normalized.
            canvas_size: Reserved denormalization canvas size.
            num_inference_steps: Reserved v1 step count.
            output_type: Return dataclass or dict.
            return_intermediates: Include raw logits/boxes in intermediates.
            scene_graph: Public relation payload.
            objects: Object-node shorthand when ``scene_graph`` is omitted.
            relations: Relation-edge shorthand when ``scene_graph`` is omitted.

        Returns:
            Layout generation output dataclass or dict.

        Raises:
            ValueError: If the condition or graph payload is unsupported.
        """
        _ = (labels, bbox, mask, num_elements, num_inference_steps)
        model_device = next(self.model.parameters()).device
        prepared_generator = self.prepare_generator(
            generator=generator, seed=seed, device=model_device
        )
        encoded = self.processor(
            scene_graph=scene_graph,
            objects=objects,
            relations=relations,
            batch_size=batch_size,
            condition_type=condition_type,
            return_tensors="pt",
        )
        model_inputs = {
            key: value.to(model_device)
            for key, value in encoded.items()
            if isinstance(value, torch.Tensor)
        }
        was_training = self.model.training
        self.model.eval()
        try:
            output = self.model._generate_boxes(
                **model_inputs,
                generator=prepared_generator,
            )
        finally:
            self.model.train(was_training)
        return self.processor.post_process_layout_generation(
            output,
            input_token=model_inputs["input_token"],
            input_obj_id=model_inputs["input_obj_id"],
            token_type=model_inputs["token_type"],
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            output_type=output_type,
            return_intermediates=return_intermediates,
        )

class LTNetProcessor(ProcessorMixin):
    """Normalize scene graphs, tokenize LT-Net inputs, and postprocess boxes."""

    attributes = ["tokenizer"]
    tokenizer_class = "LTNetRelationTokenizer"

    def __init__(
        self,
        tokenizer: LTNetRelationTokenizer,
        dataset_name: str = "coco",
        max_sequence_length: int = 128,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        relation_id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        object_reduce: Literal["first", "last", "mean"] = "first",
    ) -> None:
        """Initialize processor label maps and tokenizer component."""
        self.tokenizer = tokenizer
        self.dataset_name = dataset_name
        self.max_sequence_length = max_sequence_length
        self.id2label = {
            int(key): str(value)
            for key, value in (id2label or DEFAULT_ID2LABEL).items()
        }
        self.relation_id2label = {
            int(key): str(value)
            for key, value in (relation_id2label or DEFAULT_RELATION_ID2LABEL).items()
        }
        self.label2id = {value.lower(): key for key, value in self.id2label.items()}
        self.relation_label2id = {
            value.lower(): key for key, value in self.relation_id2label.items()
        }
        if object_reduce not in {"first", "last", "mean"}:
            raise ValueError("object_reduce must be 'first', 'last', or 'mean'")

        self.object_reduce = object_reduce
        super().__init__(tokenizer=tokenizer)

    @classmethod
    def from_config(
        cls,
        *,
        dataset_name: str = "coco",
        max_sequence_length: int = 128,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        relation_id2label: Mapping[int, str] | Mapping[str, str] | None = None,
    ) -> "LTNetProcessor":
        """Construct a processor and tokenizer without external files.

        Returns:
            Processor with synthetic vocabulary derived from the label maps.

        Examples:
            >>> processor = LTNetProcessor.from_config()
            >>> processor.tokenizer.cls_token_id
            1
        """
        object_labels = {
            int(key): str(value)
            for key, value in (id2label or DEFAULT_ID2LABEL).items()
        }
        relation_labels = {
            int(key): str(value)
            for key, value in (relation_id2label or DEFAULT_RELATION_ID2LABEL).items()
        }
        tokens = ["__image__"]
        tokens.extend(value for _, value in sorted(object_labels.items()))
        tokens.extend(value for _, value in sorted(relation_labels.items()))
        tokenizer = LTNetRelationTokenizer(tokens=tokens)
        return cls(
            tokenizer=tokenizer,
            dataset_name=dataset_name,
            max_sequence_length=max_sequence_length,
            id2label=object_labels,
            relation_id2label=relation_labels,
        )

    @classmethod
    def _load_tokenizer_from_pretrained(
        cls,
        sub_processor_type: str,
        pretrained_model_name_or_path: str | PathLike[str],
        subfolder: str = "",
        **kwargs: str | int | float | bool | None,
    ) -> LTNetRelationTokenizer:
        """Load tokenizer for ``ProcessorMixin.from_pretrained``."""
        _ = sub_processor_type
        path = Path(pretrained_model_name_or_path)
        tokenizer_path = path / subfolder if subfolder else path
        token = kwargs.get("token")
        return LTNetRelationTokenizer.from_pretrained(
            tokenizer_path,
            cache_dir=cast(str | PathLike[str] | None, kwargs.get("cache_dir")),
            force_download=bool(kwargs.get("force_download", False)),
            local_files_only=bool(kwargs.get("local_files_only", False)),
            token=token if isinstance(token, str | bool) else None,
            revision=str(kwargs.get("revision", "main")),
        )

    def normalize_condition_type(
        self, condition_type: ConditionType | str
    ) -> ConditionType:
        """Normalize and validate the LT-Net public condition type."""
        condition = normalize_condition_type(condition_type)
        if condition is not ConditionType.relation:
            raise ValueError(
                "LT-Net only supports condition_type='relation' "
                "and aliases 'scene_graph', 'graph', or 'gen_r'."
            )

        return condition

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

    def _relation_to_id(self, predicate: int | str) -> int:
        if isinstance(predicate, int):
            return predicate
        lowered = predicate.lower()
        if lowered in self.relation_label2id:
            return self.relation_label2id[lowered]
        raise ValueError(f"Unknown relation label: {predicate}")

    def _token_for_object(self, label_id: int) -> str:
        return self.id2label.get(label_id, str(label_id))

    def _token_for_relation(self, relation_id: int) -> str:
        return self.relation_id2label.get(relation_id, str(relation_id))

    def _normalize_scene_graph(
        self,
        scene_graph: SceneGraphInput | SceneGraphMapping | None,
        *,
        objects: Sequence[LayoutObject] | None,
        relations: Sequence[LayoutRelation] | None,
    ) -> SceneGraphInput:
        if isinstance(scene_graph, SceneGraphInput):
            return scene_graph
        if scene_graph is None:
            if objects is None:
                raise ValueError("scene_graph or objects must be provided")

            return SceneGraphInput(
                objects=tuple(objects),
                relations=tuple(relations or ()),
                id2label=self.id2label,
                relation_id2label=self.relation_id2label,
            )
        nodes = scene_graph.get("nodes", scene_graph.get("objects", ()))
        edges = scene_graph.get("edges", scene_graph.get("relations", ()))

        normalized_objects: list[LayoutObject] = []

        for node in cast(Sequence[SceneGraphItemMapping], nodes):
            item = node
            node_id = cast(int | str, item["id"])
            label = cast(int | str, item.get("label_id", item.get("label")))
            bbox = cast(tuple[float, float, float, float] | None, item.get("bbox"))
            normalized_objects.append(LayoutObject(id=node_id, label=label, bbox=bbox))

        normalized_relations: list[LayoutRelation] = []
        for edge in cast(Sequence[SceneGraphItemMapping], edges):
            item = edge
            subject = cast(int | str, item.get("source", item.get("subject")))
            predicate = cast(int | str, item.get("predicate_id", item.get("predicate")))
            target = cast(int | str, item.get("target", item.get("object")))
            normalized_relations.append(
                LayoutRelation(
                    subject=subject,
                    predicate=predicate,
                    object=target,
                    score=cast(float | None, item.get("score")),
                )
            )

        return SceneGraphInput(
            objects=tuple(normalized_objects),
            relations=tuple(normalized_relations),
            id2label=cast(dict[int, str] | None, scene_graph.get("id2label")),
            relation_id2label=cast(
                dict[int, str] | None,
                scene_graph.get("relation_id2label"),
            ),
        )

    def _serialize_graph(
        self,
        graph: SceneGraphInput,
    ) -> tuple[list[int], list[int], list[int], list[int], list[int]]:
        object_by_id = {item.id: item for item in graph.objects}
        object_ids = {item.id: idx + 1 for idx, item in enumerate(graph.objects)}
        tokens = [self.tokenizer.cls_token]
        input_obj_id = [0]
        segment_label = [0]
        token_type = [0]
        segment = 1
        for relation in graph.relations:
            subject = object_by_id[relation.subject]
            target = object_by_id[relation.object]
            subject_label = self._label_to_id(subject.label)
            target_label = self._label_to_id(target.label)
            relation_id = self._relation_to_id(relation.predicate)
            triple_tokens = [
                self._token_for_object(subject_label),
                self._token_for_relation(relation_id),
                self._token_for_object(target_label),
                self.tokenizer.sep_token,
            ]
            tokens.extend(triple_tokens)
            input_obj_id.extend([object_ids[subject.id], 0, object_ids[target.id], 0])
            segment_label.extend([segment] * 4)
            token_type.extend([1, 2, 3, 0])
            segment += 1
        if not graph.relations:
            for item in graph.objects:
                label_id = self._label_to_id(item.label)
                tokens.extend(
                    [self._token_for_object(label_id), self.tokenizer.sep_token]
                )
                input_obj_id.extend([object_ids[item.id], 0])
                segment_label.extend([segment, segment])
                token_type.extend([1, 0])
                segment += 1

        input_token = self.tokenizer.encode_scene_graph_tokens(tokens)
        length = min(len(input_token), self.max_sequence_length)
        input_token = input_token[:length]
        input_obj_id = input_obj_id[:length]
        segment_label = segment_label[:length]
        token_type = token_type[:length]

        src_mask = [1] * length
        pad_length = self.max_sequence_length - length
        input_token.extend([self.tokenizer.pad_token_id] * pad_length)
        input_obj_id.extend([0] * pad_length)
        segment_label.extend([0] * pad_length)
        token_type.extend([0] * pad_length)
        src_mask.extend([0] * pad_length)
        return input_token, input_obj_id, segment_label, token_type, src_mask

    def __call__(
        self,
        *,
        scene_graph: SceneGraphInput | SceneGraphMapping | None = None,
        objects: Sequence[LayoutObject] | None = None,
        relations: Sequence[LayoutRelation] | None = None,
        batch_size: int = 1,
        condition_type: ConditionType | str = ConditionType.relation,
        return_tensors: Literal["pt", "np"] = "pt",
        max_sequence_length: int | None = None,
    ) -> BatchEncoding:
        """Build LT-Net model tensors from public scene-graph inputs."""
        self.normalize_condition_type(condition_type)
        original_max_length = self.max_sequence_length
        try:
            if max_sequence_length is not None:
                self.max_sequence_length = max_sequence_length
            graph = self._normalize_scene_graph(
                scene_graph,
                objects=objects,
                relations=relations,
            )
            rows = [self._serialize_graph(graph) for _ in range(batch_size)]
        finally:
            self.max_sequence_length = original_max_length
        data = {
            "input_token": torch.tensor([row[0] for row in rows], dtype=torch.long),
            "input_obj_id": torch.tensor([row[1] for row in rows], dtype=torch.long),
            "segment_label": torch.tensor([row[2] for row in rows], dtype=torch.long),
            "token_type": torch.tensor([row[3] for row in rows], dtype=torch.long),
            "src_mask": torch.tensor(
                [row[4] for row in rows], dtype=torch.bool
            ).unsqueeze(1),
            "global_mask": torch.tensor([row[0] for row in rows], dtype=torch.long).ge(
                2
            ),
        }
        if return_tensors == "pt":
            return BatchEncoding(data)
        if return_tensors == "np":
            return BatchEncoding({key: value.numpy() for key, value in data.items()})
        raise ValueError("return_tensors must be 'pt' or 'np'")

    def post_process_layout_generation(
        self,
        model_outputs: LTNetModelOutput,
        *,
        input_token: Int[torch.Tensor, "batch sequence"] | None = None,
        input_obj_id: Int[torch.Tensor, "batch sequence"],
        token_type: Int[torch.Tensor, "batch sequence"],
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        output_type: OutputType = "dataclass",
        return_intermediates: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | dict[str, Float[torch.Tensor, "..."] | None],
        ]
    ):
        """Convert raw token-level boxes into public object-level layouts."""
        _ = (canvas_size, normalize_box_format(box_format))
        if not normalized:
            raise ValueError("LT-Net outputs normalized boxes only")

        raw_box = model_outputs.refine_box
        if raw_box is None:
            raw_box = model_outputs.coarse_box
        if raw_box is None:
            raise ValueError("model_outputs must contain coarse_box or refine_box")

        batch_boxes: list[Float[torch.Tensor, "elements 4"]] = []
        batch_labels: list[Int[torch.Tensor, "elements"]] = []
        batch_masks: list[Bool[torch.Tensor, "elements"]] = []
        token_rows = (
            [None] * raw_box.size(0)
            if input_token is None
            else list(input_token.unbind(dim=0))
        )
        for row_box, row_obj_id, row_type, row_token in zip(
            raw_box, input_obj_id, token_type, token_rows, strict=True
        ):
            object_positions = row_type.eq(1) | row_type.eq(3)
            object_ids = row_obj_id[object_positions]
            boxes = row_box[object_positions].clamp(0.0, 1.0)
            labels = (
                torch.clamp(object_ids - 1, min=0).long()
                if row_token is None
                else row_token[object_positions].long()
            )
            valid = object_ids.gt(0)
            boxes = boxes[valid]
            labels = labels[valid]
            object_ids = object_ids[valid]
            reduced_boxes: list[Float[torch.Tensor, "elements 4"]] = []
            reduced_labels: list[Int[torch.Tensor, "elements"]] = []
            for object_id in object_ids.unique(sorted=True):
                positions = object_ids.eq(object_id).nonzero().flatten()
                if self.object_reduce == "mean":
                    reduced_boxes.append(boxes[positions].mean(dim=0))
                    reduced_labels.append(labels[positions[0]])
                else:
                    selected = (
                        positions[0] if self.object_reduce == "first" else positions[-1]
                    )
                    reduced_boxes.append(boxes[selected])
                    reduced_labels.append(labels[selected])
            if reduced_boxes:
                batch_boxes.append(torch.stack(reduced_boxes))
                batch_labels.append(torch.stack(reduced_labels).long())
            else:
                batch_boxes.append(boxes)
                batch_labels.append(labels)
            batch_masks.append(
                torch.ones(len(reduced_boxes), dtype=torch.bool, device=row_box.device)
            )
        max_items = max((item.size(0) for item in batch_boxes), default=0)
        padded_boxes = raw_box.new_zeros((raw_box.size(0), max_items, 4))
        padded_labels = input_obj_id.new_zeros((raw_box.size(0), max_items))
        padded_masks = torch.zeros(
            (raw_box.size(0), max_items), dtype=torch.bool, device=raw_box.device
        )
        for idx, (boxes, labels, mask) in enumerate(
            zip(batch_boxes, batch_labels, batch_masks, strict=True)
        ):
            length = boxes.size(0)
            padded_boxes[idx, :length] = boxes
            padded_labels[idx, :length] = labels
            padded_masks[idx, :length] = mask
        intermediates = None
        if return_intermediates:
            intermediates = {
                "coarse_box": model_outputs.coarse_box,
                "refine_box": model_outputs.refine_box,
                "vocab_logits": model_outputs.vocab_logits,
                "obj_id_logits": model_outputs.obj_id_logits,
                "token_type_logits": model_outputs.token_type_logits,
            }
        output = LayoutGenerationOutput(
            bbox=padded_boxes,
            labels=padded_labels,
            mask=padded_masks,
            id2label=dict(self.id2label),
            intermediates=intermediates,
        )
        if output_type == "dict":
            return dict(output.items())
        return output

    def save_pretrained(
        self,
        save_directory: str | PathLike[str],
        push_to_hub: bool = False,
        **kwargs: str | int | float | bool | None,
    ) -> tuple[str, ...]:
        """Save processor metadata and tokenizer files."""
        paths = super().save_pretrained(
            save_directory, push_to_hub=push_to_hub, **kwargs
        )
        metadata = {
            "processor_class": self.__class__.__name__,
            "dataset_name": self.dataset_name,
            "max_sequence_length": self.max_sequence_length,
            "id2label": self.id2label,
            "relation_id2label": self.relation_id2label,
            "object_reduce": self.object_reduce,
        }
        with (Path(save_directory) / "preprocessor_config.json").open("w") as f:
            json.dump(metadata, f, indent=2, sort_keys=True)
        return paths