Layout detr

class BackgroundPreprocessing(StrEnum):
    """Supported public background preprocessing modes."""

    none = auto()
    resize_256 = "256"
    resize_128 = "128"
    blur = auto()
    edge = auto()

class LayoutDetrConfig(PretrainedConfig):
    """Configuration for LayoutDETR model, processor, and pipeline."""

    model_type = "layout-detr"

    def __init__(
        self,
        *,
        dataset_name: str = "ad_banner",
        id2label: Mapping[int | str, str] | None = None,
        max_seq_length: int = 9,
        z_dim: int = 4,
        img_channels: int = 3,
        img_height: int = 256,
        img_width: int = 256,
        background_size: int = 256,
        hidden_dim: int = 256,
        bert_f_dim: int = 768,
        bert_num_encoder_layers: int = 12,
        bert_num_decoder_layers: int = 2,
        bert_num_heads: int = 4,
        max_text_length: int = 256,
        text_vocab_size: int = 30_522,
        med_config: Mapping[str, LayoutDetrMetadataValue] | None = None,
        backbone_name: str = "resnet50",
        image_mean: Sequence[float] = (0.485, 0.456, 0.406),
        image_std: Sequence[float] = (0.229, 0.224, 0.225),
        architecture: Literal["lightweight", "reference"] = "lightweight",
        model_subfolder: str = "model",
        processor_subfolder: str = "processor",
        original_training_options: Mapping[str, LayoutDetrMetadataValue] | None = None,
        conversion_report: Mapping[str, LayoutDetrMetadataValue] | None = None,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize LayoutDETR configuration."""
        raw_id2label = id2label or DEFAULT_ID2LABEL
        normalized_id2label = {int(key): value for key, value in raw_id2label.items()}
        super().__init__(id2label=normalized_id2label, **kwargs)  # ty: ignore[invalid-argument-type]

        self.dataset_name = dataset_name
        self.id2label = normalized_id2label
        self.label2id = {value: key for key, value in self.id2label.items()}
        self.max_seq_length = int(max_seq_length)
        self.z_dim = int(z_dim)
        self.img_channels = int(img_channels)
        self.img_height = int(img_height)
        self.img_width = int(img_width)
        self.background_size = int(background_size)

        self.hidden_dim = int(hidden_dim)
        self.bert_f_dim = int(bert_f_dim)
        self.bert_num_encoder_layers = int(bert_num_encoder_layers)
        self.bert_num_decoder_layers = int(bert_num_decoder_layers)
        self.bert_num_heads = int(bert_num_heads)
        self.max_text_length = int(max_text_length)
        self.text_vocab_size = int(text_vocab_size)

        self.med_config = dict(med_config or {})
        self.backbone_name = backbone_name
        self.image_mean = tuple(float(value) for value in image_mean)
        self.image_std = tuple(float(value) for value in image_std)
        self.architecture = architecture
        self.model_subfolder = model_subfolder
        self.processor_subfolder = processor_subfolder
        self.original_training_options = dict(original_training_options or {})
        self.conversion_report = dict(conversion_report or {})

    @property
    def num_labels(self) -> int:
        """Return the number of public semantic labels."""
        return len(cast(dict[int, str], self.id2label))

    @property
    def num_bbox_labels(self) -> int:
        """Return the model bbox-label count."""
        return self.num_labels

    @property
    def pad_label_id(self) -> int:
        """Return the internal padded label id."""
        return 0

    @property
    def max_elements(self) -> int:
        """Return the maximum generated element count."""
        return self.max_seq_length

class LayoutDetrImageProcessor(BaseImageProcessor):
    """Prepare ImageNet-normalized background tensors for LayoutDETR."""

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        background_size: int = 256,
        image_mean: Sequence[float] = (0.485, 0.456, 0.406),
        image_std: Sequence[float] = (0.229, 0.224, 0.225),
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Initialize image normalization settings."""
        super().__init__(**kwargs)
        self.background_size = int(background_size)
        self.image_mean = tuple(float(value) for value in image_mean)
        self.image_std = tuple(float(value) for value in image_std)

    @classmethod
    def from_config(cls, config: LayoutDetrConfig) -> "LayoutDetrImageProcessor":
        """Build an image processor from a LayoutDETR config."""
        return cls(
            background_size=config.background_size,
            image_mean=config.image_mean,
            image_std=config.image_std,
        )

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput],
        *,
        background_preprocessing: BackgroundPreprocessing
        | str = BackgroundPreprocessing.none,
        canvas_size: tuple[int, int] | None = None,
        return_tensors: Literal["pt"] = "pt",
        **kwargs: str | int | float | bool | None,
    ) -> BatchFeature:
        """Preprocess a background image or batch.

        Args:
            images: PIL, NumPy, or torch image input.
            background_preprocessing: Released-checkpoint-compatible mode.
            canvas_size: Optional canvas metadata override.
            return_tensors: Only ``"pt"`` is supported.
            kwargs: Ignored compatibility kwargs.

        Returns:
            ``BatchFeature`` with ``pixel_values`` and ``canvas_size``.
        """
        del kwargs
        if return_tensors != "pt":
            raise ValueError(
                "LayoutDetrImageProcessor only supports return_tensors='pt'"
            )

        mode = normalize_background_preprocessing(background_preprocessing)
        tensors: list[Float[torch.Tensor, "channels height width"]] = []
        sizes: list[tuple[int, int]] = []
        for image in _ensure_pil_batch(images):
            sizes.append(canvas_size or image.size)
            processed = _apply_background_preprocessing(image.convert("RGB"), mode)
            processed = processed.resize(
                (self.background_size, self.background_size),
                Image.Resampling.BILINEAR,
            )
            array = np.asarray(processed, dtype=np.float32) / 255.0
            mean = np.asarray(self.image_mean, dtype=np.float32)
            std = np.asarray(self.image_std, dtype=np.float32)
            tensors.append(torch.from_numpy(((array - mean) / std).transpose(2, 0, 1)))
        return BatchFeature(
            {
                "pixel_values": torch.stack(tensors).float(),
                "canvas_size": torch.tensor(sizes, dtype=torch.long),
            }
        )

