Lace

@dataclass(frozen=True)
class LaceDatasetSpec:
    """Static dataset metadata used to configure a LACE checkpoint.

    Attributes:
        dataset: Canonical dataset name.
        labels: Ordered category labels without the padding class.
        max_seq_length: Maximum number of layout elements.
        dim_transformer: Transformer hidden size used by the original model.
        nhead: Number of attention heads.
        num_layers: Number of transformer blocks.
        dim_feedforward: Feed-forward hidden size.
    """

    dataset: DatasetName
    labels: tuple[LaceLabel, ...]
    max_seq_length: int = 25
    dim_transformer: int = 512
    nhead: int = 16
    num_layers: int = 4
    dim_feedforward: int = 2048

    @property
    def pad_label_id(self) -> int:
        """Return the integer id reserved for padding."""
        return len(self.labels)

    @property
    def num_classes_with_pad(self) -> int:
        """Return the number of category channels including padding."""
        return len(self.labels) + 1

    @property
    def seq_dim(self) -> int:
        """Return the latent per-element feature size."""
        return self.num_classes_with_pad + 4

    @property
    def id2label(self) -> dict[int, str]:
        """Return the category id to label mapping."""
        return dict(enumerate(str(label) for label in self.labels))

class ActivationName(StrEnum):
    """Supported feed-forward activation names.

    Origin:
        ``gelu2`` is the VQ-Diffusion GELU2/QuickGELU branch used by the
        LayoutDM, LACE, and LayoutFlow transformer utilities.
    """

    relu = auto()
    gelu = auto()
    gelu2 = auto()

@dataclass
class LaceModelOutput(BaseOutput):
    """Output returned by the LACE transformer.

    Attributes:
        sample: Predicted noise tensor with the same shape as the input sample.
    """

    sample: Float[torch.Tensor, "batch elements channels"]

