Cgb dm

class CGBDMConfig(ConfigMixin):
    """Store CGB-DM architecture, schedule, and dataset metadata.

    Args:
        dataset_name: Poster/content dataset key.
        num_labels: Internal class-channel count, including invalid/pad.
        max_seq_length: Maximum number of layout elements.
        image_size: Model image size as ``(height, width)``.
        canvas_size: Dataset canvas size as ``(width, height)``.
        num_train_timesteps: DDPM training timesteps.
        ddim_num_steps: Default DDIM inference steps.
        dim_model: Transformer hidden dimension.
        n_head: Attention head count.
        num_layers: Number of layout decoder layers.
        feature_dim: Feed-forward hidden dimension.
        id2label: Public id-to-label mapping, excluding invalid/pad.

    Examples:
        >>> CGBDMConfig().dataset_name
        'pku_posterlayout'
    """

    config_name = "cgb_dm_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        dataset_name: DatasetName | str = DatasetName.pku_posterlayout,
        num_labels: int | None = None,
        max_seq_length: int = 16,
        image_size: tuple[int, int] | list[int] = (384, 256),
        canvas_size: tuple[int, int] | list[int] = (513, 750),
        num_train_timesteps: int = 1000,
        ddim_num_steps: int = 100,
        dim_model: int = 512,
        n_head: int = 8,
        num_layers: int = 4,
        feature_dim: int = 1024,
        id2label: Id2LabelMapping | None = None,
        condition_types: list[str] | tuple[str, ...] | None = None,
        train_beta_schedule: str = "cosine",
        sampling_beta_schedule: str = "linear",
        model_subfolder: str = "model",
        scheduler_subfolder: str = "scheduler",
        processor_subfolder: str = "processor",
    ) -> None:
        """Initialize CGB-DM configuration."""
        dataset = normalize_dataset_name(dataset_name)
        spec = DATASET_SPECS.get(dataset)
        if spec is None:
            raise ValueError(f"Unsupported CGB-DM dataset_name: {dataset_name}")

        self.dataset_name = str(dataset)
        self.num_labels = int(num_labels or spec.num_labels)
        self.max_seq_length = int(max_seq_length)
        self.image_size: tuple[int, int] = (int(image_size[0]), int(image_size[1]))
        self.canvas_size: tuple[int, int] = (int(canvas_size[0]), int(canvas_size[1]))
        self.num_train_timesteps = int(num_train_timesteps)
        self.ddim_num_steps = int(ddim_num_steps)
        self.dim_model = int(dim_model)
        self.n_head = int(n_head)
        self.num_layers = int(num_layers)
        self.feature_dim = int(feature_dim)
        self.id2label = {int(k): v for k, v in (id2label or spec.id2label).items()}
        self.condition_types = list(
            condition_types
            or ["content_image", "label", "label_size", "completion", "refinement"]
        )
        self.train_beta_schedule = train_beta_schedule
        self.sampling_beta_schedule = sampling_beta_schedule
        self.model_subfolder = model_subfolder
        self.scheduler_subfolder = scheduler_subfolder
        self.processor_subfolder = processor_subfolder

    @property
    def seq_dim(self) -> int:
        """Return the internal layout channel count."""
        return self.num_labels + 4

    @property
    def public_num_labels(self) -> int:
        """Return the public semantic label count."""
        return len(self.id2label)

@dataclass
class CGBDMModelOutput(BaseOutput):
    """Output returned by the CGB-DM denoiser.

    Attributes:
        sample: Predicted epsilon tensor with the same shape as the input layout.
        cgb_weight: Content-graphic balance weight estimated from image tokens.
    """

    sample: Float[torch.Tensor, "batch elements channels"]
    cgb_weight: Float[torch.Tensor, "batch 1 1"] | None = None