class LayoutDetrForConditionalGeneration(PreTrainedModel):
    """A standard ``PreTrainedModel`` wrapper for LayoutDETR forward inference."""

    config_class = LayoutDetrConfig
    base_model_prefix = "layout_detr"
    main_input_name = "pixel_values"
    supports_gradient_checkpointing = False

    def __init__(self, config: LayoutDetrConfig) -> None:
        """Initialize LayoutDETR layers."""
        super().__init__(config)
        self._is_reference_architecture = config.architecture == "reference"
        if self._is_reference_architecture:  # pragma: no cover
            self._init_reference_layers(config)
        else:
            self._init_lightweight_layers(config)
        self.post_init()

    def _init_lightweight_layers(self, config: LayoutDetrConfig) -> None:
        self.background_encoder = nn.Sequential(
            nn.Conv2d(config.img_channels, config.hidden_dim, kernel_size=3, padding=1),
            nn.GELU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.fc_z = nn.Linear(config.z_dim, config.bert_f_dim)
        self.emb_label = nn.Embedding(config.num_bbox_labels, config.bert_f_dim)
        self.text_embeddings = nn.Embedding(config.text_vocab_size, config.bert_f_dim)
        self.text_len_embeddings = nn.Embedding(
            config.max_text_length, config.bert_f_dim
        )
        self.background_proj = nn.Linear(config.hidden_dim, config.hidden_dim)
        self.fc_in = nn.Sequential(
            nn.Linear(config.bert_f_dim * 4, config.hidden_dim),
            nn.GELU(),
            nn.Linear(config.hidden_dim, config.hidden_dim),
        )
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.hidden_dim,
            nhead=max(1, min(8, config.hidden_dim // 8)),
            dim_feedforward=max(config.hidden_dim * 4, 64),
            batch_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.bbox_embed = nn.Sequential(
            nn.Linear(config.hidden_dim, config.hidden_dim),
            nn.GELU(),
            nn.Linear(config.hidden_dim, 4),
        )

    def _init_reference_layers(
        self, config: LayoutDetrConfig
    ) -> None:  # pragma: no cover
        self.backbone = _build_backbone(config.backbone_name)
        self.input_proj = nn.Conv2d(self.backbone.num_channels, config.hidden_dim, 1)
        self.fc_z = nn.Linear(config.z_dim * config.max_seq_length, config.bert_f_dim)
        self.emb_label = nn.Embedding(config.num_bbox_labels, config.bert_f_dim)
        self.text_encoder = _build_reference_bert_model(
            config,
            num_hidden_layers=config.bert_num_encoder_layers,
            encoder_width=config.bert_f_dim,
            add_pooling_layer=False,
            use_cross_attention=True,
            text_encoder=True,
        )
        self.enc_text_len = nn.Embedding(config.max_text_length, config.bert_f_dim)
        self.fc_in = _ReferenceMLP(
            input_dim=config.bert_f_dim * 4,
            hidden_dim=config.bert_f_dim,
            output_dim=config.hidden_dim,
            num_layers=3,
        )
        self.transformer = _DetrTransformer(
            d_model=config.hidden_dim,
            dropout=0.1,
            nhead=8,
            dim_feedforward=2048,
            num_encoder_layers=6,
            num_decoder_layers=6,
        )
        self.bbox_embed = _ReferenceMLP(
            input_dim=config.hidden_dim,
            hidden_dim=config.hidden_dim,
            output_dim=4,
            num_layers=3,
        )
        self.fc_z_rec = nn.Linear(
            config.hidden_dim, config.z_dim * config.max_seq_length
        )
        self.fc_out_cls = nn.Linear(config.hidden_dim, config.num_bbox_labels)
        self.text_decoder = _build_reference_bert_lm_head(
            config,
            num_hidden_layers=config.bert_num_decoder_layers,
            encoder_width=512,
        )
        self.fc_text_len_rec = nn.Linear(config.hidden_dim, config.max_text_length)

    def forward(
        self,
        *,
        pixel_values: Float[torch.Tensor, "batch channels height width"],
        input_ids: Int[torch.Tensor, "batch elements tokens"],
        text_attention_mask: Bool[torch.Tensor, "batch elements tokens"],
        bbox_labels: Int[torch.Tensor, "batch elements"],
        layout_mask: Bool[torch.Tensor, "batch elements"],
        latents: Float[torch.Tensor, "batch elements latent"],
        text_lengths: Int[torch.Tensor, "batch elements"] | None = None,
        return_dict: bool | None = None,
    ) -> (
        LayoutDetrModelOutput
        | tuple[
            Float[torch.Tensor, "batch elements 4"],
            Int[torch.Tensor, "batch elements"],
            Bool[torch.Tensor, "batch elements"],
        ]
    ):
        """Run the LayoutDETR conditional forward pass."""
        return_dict = (
            self.config.use_return_dict if return_dict is None else return_dict
        )
        if bbox_labels.ndim != 2:
            raise ValueError("bbox_labels must have shape (batch, elements)")

        if (
            latents.shape[:2] != bbox_labels.shape
            or latents.shape[-1] != self.config.z_dim
        ):
            raise ValueError("latents must have shape (batch, elements, z_dim)")

        if input_ids.shape[:2] != bbox_labels.shape:
            raise ValueError("input_ids must have shape (batch, elements, tokens)")

        labels = bbox_labels.to(dtype=torch.long)
        if labels.numel() and (
            int(labels.min().item()) < 0
            or int(labels.max().item()) >= self.config.num_bbox_labels
        ):
            raise ValueError("bbox_labels contain ids outside config.num_bbox_labels")

        device = labels.device
        pixel_values = pixel_values.to(device=device, dtype=self.dtype)
        latents = latents.to(device=device, dtype=self.dtype)
        input_ids = input_ids.to(device=device, dtype=torch.long)
        text_attention_mask = text_attention_mask.to(device=device, dtype=torch.bool)
        layout_mask = layout_mask.to(device=device, dtype=torch.bool)

        if self._is_reference_architecture:
            bbox, hidden = self._forward_reference(
                pixel_values=pixel_values,
                input_ids=input_ids,
                text_attention_mask=text_attention_mask,
                labels=labels,
                layout_mask=layout_mask,
                latents=latents,
                text_lengths=text_lengths,
            )
            if not return_dict:
                return bbox, labels, layout_mask
            return LayoutDetrModelOutput(
                bbox=bbox,
                labels=labels,
                mask=layout_mask,
                latents=latents,
                hidden_states=hidden,
            )

        bg = self.background_proj(self.background_encoder(pixel_values)).unsqueeze(1)
        z = self.fc_z(latents)
        label_features = self.emb_label(labels)
        text_tokens = self.text_embeddings(input_ids)
        token_mask = text_attention_mask.unsqueeze(-1).to(dtype=text_tokens.dtype)
        denom = token_mask.sum(dim=2).clamp_min(1.0)
        text_features = (text_tokens * token_mask).sum(dim=2) / denom
        lengths = text_attention_mask.sum(dim=-1).clamp_max(
            self.config.max_text_length - 1
        )
        text_len_features = self.text_len_embeddings(lengths)
        hidden = self.fc_in(
            torch.cat([z, label_features, text_features, text_len_features], dim=-1)
        )
        hidden = hidden + bg
        hidden = self.transformer(hidden, src_key_padding_mask=~layout_mask)
        bbox = torch.sigmoid(self.bbox_embed(hidden))
        if not return_dict:
            return bbox, labels, layout_mask
        return LayoutDetrModelOutput(
            bbox=bbox,
            labels=labels,
            mask=layout_mask,
            latents=latents,
            hidden_states=hidden,
        )

    def _forward_reference(  # pragma: no cover
        self,
        *,
        pixel_values: Float[torch.Tensor, "batch channels height width"],
        input_ids: Int[torch.Tensor, "batch elements tokens"],
        text_attention_mask: Bool[torch.Tensor, "batch elements tokens"],
        labels: Int[torch.Tensor, "batch elements"],
        layout_mask: Bool[torch.Tensor, "batch elements"],
        latents: Float[torch.Tensor, "batch elements latent"],
        text_lengths: Int[torch.Tensor, "batch elements"] | None = None,
    ) -> tuple[
        Float[torch.Tensor, "batch elements 4"],
        Float[torch.Tensor, "batch elements hidden"],
    ]:
        bg_nested = _NestedTensor(
            pixel_values,
            torch.zeros(
                pixel_values.shape[0],
                pixel_values.shape[2],
                pixel_values.shape[3],
                dtype=torch.bool,
                device=pixel_values.device,
            ),
        )
        bg_feat, pos = self.backbone(bg_nested)
        bg_tensor, bg_mask = bg_feat[-1].decompose()
        z0 = _normalize_2nd_moment(latents.reshape(latents.shape[0], -1))
        z = self.fc_z(z0).unsqueeze(1).expand(-1, labels.shape[1], -1)
        label_features = self.emb_label(labels)
        flat_input_ids = input_ids.reshape(-1, input_ids.shape[-1])
        flat_attention = text_attention_mask.reshape(-1, text_attention_mask.shape[-1])
        text_output = self.text_encoder(
            flat_input_ids,
            attention_mask=flat_attention,
            return_dict=True,
            mode="text",
        )
        text_features = text_output.last_hidden_state[:, 0, :].view(
            labels.shape[0], labels.shape[1], -1
        )
        if text_lengths is None:
            lengths = flat_attention.sum(dim=-1)
        else:
            lengths = text_lengths.to(device=labels.device, dtype=torch.long).reshape(
                -1
            )
        lengths = lengths.clamp_max(self.config.max_text_length - 1)
        text_len_features = self.enc_text_len(lengths.view(labels.shape))
        hidden = torch.cat(
            [z, label_features, text_features, text_len_features], dim=-1
        )
        hidden = torch.relu(self.fc_in(hidden)).permute(1, 0, 2)
        hidden = self.transformer(
            src=self.input_proj(bg_tensor),
            mask=bg_mask,
            pos_embed=pos[-1],
            tgt=hidden,
            tgt_key_padding_mask=~layout_mask,
        )[0]
        return torch.sigmoid(self.bbox_embed(hidden)), hidden

@dataclass
class LayoutDetrModelOutput(ModelOutput):
    """Raw LayoutDETR model output."""

    bbox: Float[torch.Tensor, "batch elements 4"]
    labels: Int[torch.Tensor, "batch elements"] = cast(
        Int[torch.Tensor, "batch elements"], None
    )
    mask: Bool[torch.Tensor, "batch elements"] = cast(
        Bool[torch.Tensor, "batch elements"], None
    )
    latents: Float[torch.Tensor, "batch elements latent"] | None = None
    hidden_states: Float[torch.Tensor, "batch elements hidden"] | None = None

class LayoutDetrPipeline(LayoutGenerationPipeline):
    """Transformers-side LayoutDETR pipeline."""

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

    config: LayoutDetrConfig
    model: LayoutDetrForConditionalGeneration
    processor: LayoutDetrProcessor

    def __init__(
        self,
        model: LayoutDetrForConditionalGeneration,
        processor: LayoutDetrProcessor | None = None,
        config: LayoutDetrConfig | None = None,
        device: str | torch.device | None = None,
    ) -> None:
        """Initialize a LayoutDETR pipeline."""
        super().__init__(config or model.config)
        self.config = config or model.config
        self.model = model
        self.processor = processor or LayoutDetrProcessor(config=self.config)
        self.model.eval()
        if device is not None:
            self.to(device)

    @classmethod
    def _from_pretrained_components(
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PipelineComponent | None],
    ) -> "LayoutDetrPipeline":
        """Build a pipeline from saved components."""
        return cls(
            config=cast(LayoutDetrConfig, config),
            model=cast(LayoutDetrForConditionalGeneration, components["model"]),
            processor=cast(LayoutDetrProcessor, components["processor"]),
        )

    @torch.no_grad()
    def __call__(
        self,
        images: ImageInput
        | Sequence[ImageInput]
        | Shaped[torch.Tensor, "..."]
        | None = None,
        *,
        content: Mapping[
            str,
            ImageInput
            | Sequence[ImageInput]
            | Sequence[Sequence[str]]
            | Sequence[str]
            | Sequence[Sequence[int | str]]
            | Sequence[int | str],
        ]
        | None = None,
        prompt: str | Sequence[str] | None = None,
        texts: Sequence[Sequence[str]] | Sequence[str] | None = None,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.content_image,
        labels: Int[torch.Tensor, "batch elements"]
        | Sequence[Sequence[int | str]]
        | Sequence[int | str]
        | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"] | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Sequence[Sequence[bool]]
        | Sequence[bool]
        | None = None,
        num_elements: int | Sequence[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,
        background_preprocessing: BackgroundPreprocessing
        | str = BackgroundPreprocessing.none,
        out_jittering_strength: float = 0.0,
        out_postprocessing: PostprocessingMode | str = PostprocessingMode.none,
        latents: Float[torch.Tensor, "batch elements latent"] | None = None,
    ) -> LayoutGenerationOutput:
        """Generate layouts for a background image and per-element text labels."""
        normalize_condition_type(condition_type)
        if bbox is not None:
            raise ValueError("LayoutDETR does not condition on existing bbox")

        if num_elements is not None:
            raise ValueError("LayoutDETR infers num_elements from labels/mask")

        if box_format != BoxFormat.xywh and box_format != "xywh":
            raise ValueError("LayoutDETR outputs normalized xywh boxes only")

        if not normalized:
            raise ValueError("LayoutDETR expects normalized public boxes")

        if num_inference_steps is not None:
            raise ValueError("LayoutDETR is a single forward pass, not iterative")

        encoded = self.processor(
            images=images,
            content=content,
            prompt=prompt,
            texts=texts,
            labels=labels,
            mask=mask,
            condition_type=str(ConditionType.content_image),
            background_preprocessing=background_preprocessing,
            batch_size=batch_size,
            canvas_size=canvas_size,
        )
        device = self.device or next(self.model.parameters()).device
        encoded = encoded.to(device)
        batch, elements = encoded["bbox_labels"].shape
        runtime_generator = self.prepare_generator(
            generator=generator,
            seed=seed,
            device=device,
        )
        if latents is None:
            latents = torch.randn(
                (batch, elements, self.config.z_dim),
                generator=runtime_generator,
                device=device,
            )
        else:
            latents = latents.to(device=device)
        model_output = self.model(
            pixel_values=encoded["pixel_values"],
            input_ids=encoded["input_ids"],
            text_attention_mask=encoded["text_attention_mask"],
            bbox_labels=encoded["bbox_labels"],
            layout_mask=encoded["layout_mask"],
            latents=latents,
            text_lengths=encoded["text_lengths"],
        )
        bbox_out = apply_postprocessing(
            model_output.bbox,
            model_output.mask,
            mode=out_postprocessing,
            jitter_strength=out_jittering_strength,
            generator=runtime_generator,
        )
        intermediates = {
            "latents": latents.detach().cpu(),
            "texts": encoded["texts"],
            "background_preprocessing": str(background_preprocessing),
            "postprocessing": str(out_postprocessing),
        }
        return cast(
            LayoutGenerationOutput,
            self.processor.post_process_layouts(
                bbox_out,
                model_output.labels,
                model_output.mask,
                output_type=output_type,
                return_intermediates=return_intermediates,
                intermediates=intermediates,
            ),
        )