class LaceTransformerModel(ModelMixin, ConfigMixin):
    """Transformer denoiser for continuous LACE layout tensors.

    Args:
        seq_dim: Number of channels per layout element.
        max_seq_length: Maximum number of layout elements.
        num_layers: Number of transformer blocks.
        dim_transformer: Hidden dimension.
        nhead: Number of attention heads.
        dim_feedforward: Feed-forward hidden dimension.
        diffusion_step: Maximum diffusion timestep.
        timestep_type: Timestep-conditioned normalization variant.
    """

    config_name = "model_config.json"

    pos_embed: Float[torch.Tensor, "max_seq_length dim_transformer"]

    @register_to_config
    def __init__(
        self,
        *,
        seq_dim: int,
        max_seq_length: int = 25,
        num_layers: int = 4,
        dim_transformer: int = 512,
        nhead: int = 16,
        dim_feedforward: int = 2048,
        diffusion_step: int = 1000,
        timestep_type: TimestepEmbeddingType
        | str
        | None = TimestepEmbeddingType.adalayernorm,
    ) -> None:
        """Initialize a LACE transformer denoiser."""
        super().__init__()
        self.seq_dim = seq_dim
        self.max_seq_length = max_seq_length
        self.pos_encoder = SinusoidalPosEmb(max_seq_length, dim_transformer)
        pos_i = torch.arange(max_seq_length)
        self.register_buffer("pos_embed", self.pos_encoder(pos_i), persistent=False)
        self.layer_in = nn.Linear(seq_dim, dim_transformer)
        encoder_layer = Block(
            d_model=dim_transformer,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            diffusion_step=diffusion_step,
            timestep_type=timestep_type,
        )
        self.layers = clone_module_list(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.layer_out = nn.Linear(dim_transformer, seq_dim)

    def forward(
        self,
        sample: Float[torch.Tensor, "batch elements channels"],
        timestep: Int[torch.Tensor, "batch"],
        attention_mask: Bool[torch.Tensor, "batch elements"] | None = None,
        return_dict: bool = True,
    ) -> LaceModelOutput | tuple[Float[torch.Tensor, "batch elements channels"]]:
        """Predict denoising residuals for a layout sample.

        Args:
            sample: Noisy layout tensor.
            timestep: Diffusion timestep per sample.
            attention_mask: Optional valid-element mask.
            return_dict: Whether to return ``LaceModelOutput``.

        Returns:
            Output dataclass or a one-item tuple containing the prediction.
        """
        output = F.softplus(self.layer_in(sample))
        pos_i = torch.arange(output.shape[1], device=output.device)
        output = output + self.pos_encoder(pos_i).to(output)
        key_padding_mask = None if attention_mask is None else ~attention_mask.bool()
        for i, layer in enumerate(self.layers):
            output = layer(
                output,
                src_key_padding_mask=key_padding_mask,
                timestep=timestep,
            )
            if i < self.num_layers - 1:
                output = F.softplus(output)
        output = self.layer_out(output)
        if not return_dict:
            return (output,)
        return LaceModelOutput(sample=output)

class TimestepEmbeddingType(StrEnum):
    """Supported timestep-conditioned normalization variants.

    Origin:
        These names come from VQ-Diffusion-derived adaptive normalization modes
        used by the LayoutDM and LACE checkpoint backbones.
    """

    adalayernorm = auto()
    adainnorm = auto()
    adalayernorm_abs = auto()
    adainnorm_abs = auto()
    adalayernorm_mlp = auto()
    adainnorm_mlp = auto()

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

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

class LacePipeline(DiffusionPipeline):
    """Generate layouts with a converted LACE checkpoint.

    Args:
        model: LACE transformer denoiser.
        scheduler: DDIM-style scheduler.
        processor: Processor that encodes and decodes layout tensors.

    Examples:
        >>> from lace import LaceProcessor, LaceScheduler, LaceTransformerModel
        >>> model = LaceTransformerModel(seq_dim=10, max_seq_length=2, num_layers=1, dim_transformer=8, nhead=2, dim_feedforward=16)
        >>> pipe = LacePipeline(model=model, scheduler=LaceScheduler(ddim_num_steps=1), processor=LaceProcessor.from_dataset("publaynet"))
        >>> pipe.processor.max_seq_length
        25
    """

    model_cpu_offload_seq = "model"
    _optional_components = ["processor"]

    def __init__(
        self,
        model: LaceTransformerModel,
        scheduler: LaceScheduler,
        processor: LaceProcessor,
    ) -> None:
        """Attach the converted LACE denoiser, scheduler, and processor."""
        super().__init__()
        self.register_modules(model=model, scheduler=scheduler)
        self.model = model
        self.scheduler = scheduler
        self.processor = processor
        self.model.eval()

    @property
    def components(
        self,
    ) -> dict[str, LaceTransformerModel | LaceScheduler | LaceProcessor]:
        """Expose modules and processor metadata for Diffusers serialization."""
        return dict(
            model=self.model, scheduler=self.scheduler, processor=self.processor
        )

    @torch.no_grad()
    def __call__(
        self,
        *,
        batch_size: int = 1,
        num_inference_steps: int | None = None,
        generator: torch.Generator | None = None,
        seed: int | None = None,
        condition_type: ConditionType | str | None = ConditionType.unconditional,
        bbox: Float[torch.Tensor, "batch elements 4"] | None = None,
        labels: Int[torch.Tensor, "batch elements"] | None = None,
        mask: Bool[torch.Tensor, "batch elements"] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        completion_ratio: float = 0.2,
        refinement_noise: float = 0.1,
        beautify: bool = False,
        beautify_overlap_weight: float | None = None,
        beautify_alignment_weight: float = 1.0,
        output_type: PipelineOutputType | str = PipelineOutputType.dataclass,
        return_intermediates: bool = False,
    ) -> LayoutGenerationOutput | dict[str, Shaped[torch.Tensor, "..."]]:
        """Run LACE denoising and return generated layouts.

        Args:
            batch_size: Number of layouts to generate for unconditional calls.
            num_inference_steps: Number of DDIM steps. Uses scheduler default if
                omitted.
            generator: Optional torch generator. Takes precedence over ``seed``.
            seed: Convenience seed used only when ``generator`` is absent.
            condition_type: Conditioning mode or alias.
            bbox: Conditioning boxes for non-unconditional modes.
            labels: Conditioning labels for non-unconditional modes.
            mask: Optional conditioning mask.
            box_format: Input box format for conditioning boxes.
            normalized: Whether conditioning boxes are normalized.
            canvas_size: Pixel canvas size required when ``normalized`` is false.
            completion_ratio: Maximum random completion fraction.
            refinement_noise: Noise scale for refinement conditioning.
            beautify: Whether to run the aesthetic post-optimization.
            beautify_overlap_weight: Optional overlap penalty override.
            beautify_alignment_weight: Alignment penalty weight.
            output_type: ``"dataclass"`` or ``"dict"``.
            return_intermediates: Whether to return the denoising trajectory.

        Returns:
            Layout output dataclass or a dictionary.

        Raises:
            ValueError: If a condition/output mode is unsupported or required
                conditioning tensors are missing.
        """
        if generator is None and seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        canonical = normalize_condition_type(condition_type)
        encoded = None
        if canonical is not ConditionType.unconditional:
            if bbox is None or labels is None:
                raise ValueError(
                    f"bbox and labels are required for condition_type={condition_type}"
                )

            encoded = self.processor(
                bbox=bbox,
                labels=labels,
                mask=mask,
                box_format=box_format,
                normalized=normalized,
                canvas_size=canvas_size,
            )
            batch_size = encoded[LACE_LAYOUT_KEY].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,
        )
        if canonical is ConditionType.refinement:
            assert encoded is not None
            noise = torch.randn(
                encoded[LACE_BBOX_KEY].shape,
                dtype=encoded[LACE_BBOX_KEY].dtype,
                device=encoded[LACE_BBOX_KEY].device,
                generator=generator,
            )
            noisy_bbox = (encoded[LACE_BBOX_KEY] + refinement_noise * noise).clamp(0, 1)
            sample = self.processor.encode(
                noisy_bbox.to(self.device),
                encoded[LACE_LABELS_KEY].to(self.device),
                encoded[LACE_MASK_KEY].to(self.device),
            )
        real_layout = (
            None if encoded is None else encoded[LACE_LAYOUT_KEY].to(self.device)
        )
        fix_mask = self._build_fix_mask(
            canonical, real_layout, completion_ratio, generator
        )
        trajectory = [] if return_intermediates else None
        step_indices = (
            self.scheduler.refinement_indices()
            if canonical is ConditionType.refinement
            else list(range(len(self.scheduler.timesteps) - 1, -1, -1))
        )
        for index in step_indices:
            step = self.scheduler.ddim_timesteps[index]
            timestep = torch.full(
                (batch_size,), int(step.item()), device=self.device, dtype=torch.long
            )
            if real_layout is not None and fix_mask is not None:
                sample[fix_mask] = real_layout[fix_mask]
            model_output = self.model(sample=sample, timestep=timestep).sample
            out = self.scheduler.step(
                model_output, timestep, sample, index=index, generator=generator
            )
            sample = out.prev_sample
            if real_layout is not None and fix_mask is not None:
                sample[fix_mask] = real_layout[fix_mask]
            if trajectory is not None:
                trajectory.append(sample.detach().cpu())
        decoded = self.processor.decode(sample.detach().cpu())
        bbox_out = decoded.bbox
        mask_out = decoded.mask
        if beautify:
            overlap = beautify_overlap_weight
            if overlap is None:
                overlap = (
                    1.0
                    if normalize_dataset(self.processor.dataset)
                    is DatasetName.publaynet
                    else 0.0
                )
            bbox_out, mask_out = beautify_layout(
                bbox_out,
                mask_out,
                overlap_weight=overlap,
                alignment_weight=beautify_alignment_weight,
            )
        output = LayoutGenerationOutput(
            bbox=bbox_out,
            labels=decoded.labels,
            mask=mask_out,
            id2label=decoded.id2label,
            trajectory=trajectory,
            intermediates={"condition_type": str(canonical)}
            if return_intermediates
            else None,
        )
        out_type = normalize_output_type(output_type)
        if out_type is PipelineOutputType.dict:
            return dict(output)
        if out_type is PipelineOutputType.dataclass:
            return output
        assert_never(out_type)

    generate = __call__

    def save_pretrained(self, save_directory: str | Path) -> None:
        """Save pipeline components.

        Args:
            save_directory: Output directory.
        """
        super().save_pretrained(save_directory)
        self.processor.save_pretrained(save_directory)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        processor: LaceProcessor | None = None,
    ) -> "LacePipeline":
        """Load a saved LACE pipeline.

        Args:
            pretrained_model_name_or_path: Local path or Hub id.
            processor: Optional processor override.

        Returns:
            Loaded pipeline with the serialized processor attached.
        """
        loaded_processor = (
            LaceProcessor.from_pretrained(pretrained_model_name_or_path)
            if processor is None
            else processor
        )
        pipe = super().from_pretrained(pretrained_model_name_or_path)
        pipe.processor = loaded_processor
        return pipe

    def _build_fix_mask(
        self,
        condition_type: ConditionType,
        real_layout: Float[torch.Tensor, "batch elements channels"] | None,
        completion_ratio: float,
        generator: torch.Generator | None,
    ) -> Bool[torch.Tensor, "batch elements channels"] | None:
        """Build the fixed-channel mask for conditional generation."""
        if real_layout is None or condition_type is ConditionType.refinement:
            return None

        batch_size, seq_len, seq_dim = real_layout.shape
        num_class = seq_dim - 4

        if condition_type is ConditionType.label:
            fix_mask = torch.zeros_like(real_layout, dtype=torch.bool)
            fix_mask[:, :, :num_class] = True
            return fix_mask

        if condition_type is ConditionType.label_size:
            fix_mask = torch.zeros_like(real_layout, dtype=torch.bool)
            fix_indices = list(range(num_class)) + [num_class + 2, num_class + 3]
            fix_mask[:, :, fix_indices] = True
            return fix_mask

        if condition_type is ConditionType.completion:
            labels = real_layout[:, :, :num_class].argmax(dim=2)
            real_mask = labels != (num_class - 1)
            cutoff = torch.rand(
                (), device=real_layout.device, generator=generator
            ).item()
            element_mask = (
                torch.rand(
                    batch_size,
                    seq_len,
                    device=real_layout.device,
                    generator=generator,
                )
                <= cutoff * completion_ratio
            ) & real_mask
            return element_mask.unsqueeze(-1).expand(-1, -1, seq_dim)

        if condition_type is ConditionType.unconditional:
            return None
        raise ValueError(f"Unsupported LACE condition_type: {condition_type}")