class CGBDMTransformerModel(ModelMixin, ConfigMixin):
    """CGB-DM transformer denoiser with image and saliency conditioning.

    Args:
        num_labels: Internal class-channel count including invalid/pad.
        max_seq_length: Maximum number of layout elements.
        image_size: Image tensor size as ``(height, width)``.
        patch_size: Image patch size.
        dim_model: Hidden dimension.
        n_head: Attention head count.
        feature_dim: Feed-forward hidden dimension.
        num_layers: Number of decoder layers.
        num_train_timesteps: Number of training diffusion steps.

    Examples:
        >>> model = CGBDMTransformerModel(num_labels=4, max_seq_length=2, image_size=(32, 32), dim_model=16, n_head=2, feature_dim=32, num_layers=1)
        >>> model.seq_dim
        8
    """

    config_name = "model_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        num_labels: int = 4,
        max_seq_length: int = 16,
        image_size: tuple[int, int] | list[int] = (384, 256),
        patch_size: int = 32,
        dim_model: int = 512,
        n_head: int = 8,
        feature_dim: int = 1024,
        num_layers: int = 4,
        num_train_timesteps: int = 1000,
    ) -> None:
        """Initialize the CGB-DM denoising network."""
        super().__init__()
        self.num_labels = int(num_labels)
        self.max_seq_length = int(max_seq_length)
        self.image_size: tuple[int, int] = (int(image_size[0]), int(image_size[1]))
        self.seq_dim = self.num_labels + 4
        self.img_encoder = CGBDMImageEncoder(
            image_size=self.image_size,
            patch_size=patch_size,
            in_channels=4,
            dim_model=dim_model,
            depth=6,
            heads=8,
            mlp_dim=2048,
            dropout=0.1,
            emb_dropout=0.1,
        )
        self.layout_encoder = CGBDMLayoutModule(
            seq_dim=self.seq_dim,
            dim_model=dim_model,
            n_head=n_head,
            feature_dim=feature_dim,
            num_layers=num_layers // 2,
            num_train_timesteps=num_train_timesteps,
            max_seq_length=self.max_seq_length,
            if_encoder=True,
        )
        self.layout_decoder = CGBDMLayoutModule(
            seq_dim=self.seq_dim,
            dim_model=dim_model,
            n_head=n_head,
            feature_dim=feature_dim,
            num_layers=num_layers,
            num_train_timesteps=num_train_timesteps,
            max_seq_length=self.max_seq_length,
            if_encoder=False,
        )
        self.cgbwp = CGBDMQFormer(in_dim=dim_model, num_tokens=1)
        self.salbox_encoder = CGBDMMLP(4, dim_model, dim_model)

    def forward(
        self,
        sample: Float[torch.Tensor, "batch elements channels"],
        image: Float[torch.Tensor, "batch 4 height width"],
        saliency_box: Float[torch.Tensor, "batch 1 4"],
        timestep: Int[torch.Tensor, "batch"],
        return_dict: bool = True,
    ) -> CGBDMModelOutput | tuple[Float[torch.Tensor, "batch elements channels"]]:
        """Predict epsilon for a noisy layout tensor.

        Args:
            sample: Noisy class-plus-box layout tensor.
            image: Four-channel RGB/saliency tensor in ``[-1, 1]``.
            saliency_box: Saliency box tensor in internal ``[-1, 1]`` center xywh.
            timestep: Per-example diffusion timestep ids.
            return_dict: Whether to return ``CGBDMModelOutput``.

        Returns:
            Output dataclass or one-item tuple containing predicted epsilon.
        """
        image_tokens = self.img_encoder(image)
        saliency_tokens = self.salbox_encoder(saliency_box)
        encoded = self.layout_encoder(sample, None, None, None, timestep)
        cgb_weight = self.cgbwp(image_tokens)
        pred = self.layout_decoder(
            encoded,
            image_tokens,
            cgb_weight,
            saliency_tokens,
            timestep,
        )
        if not return_dict:
            return (pred,)
        return CGBDMModelOutput(sample=pred, cgb_weight=cgb_weight)