class PipelineOutputType(StrEnum):
    """Supported LACE pipeline output containers."""

    dataclass = auto()
    dict = auto()

class LaceProcessor(ProcessorMixin):
    """Encode public layout tensors into the continuous LACE sequence format.

    Args:
        dataset: Canonical dataset name serialized with the processor config.
        labels: Ordered category labels without the padding label.
        max_seq_length: Maximum number of layout elements.

    Examples:
        >>> processor = LaceProcessor.from_dataset("publaynet")
        >>> processor.seq_dim
        10
    """

    config_name = "processor_config.json"

    def __init__(
        self,
        dataset: DatasetName | str,
        labels: list[str],
        max_seq_length: int = 25,
    ) -> None:
        """Initialize processor metadata.

        Args:
            dataset: Canonical dataset name.
            labels: Ordered category labels without padding.
            max_seq_length: Maximum number of layout elements.
        """
        super().__init__()
        self.dataset = str(normalize_dataset(dataset))
        self.labels = tuple(labels)
        self.max_seq_length = max_seq_length

    @classmethod
    def from_dataset(cls, dataset: DatasetName | str) -> "LaceProcessor":
        """Create a processor from built-in dataset metadata.

        Args:
            dataset: LACE dataset name or alias.

        Returns:
            Processor configured for the dataset.

        Raises:
            ValueError: If the dataset is unsupported.

        Examples:
            >>> LaceProcessor.from_dataset("rico13").pad_label_id
            13
        """
        spec = get_dataset_spec(dataset)
        return cls(
            dataset=str(spec.dataset),
            labels=[str(label) for label in spec.labels],
            max_seq_length=spec.max_seq_length,
        )

    @property
    def id2label(self) -> dict[int, str]:
        """Return the category id to label mapping."""
        return dict(enumerate(self.labels))

    @property
    def pad_label_id(self) -> int:
        """Return the padding label id."""
        return len(self.labels)

    @property
    def num_classes_with_pad(self) -> int:
        """Return the label-channel count including padding."""
        return len(self.labels) + 1

    @property
    def seq_dim(self) -> int:
        """Return the latent per-element feature size."""
        return self.num_classes_with_pad + 4

    def __call__(
        self,
        *,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput],
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput],
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
    ) -> dict[str, Shaped[torch.Tensor, "..."]]:
        """Encode a public layout batch.

        Args:
            bbox: Boxes in ``box_format``.
            labels: Integer category labels.
            mask: Optional valid-element mask.
            box_format: Input box format.
            normalized: Whether input coordinates are already normalized.
            canvas_size: Pixel canvas size required when ``normalized`` is false.

        Returns:
            Dictionary containing encoded layout, normalized boxes, labels, and mask.

        Raises:
            ValueError: If pixel boxes are passed without ``canvas_size`` or if
                ``box_format`` is unsupported.

        Examples:
            >>> processor = LaceProcessor.from_dataset("publaynet")
            >>> out = processor(bbox=[[[0.5, 0.5, 0.2, 0.2]]], labels=[[0]])
            >>> tuple(out["layout"].shape)
            (1, 25, 10)
        """
        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,
        )
        bbox_t, labels_t, mask_t = self.pad(bbox_t, labels_t, mask_t)
        return {
            LACE_LAYOUT_KEY: self.encode(bbox_t, labels_t, mask_t),
            LACE_BBOX_KEY: bbox_t,
            LACE_LABELS_KEY: labels_t,
            LACE_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,
        max_seq_length: int | None = None,
    ) -> tuple[
        Float[torch.Tensor, "batch padded_elements 4"],
        Int[torch.Tensor, "batch padded_elements"],
        Bool[torch.Tensor, "batch padded_elements"],
    ]:
        """Pad a batch to ``max_seq_length``.

        Args:
            bbox: Normalized boxes with shape ``(batch, seq, 4)``.
            labels: Integer labels with shape ``(batch, seq)``.
            mask: Optional valid-element mask.
            max_seq_length: Optional override for output length.

        Returns:
            Padded boxes, labels, and mask.

        Raises:
            ValueError: If the input has too many elements.
        """
        max_len = max_seq_length or self.max_seq_length
        if bbox.shape[1] > max_len:
            raise ValueError(f"LACE supports at most {max_len} elements")

        if mask is None:
            mask = torch.ones(labels.shape, dtype=torch.bool, device=labels.device)
        pad_count = max_len - bbox.shape[1]
        if pad_count:
            bbox_pad = torch.zeros(
                bbox.shape[0], pad_count, 4, dtype=bbox.dtype, device=bbox.device
            )
            label_pad = torch.full(
                (labels.shape[0], pad_count),
                self.pad_label_id,
                dtype=labels.dtype,
                device=labels.device,
            )
            mask_pad = torch.zeros(
                mask.shape[0], pad_count, dtype=torch.bool, device=mask.device
            )
            bbox = torch.cat((bbox, bbox_pad), dim=1)
            labels = torch.cat((labels, label_pad), dim=1)
            mask = torch.cat((mask, mask_pad), dim=1)
        labels = labels.clone()
        labels[~mask] = self.pad_label_id
        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 padded_elements channels"]:
        """Encode normalized boxes and labels into the LACE latent range.

        Args:
            bbox: Normalized center ``xywh`` boxes.
            labels: Integer labels.
            mask: Optional valid-element mask.

        Returns:
            Tensor with one-hot labels followed by box channels in ``[-1, 1]``.
        """
        bbox, labels, mask = self.pad(bbox, labels, mask)
        bbox_in = 2 * (bbox.clamp(0.0, 1.0) - 0.5)
        labels = labels.clamp(0, self.pad_label_id)
        labels[~mask] = self.pad_label_id
        one_hot = torch.nn.functional.one_hot(
            labels, num_classes=self.num_classes_with_pad
        ).to(dtype=bbox.dtype, device=bbox.device)
        return torch.cat((one_hot, bbox_in), dim=-1)

    def decode(
        self, layout: Float[torch.Tensor, "batch elements channels"], clamp: bool = True
    ) -> LayoutGenerationOutput:
        """Decode a LACE layout tensor into public output fields.

        Args:
            layout: Tensor with one-hot label channels and box channels.
            clamp: Whether to clamp latent box channels before conversion.

        Returns:
            Layout generation output with boxes in normalized center ``xywh``.
        """
        decoded = layout.clone()
        bbox_latent = decoded[:, :, self.num_classes_with_pad :]
        if clamp:
            bbox_latent = bbox_latent.clamp(-1.0, 1.0)
        bbox = bbox_latent / 2 + 0.5
        labels = decoded[:, :, : self.num_classes_with_pad].argmax(dim=2).long()
        mask = labels != self.pad_label_id
        return LayoutGenerationOutput(
            bbox=bbox,
            labels=labels,
            mask=mask,
            id2label=self.id2label,
            intermediates={"dataset": self.dataset},
        )

    def save_pretrained(  # ty: ignore[invalid-method-override]
        self, save_directory: str | Path
    ) -> None:
        """Save processor config to a Diffusers directory.

        Args:
            save_directory: Directory where ``processor_config.json`` is written.
        """
        super().save_pretrained(save_directory)

    @classmethod
    def from_pretrained(  # ty: ignore[invalid-method-override]
        cls,
        pretrained_model_name_or_path: str | Path,
        cache_dir: str | Path | None = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: str | bool | None = None,
        revision: str = "main",
    ) -> "LaceProcessor":
        """Load processor config from a Diffusers directory.

        Args:
            pretrained_model_name_or_path: Directory or Hub id containing
                ``processor_config.json``.
            cache_dir: Optional Hugging Face cache directory.
            force_download: Whether to force a fresh download.
            local_files_only: Whether to avoid network access.
            token: Optional Hugging Face token.
            revision: Hub revision to load.

        Returns:
            Loaded processor.
        """
        return super().from_pretrained(
            pretrained_model_name_or_path,
            cache_dir=cache_dir,
            force_download=force_download,
            local_files_only=local_files_only,
            token=token,
            revision=revision,
        )