class CGBDMPipeline(DiffusionPipeline):
    """Generate content-aware poster layouts with CGB-DM.

    Args:
        model: CGB-DM denoiser.
        scheduler: CGB-DM scheduler.
        processor: Processor for images and layouts.

    Examples:
        >>> model = CGBDMTransformerModel(num_labels=4, max_seq_length=2, image_size=(32, 32), dim_model=16, n_head=2, feature_dim=32, num_layers=1)
        >>> pipe = CGBDMPipeline(model=model, scheduler=CGBDMScheduler(num_train_timesteps=10, ddim_num_steps=1), processor=CGBDMProcessor(max_seq_length=2, image_size=(32, 32)))
        >>> pipe.processor.seq_dim
        8
    """

    model_cpu_offload_seq = "model"

    def __init__(
        self,
        model: CGBDMTransformerModel,
        scheduler: CGBDMScheduler,
        processor: CGBDMProcessor,
    ) -> None:
        """Initialize pipeline components."""
        super().__init__()
        self.register_modules(model=model, scheduler=scheduler, processor=processor)
        self.model = model
        self.scheduler = scheduler
        self.processor = processor
        self.model.eval()

    @property
    def components(
        self,
    ) -> dict[str, CGBDMTransformerModel | CGBDMScheduler | CGBDMProcessor]:
        """Return serializable pipeline components."""
        return {
            "model": self.model,
            "scheduler": self.scheduler,
            "processor": self.processor,
        }

    @torch.no_grad()
    def __call__(
        self,
        *,
        image: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | None = None,
        content: dict[
            str,
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes],
        ]
        | None = None,
        saliency: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | None = None,
        saliency_isnet: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | None = None,
        saliency_basnet: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | None = None,
        saliency_box: Float[torch.Tensor, "..."] | None = None,
        pixel_values: Float[torch.Tensor, "batch channels height width"] | 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, "..."]
        | Int[np.ndarray, "..."]
        | list[list[int]]
        | list[int]
        | list[ArrayLikeInput]
        | None = None,
        bbox: Float[torch.Tensor, "..."]
        | Float[np.ndarray, "..."]
        | list[list[list[float]]]
        | list[list[float]]
        | list[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "..."]
        | Bool[np.ndarray, "..."]
        | list[list[bool]]
        | list[bool]
        | list[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "..."] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        num_inference_steps: int | None = None,
        completion_ratio: float = 0.2,
        output_type: OutputType | str = OutputType.dataclass,
        return_intermediates: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | list[Float[torch.Tensor, "..."]]
            | dict[str, str | Float[torch.Tensor, "..."] | None]
            | None,
        ]
    ):
        """Run DDIM sampling and return generated layouts.

        Args:
            image: RGB image or batch of images.
            content: Optional content container with ``image`` and saliency keys.
            saliency: Optional merged saliency map.
            saliency_isnet: Optional first saliency map.
            saliency_basnet: Optional second saliency map.
            saliency_box: Optional normalized center ``xywh`` saliency box.
            pixel_values: Preprocessed four-channel image tensor.
            batch_size: Number of layouts when ``pixel_values`` is synthetic.
            seed: Convenience seed used when ``generator`` is absent.
            generator: Optional torch generator. Takes precedence over ``seed``.
            condition_type: Canonical condition mode or alias.
            labels: Conditioning labels for constrained modes.
            bbox: Conditioning boxes for constrained modes.
            mask: Optional valid-element mask.
            num_elements: Accepted for interface compatibility.
            box_format: Input box format.
            normalized: Whether input boxes are normalized.
            canvas_size: Canvas size required for pixel boxes.
            num_inference_steps: DDIM step count.
            completion_ratio: Completion conditioning keep ratio.
            output_type: ``"dataclass"`` or ``"dict"``.
            return_intermediates: Whether to include trajectory/debug tensors.

        Returns:
            Layout output dataclass or dictionary.

        Raises:
            ValueError: If required content or conditioning inputs are absent.
        """
        del num_elements
        canonical = normalize_condition_type(condition_type)
        out_type = normalize_output_type(output_type)
        if generator is None and seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        if pixel_values is None:
            if content is not None:
                image = content.get("image", image)
                saliency = content.get("saliency", saliency)
            encoded_content = self.processor(
                image,
                saliency=saliency,
                saliency_isnet=saliency_isnet,
                saliency_basnet=saliency_basnet,
                saliency_box=saliency_box,
            )
            pixel_values = encoded_content["pixel_values"]
            resolved_saliency_box = encoded_content["saliency_box"]
        else:
            if saliency_box is None:
                resolved_saliency_box = torch.zeros(pixel_values.shape[0], 1, 4)
            else:
                resolved_saliency_box = 2 * (
                    torch.as_tensor(saliency_box, dtype=torch.float32).clamp(0.0, 1.0)
                    - 0.5
                )
                if resolved_saliency_box.ndim == 2:
                    resolved_saliency_box = resolved_saliency_box.unsqueeze(1)
        batch_size = int(
            pixel_values.shape[0] if pixel_values is not None else batch_size
        )
        encoded_layout = None
        if canonical is not ConditionType.content_image:
            if bbox is None or labels is None:
                raise ValueError(
                    f"bbox and labels are required for condition_type={condition_type}"
                )

            encoded_layout = self.processor.encode_layout(
                bbox=bbox,
                labels=labels,
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
            )
            batch_size = encoded_layout["layout"].shape[0]
        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
        sample = self.scheduler.initial_sample(
            batch_size,
            self.processor.max_seq_length,
            self.processor.seq_dim,
            device=self.device,
            generator=generator,
        )
        real_layout = (
            None if encoded_layout is None else encoded_layout["layout"].to(self.device)
        )
        fix_mask = None
        if real_layout is not None:
            fix_mask = self.scheduler.condition_mask(
                real_layout,
                canonical,
                completion_ratio=completion_ratio,
                generator=generator,
            )
            sample = torch.where(fix_mask, real_layout, sample)
        pixel_values = pixel_values.to(self.device)
        resolved_saliency_box = resolved_saliency_box.to(self.device)
        trajectory: list[Float[torch.Tensor, "..."]] = []
        cgb_weights: list[Float[torch.Tensor, "..."]] = []
        for index, timestep in enumerate(self.scheduler.timesteps):
            timestep_batch = torch.full(
                (batch_size,),
                int(timestep.item()),
                device=self.device,
                dtype=torch.long,
            )
            model_out = self.model(
                sample, pixel_values, resolved_saliency_box, timestep_batch
            )
            step = self.scheduler.step(
                model_out.sample,
                timestep_batch,
                sample,
                len(self.scheduler.timesteps) - index - 1,
                generator=generator,
            )
            sample = step.prev_sample
            if real_layout is not None and fix_mask is not None:
                sample = torch.where(fix_mask, real_layout, sample)
            if return_intermediates:
                trajectory.append(sample.detach().cpu())
                if model_out.cgb_weight is not None:
                    cgb_weights.append(model_out.cgb_weight.detach().cpu())
        intermediates = None
        if return_intermediates:
            intermediates = {
                "condition_type": str(canonical),
                "saliency_box": resolved_saliency_box.detach().cpu(),
                "cgb_weight": cgb_weights[-1] if cgb_weights else None,
            }
        return self._decode(sample, out_type, trajectory, intermediates)

    def _decode(
        self,
        sample: Float[torch.Tensor, "..."],
        output_type: OutputType,
        trajectory: list[Float[torch.Tensor, "..."]],
        intermediates: dict[str, str | Float[torch.Tensor, "..."] | None] | None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | list[Float[torch.Tensor, "..."]]
            | dict[str, str | Float[torch.Tensor, "..."] | None]
            | None,
        ]
    ):
        decoded = self.processor.decode(
            sample.detach().cpu(),
            output_type="dataclass",
            intermediates=intermediates,
        )
        output = LayoutGenerationOutput(
            bbox=decoded.bbox,
            labels=decoded.labels,
            mask=decoded.mask,
            id2label=decoded.id2label,
            sequences=decoded.sequences,
            scores=decoded.scores,
            trajectory=trajectory or None,
            intermediates=decoded.intermediates,
        )
        if output_type is OutputType.dict:
            return dict(output)
        return output

class OutputType(StrEnum):
    """Supported CGB-DM pipeline output containers."""

    dataclass = auto()
    dict = auto()

class CGBDMProcessor(ProcessorMixin):
    """Prepare RGB/saliency inputs and decode CGB-DM layout tensors.

    Args:
        dataset_name: Poster/content dataset key.
        id2label: Public id-to-label mapping excluding invalid/pad.
        num_labels: Internal class-channel count.
        max_seq_length: Maximum number of elements.
        image_size: Resize target as ``(height, width)``.

    Examples:
        >>> CGBDMProcessor().seq_dim
        8
    """

    config_name = "processor_config.json"

    def __init__(
        self,
        dataset_name: DatasetName | str = DatasetName.pku_posterlayout,
        id2label: Mapping[int, str] | Mapping[str, str] | None = None,
        num_labels: int = 4,
        max_seq_length: int = 16,
        image_size: tuple[int, int] | list[int] = (384, 256),
    ) -> None:
        """Initialize processor metadata."""
        super().__init__()
        dataset = normalize_dataset_name(dataset_name)
        labels = id2label_for_dataset(dataset)
        public_labels = {
            key: value for key, value in labels.items() if value != "INVALID"
        }
        self.dataset_name = str(dataset)
        self.id2label = {int(k): v for k, v in (id2label or public_labels).items()}
        self.label2id = {v: k for k, v in self.id2label.items()}
        self.num_labels = int(num_labels)
        self.max_seq_length = int(max_seq_length)
        self.image_size: tuple[int, int] = (int(image_size[0]), int(image_size[1]))
        self.chat_template = None

    @property
    def seq_dim(self) -> int:
        """Return the internal class-plus-box channel count."""
        return self.num_labels + 4

    def __call__(
        self,
        images: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None = None,
        *,
        saliency: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None = None,
        saliency_isnet: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None = None,
        saliency_basnet: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None = None,
        saliency_box: Float[torch.Tensor, "..."] | None = None,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Encode image and saliency inputs into model tensors."""
        if return_tensors != "pt":
            raise ValueError("CGBDMProcessor only supports return_tensors='pt'")

        image_rows = _ensure_batch(images)
        saliency_rows = self._resolve_saliency(
            len(image_rows),
            saliency=saliency,
            saliency_isnet=saliency_isnet,
            saliency_basnet=saliency_basnet,
        )
        pixel_values = []
        boxes = []
        for image, sal in zip(image_rows, saliency_rows, strict=True):
            rgb = _to_rgb_tensor(image, self.image_size)
            sal_tensor = (
                torch.zeros(1, *self.image_size)
                if sal is None
                else _to_l_tensor(sal, self.image_size)
            )
            pixel_values.append(torch.cat((rgb, sal_tensor), dim=0))
            boxes.append(_saliency_box_from_tensor(sal_tensor))
        resolved_box = (
            torch.stack(boxes)
            if saliency_box is None
            else torch.as_tensor(saliency_box, dtype=torch.float32)
        )
        if resolved_box.ndim == 1:
            resolved_box = resolved_box.reshape(1, 1, 4)
        elif resolved_box.ndim == 2:
            resolved_box = resolved_box.unsqueeze(1)
        return BatchEncoding(
            {
                "pixel_values": torch.stack(pixel_values),
                "saliency_box": 2 * (resolved_box.clamp(0.0, 1.0) - 0.5),
            }
        )

    def encode_layout(
        self,
        *,
        bbox: Float[torch.Tensor, "..."]
        | Float[np.ndarray, "..."]
        | Sequence[Sequence[Sequence[float]]]
        | Sequence[Sequence[float]]
        | Sequence[ArrayLikeInput],
        labels: Int[torch.Tensor, "..."]
        | Int[np.ndarray, "..."]
        | Sequence[Sequence[int]]
        | Sequence[int]
        | Sequence[ArrayLikeInput],
        mask: Bool[torch.Tensor, "..."]
        | Bool[np.ndarray, "..."]
        | Sequence[Sequence[bool]]
        | Sequence[bool]
        | Sequence[ArrayLikeInput]
        | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
    ) -> CGBDMEncodedLayout:
        """Encode public layout tensors into CGB-DM latent layout format."""
        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,
        )
        bbox_t, labels_t, mask_t = self.pad(bbox_t, labels_t, mask_t)
        layout = self.encode(bbox_t, labels_t, mask_t)
        return {
            CGBDM_LAYOUT_KEY: layout,
            CGBDM_BBOX_KEY: bbox_t,
            CGBDM_LABELS_KEY: labels_t,
            CGBDM_MASK_KEY: mask_t,
        }

    def pad(
        self,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
    ) -> tuple[
        Float[torch.Tensor, "batch elements 4"],
        Int[torch.Tensor, "batch elements"],
        Bool[torch.Tensor, "batch elements"],
    ]:
        """Pad public layout tensors to ``max_seq_length``."""
        if mask is None:
            mask = torch.ones(labels.shape, dtype=torch.bool, device=labels.device)
        if bbox.shape[1] > self.max_seq_length:
            raise ValueError(f"CGB-DM supports at most {self.max_seq_length} elements")

        pad_count = self.max_seq_length - bbox.shape[1]
        if pad_count:
            bbox = torch.nn.functional.pad(bbox, (0, 0, 0, pad_count))
            labels = torch.nn.functional.pad(labels, (0, pad_count))
            mask = torch.nn.functional.pad(mask, (0, pad_count))
        labels = labels.clone()
        labels[~mask] = 0
        return bbox, labels, mask

    def encode(
        self,
        bbox: Float[torch.Tensor, "batch elements 4"],
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
    ) -> Float[torch.Tensor, "batch elements channels"]:
        """Encode normalized boxes and public labels into internal tensors."""
        bbox, labels, mask = self.pad(bbox, labels, mask)
        if self.dataset_name == str(DatasetName.pku_posterlayout):
            internal_labels = labels.clone().clamp_min(0) + 1
            internal_labels[~mask] = 0
        else:
            internal_labels = labels.clone().clamp_min(0)
            internal_labels[~mask] = 0
        one_hot = torch.nn.functional.one_hot(
            internal_labels.clamp(0, self.num_labels - 1),
            num_classes=self.num_labels,
        ).to(dtype=bbox.dtype, device=bbox.device)
        bbox_in = 2 * (bbox.clamp(0.0, 1.0) - 0.5)
        return torch.cat((one_hot, bbox_in), dim=-1)

    def decode(
        self,
        layout: Float[torch.Tensor, "batch elements channels"],
        *,
        output_type: Literal["dataclass", "dict"] = "dataclass",
        scores: Float[torch.Tensor, "..."] | None = None,
        intermediates: dict[str, str | Float[torch.Tensor, "..."] | None] | None = None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | dict[str, str | Float[torch.Tensor, "..."] | None]
            | None,
        ]
    ):
        """Decode internal layout tensors into the public schema."""
        bbox = (layout[:, :, self.num_labels :].clamp(-1.0, 1.0) / 2 + 0.5).cpu()
        class_logits = layout[:, :, : self.num_labels]
        class_ids = class_logits.argmax(dim=-1).long().cpu()
        mask = class_ids != 0
        if self.dataset_name == str(DatasetName.pku_posterlayout):
            labels = (class_ids - 1).clamp(0, max(self.id2label)).cpu()
        else:
            labels = class_ids.clamp(0, max(self.id2label)).cpu()
        resolved_scores = (
            scores
            if scores is not None
            else class_logits.softmax(dim=-1).max(dim=-1).values
        )
        output = LayoutGenerationOutput(
            bbox=bbox,
            labels=labels,
            mask=mask,
            id2label=dict(self.id2label),
            scores=resolved_scores.detach().cpu(),
            intermediates=intermediates,
        )
        if output_type == "dict":
            return dict(output)
        if output_type == "dataclass":
            return output
        raise ValueError(f"Unsupported output_type: {output_type}")

    def _resolve_saliency(
        self,
        batch_size: int,
        *,
        saliency: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None,
        saliency_isnet: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None,
        saliency_basnet: Float[torch.Tensor, "..."]
        | Image.Image
        | str
        | bytes
        | Path
        | IO[bytes]
        | Sequence[
            Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes]
        ]
        | None,
    ) -> list[
        Float[torch.Tensor, "..."] | Image.Image | str | bytes | Path | IO[bytes] | None
    ]:
        if saliency is not None:
            rows = _ensure_batch(saliency)
            if len(rows) != batch_size:
                raise ValueError("saliency batch size must match images")

            return cast(
                list[
                    Float[torch.Tensor, "..."]
                    | Image.Image
                    | str
                    | bytes
                    | Path
                    | IO[bytes]
                    | None
                ],
                rows,
            )
        if saliency_isnet is None and saliency_basnet is None:
            return [None] * batch_size
        return cast(
            list[
                Float[torch.Tensor, "..."]
                | Image.Image
                | str
                | bytes
                | Path
                | IO[bytes]
                | None
            ],
            [
                _merge_saliency_pair(left, right, self.image_size)
                for left, right in zip(
                    _optional_batch(saliency_isnet, batch_size),
                    _optional_batch(saliency_basnet, batch_size),
                    strict=True,
                )
            ],
        )