Layousyn

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()

@dataclass
class LayoutGenerationOutput(BaseOutput):
    """Layout-generation output for Diffusers pipelines.

    Attributes:
        bbox: Normalized center ``xywh`` boxes with shape
            ``(batch, elements, 4)``.
        labels: Dataset-local integer labels with shape ``(batch, elements)``.
        mask: Boolean valid-element mask with shape ``(batch, elements)``.
        id2label: Mapping from integer label ids to display names.
        sequences: Optional raw token sequences.
        scores: Optional per-token or per-element scores.
        trajectory: Optional sampling trajectory.
        intermediates: Optional model-specific debug or auxiliary data.

    Examples:
        >>> import torch
        >>> output = LayoutGenerationOutput(
        ...     bbox=torch.zeros(1, 1, 4),
        ...     labels=torch.zeros(1, 1, dtype=torch.long),
        ...     mask=torch.ones(1, 1, dtype=torch.bool),
        ...     id2label={0: "text"},
        ... )
        >>> output.to_tuple()[0].shape
        torch.Size([1, 1, 4])
    """

    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
    )
    id2label: dict[int, str] = cast(dict[int, str], None)
    sequences: object | None = None
    scores: object | None = None
    trajectory: object | None = None
    intermediates: object | None = None

class LayouSynConfig(ConfigMixin):
    """Serializable LayouSyn configuration.

    Args:
        model_name: Reference DiT architecture key.
        in_channels: Layout coordinate channels.
        concept_in_channels: Concept embedding width.
        y_in_channels: Caption embedding width.
        max_in_len: Maximum number of object slots.
        max_y_len: Maximum number of caption tokens.
        layout_type: Reference layout coordinate type.
        t5_size: Reference T5 size suffix.
        scale: Default classifier-free guidance scale.
        noise_schedule: Reference diffusion beta schedule.
        diffusion_steps: Number of diffusion training timesteps.
        hidden_size: Optional resolved hidden width override.
        depth: Optional resolved transformer depth override.
        num_heads: Optional resolved attention head override.
        license: Upstream checkpoint license identifier.
    """

    config_name = "config.json"

    @register_to_config
    def __init__(
        self,
        *,
        model_name: str = "DiT-S",
        in_channels: int = 4,
        concept_in_channels: int = 768,
        y_in_channels: int | None = 768,
        max_in_len: int = 60,
        max_y_len: int | None = 120,
        layout_type: LayoutType = "xyxy",
        t5_size: str | None = "base",
        scale: float = 2.0,
        noise_schedule: str = "linear",
        diffusion_steps: int = 100,
        hidden_size: int | None = None,
        depth: int | None = None,
        num_heads: int | None = None,
        license: str = "cc-by-nc-4.0",
    ) -> None:
        """Initialize configuration fields."""
        shape = resolve_model_shape(
            model_name,
            hidden_size=hidden_size,
            depth=depth,
            num_heads=num_heads,
        )
        self.model_name = model_name
        self.in_channels = in_channels
        self.concept_in_channels = concept_in_channels
        self.y_in_channels = y_in_channels
        self.max_in_len = max_in_len
        self.max_y_len = max_y_len
        self.layout_type = layout_type
        self.t5_size = t5_size
        self.scale = scale
        self.noise_schedule = noise_schedule
        self.diffusion_steps = diffusion_steps
        self.hidden_size = shape["hidden_size"]
        self.depth = shape["depth"]
        self.num_heads = shape["num_heads"]
        self.license = license

    @classmethod
    def from_reference_json(cls, path: str | Path) -> "LayouSynConfig":
        """Load a reference JSON config.

        Args:
            path: Path to a Lay-Your-Scene JSON config.

        Returns:
            Converted configuration object.
        """
        data = json.loads(Path(path).read_text())
        layout_type = data.get("layout_type", "xyxy")
        if not isinstance(layout_type, str):
            layout_type = str(layout_type)
        return cls(
            model_name=data.get("model", "DiT-S"),
            in_channels=data.get("in_channel", 4),
            concept_in_channels=data.get("concept_in_channel", 768),
            y_in_channels=data.get("y_in_channel"),
            max_in_len=data.get("max_in_len", 60),
            max_y_len=data.get("max_y_len"),
            layout_type="cxcywh" if "cxcywh" in layout_type.lower() else "xyxy",
            t5_size=data.get("t5_size"),
            scale=data.get("scale", 1.0),
            noise_schedule=data.get("noise_schedule", "linear"),
            diffusion_steps=data.get("diffusion_steps", 1000),
        )

    def to_reference_dict(self) -> LayouSynReferenceConfig:
        """Return the config keys expected by the original repository."""
        return {
            "model": self.model_name,
            "in_channel": self.in_channels,
            "concept_in_channel": self.concept_in_channels,
            "y_in_channel": self.y_in_channels,
            "max_in_len": self.max_in_len,
            "max_y_len": self.max_y_len,
            "scale": self.scale,
            "noise_schedule": self.noise_schedule,
            "layout_type": self.layout_type,
            "diffusion_steps": self.diffusion_steps,
            "t5_size": self.t5_size,
        }

class LayouSynDiTModel(ModelMixin, ConfigMixin):
    """Converted LayouSyn DiT denoiser."""

    config_name = "model_config.json"

    @register_to_config
    def __init__(
        self,
        *,
        in_channels: int = 4,
        max_in_len: int = 60,
        concept_in_channels: int = 768,
        y_in_channels: int | None = 768,
        max_y_len: int | None = 120,
        model_name: str = "DiT-S",
        hidden_size: int | None = None,
        depth: int | None = None,
        num_heads: int | None = None,
        mlp_ratio: float = 4.0,
        class_dropout_prob: float = 0.1,
        learn_sigma: bool = True,
        is_unconditional: bool = False,
    ) -> None:
        """Initialize the converted DiT model."""
        super().__init__()
        shape = resolve_model_shape(
            model_name,
            hidden_size=hidden_size,
            depth=depth,
            num_heads=num_heads,
        )
        self.in_channels = in_channels
        self.learn_sigma = learn_sigma
        self.num_heads = shape["num_heads"]
        self.max_len = max_in_len
        self.is_unconditional = is_unconditional
        self.x_embedder = InputEmbedder(in_channels, shape["hidden_size"])
        self.concept_embedder = ConceptEmbedder(
            concept_in_channels, shape["hidden_size"]
        )
        self.t_embedder = ScalarEmbedder(shape["hidden_size"])
        self.ar_embedder = ScalarEmbedder(shape["hidden_size"])
        if is_unconditional:
            self.y_embedder = CaptionEmbedderIdentity()
        else:
            if y_in_channels is None or max_y_len is None:
                raise ValueError("y_in_channels and max_y_len are required")

            y_null = torch.zeros(max_y_len, y_in_channels)
            y_mask = torch.ones(max_y_len, dtype=torch.bool)
            self.y_embedder = CaptionEmbedder(
                y_in_channels,
                shape["hidden_size"],
                class_dropout_prob,
                y_null,
                y_mask,
            )
        self.pos_embed = nn.Parameter(
            torch.zeros(1, max_in_len, shape["hidden_size"]), requires_grad=False
        )
        block_cls = DiTUCBlock if is_unconditional else DiTBlock
        self.blocks = nn.ModuleList(
            [
                block_cls(shape["hidden_size"], shape["num_heads"], mlp_ratio=mlp_ratio)
                for _ in range(shape["depth"])
            ]
        )
        self.final_layer = FinalLayer(shape["hidden_size"], 2 * in_channels)
        self.initialize_weights()

    def initialize_weights(self) -> None:
        """Initialize weights with the reference policy."""

        def _basic_init(module: nn.Module) -> None:
            if isinstance(module, nn.Linear):
                torch.nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)

        self.apply(_basic_init)
        pos_embed = get_1d_sincos_pos_embed(self.pos_embed.shape[-1], self.max_len)
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
        nn.init.xavier_uniform_(self.x_embedder.proj.weight)
        nn.init.constant_(self.x_embedder.proj.bias, 0)

        if not self.is_unconditional and isinstance(self.y_embedder, CaptionEmbedder):
            nn.init.normal_(self.y_embedder.proj.fc1.weight, std=0.02)
            nn.init.normal_(self.y_embedder.proj.fc2.weight, std=0.02)

        nn.init.normal_(self.concept_embedder.proj.fc1.weight, std=0.02)
        nn.init.normal_(self.concept_embedder.proj.fc2.weight, std=0.02)
        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
        nn.init.normal_(self.ar_embedder.mlp[0].weight, std=0.02)
        nn.init.normal_(self.ar_embedder.mlp[2].weight, std=0.02)

        for block in self.blocks:
            block.initialize_weights()

        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
        nn.init.constant_(self.final_layer.linear.weight, 0)
        nn.init.constant_(self.final_layer.linear.bias, 0)

    def forward(
        self,
        sample: Float[torch.Tensor, "batch elements channels"],
        timestep: Int[torch.Tensor, "batch"],
        *,
        x_padding_mask: Bool[torch.Tensor, "batch elements"],
        aspect_ratio: Float[torch.Tensor, "batch"],
        concept_embeds: Float[torch.Tensor, "batch elements embedding_dim"],
        caption_embeds: Float[torch.Tensor, "batch tokens embedding_dim"] | None = None,
        caption_padding_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
    ) -> Float[torch.Tensor, "batch seq channels"]:
        """Predict epsilon and variance channels for one timestep."""
        x = self.x_embedder(sample)
        x_enc = self.concept_embedder(concept_embeds)
        c = self.t_embedder(timestep) + self.ar_embedder(aspect_ratio)
        if self.is_unconditional:
            for block in self.blocks:
                x = block(x, x_padding_mask=x_padding_mask, c=c)
        else:
            if caption_embeds is None or caption_padding_mask is None:
                raise ValueError("caption_embeds and caption_padding_mask are required")

            y, y_padding_mask = self.y_embedder(
                caption_embeds, caption_padding_mask, self.training
            )
            for block in self.blocks:
                x, x_enc = block(
                    x,
                    x_enc,
                    x_padding_mask,
                    c,
                    y=y,
                    y_padding_mask=y_padding_mask,
                    pos_embed=self.pos_embed[:, : x.shape[1]],
                )
        out = self.final_layer(x, c).chunk(2, dim=-1)
        return torch.cat([out[0], out[1]], dim=1)

    def forward_with_cfg(
        self,
        sample: Float[torch.Tensor, "batch elements channels"],
        timestep: Int[torch.Tensor, "batch"],
        *,
        x_padding_mask: Bool[torch.Tensor, "batch elements"],
        aspect_ratio: Float[torch.Tensor, "batch"],
        concept_embeds: Float[torch.Tensor, "batch elements embedding_dim"],
        caption_embeds: Float[torch.Tensor, "batch tokens embedding_dim"],
        caption_padding_mask: Bool[torch.Tensor, "batch tokens"],
        guidance_scale: float,
    ) -> Float[torch.Tensor, "batch seq channels"]:
        """Run reference classifier-free guidance batching."""
        half = sample[: len(sample) // 2]
        combined = torch.cat([half, half], dim=0)
        model_out = self.forward(
            combined,
            timestep,
            x_padding_mask=x_padding_mask,
            aspect_ratio=aspect_ratio,
            concept_embeds=concept_embeds,
            caption_embeds=caption_embeds,
            caption_padding_mask=caption_padding_mask,
        )
        eps, rest = model_out[:, : sample.shape[1]], model_out[:, sample.shape[1] :]
        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
        half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
        eps = torch.cat([half_eps, half_eps], dim=0)
        return torch.cat([eps, rest], dim=1)

class LayouSynPipeline(DiffusionPipeline):
    """Generate open-vocabulary scene layouts with LayouSyn.

    Args:
        model: Converted DiT denoiser.
        scheduler: LayouSyn Gaussian/DDIM scheduler.
        processor: Processor for prompt/concept inputs and postprocessing.
    """

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

    def __init__(
        self,
        model: LayouSynDiTModel,
        scheduler: LayouSynScheduler,
        processor: LayouSynProcessor,
    ) -> None:
        """Initialize the pipeline."""
        super().__init__()
        self._register_layousyn_modules(model, scheduler, processor)
        self.model.eval()

    def _register_layousyn_modules(
        self,
        model: LayouSynDiTModel,
        scheduler: LayouSynScheduler,
        processor: LayouSynProcessor,
    ) -> None:
        """Register pipeline modules and keep concrete attributes typed."""
        self.register_modules(model=model, scheduler=scheduler)
        self.model = model
        self.scheduler = scheduler
        self.processor = processor

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

    def save_pretrained(
        self,
        save_directory: str | os.PathLike[str],
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Save pipeline components plus processor metadata."""
        super().save_pretrained(save_directory, **kwargs)
        self.processor.save_pretrained(save_directory)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | os.PathLike[str],
        **kwargs: str | int | float | bool | None,
    ) -> LayouSynPipeline:
        """Load pipeline and restore local processor metadata."""
        pipe = super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        if not isinstance(pipe, cls):
            raise TypeError(f"Expected {cls.__name__}, got {type(pipe).__name__}")

        processor_config = (
            Path(pretrained_model_name_or_path) / LayouSynProcessor.config_name
        )
        if pipe.processor is None and processor_config.exists():
            pipe.processor = LayouSynProcessor.from_pretrained(
                pretrained_model_name_or_path
            )
        return pipe

    @torch.no_grad()
    def __call__(
        self,
        *,
        prompt: str | list[str] | None = None,
        batch_size: int = 1,
        seed: int | None = None,
        generator: torch.Generator | None = None,
        condition_type: ConditionType | str = ConditionType.text,
        labels: Int[torch.Tensor, "batch elements"]
        | Int[np.ndarray, "batch elements"]
        | list[str]
        | list[list[str]]
        | None = None,
        id2label: dict[int, str] | None = None,
        bbox: Float[torch.Tensor, "batch elements 4"]
        | Float[np.ndarray, "batch elements 4"]
        | Sequence[ArrayLikeInput]
        | None = None,
        mask: Bool[torch.Tensor, "batch elements"]
        | Bool[np.ndarray, "batch elements"]
        | Sequence[ArrayLikeInput]
        | None = None,
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None = None,
        box_format: BoxFormat | str = BoxFormat.xywh,
        normalized: bool = True,
        canvas_size: tuple[int, int] | None = None,
        aspect_ratio: float | list[float] | Float[torch.Tensor, "batch"] = 1.0,
        num_inference_steps: int | None = None,
        guidance_scale: float = 2.0,
        sampling_type: Literal["ddim", "ddpm"] = "ddim",
        output_type: Literal["dataclass", "dict"] = "dataclass",
        return_intermediates: bool = False,
        caption_embeds: Float[torch.Tensor, "batch tokens embedding_dim"] | None = None,
        caption_padding_mask: Bool[torch.Tensor, "batch tokens"] | None = None,
        concept_embeds: Float[torch.Tensor, "batch elements embedding_dim"]
        | None = None,
    ) -> LayoutGenerationOutput | LayouSynOutputDict:
        """Run LayouSyn denoising.

        Args:
            prompt: Caption text.
            batch_size: Number of generated layouts when labels are unbatched.
            seed: Convenience seed used only if ``generator`` is absent.
            generator: Exact reproducibility API.
            condition_type: Canonical condition name. First-class public mode is
                ``text``; unsupported modes fail explicitly.
            labels: String concepts or integer ids.
            id2label: Mapping for integer labels.
            bbox: Reserved for future initialization/refinement support.
            mask: Optional valid concept mask.
            num_elements: Optional expected element count. It is validated
                against labels when supplied.
            box_format: Public input bbox format.
            normalized: Whether input boxes are normalized.
            canvas_size: Required for pixel boxes.
            aspect_ratio: Scalar or per-example aspect ratio.
            num_inference_steps: Number of reverse diffusion steps.
            guidance_scale: Classifier-free guidance scale.
            sampling_type: ``ddim`` or ``ddpm``.
            output_type: ``dataclass`` or ``dict``.
            return_intermediates: Whether to return denoising trajectory.
            caption_embeds: Precomputed caption embeddings.
            caption_padding_mask: Precomputed caption padding mask.
            concept_embeds: Precomputed concept embeddings.

        Returns:
            Public layout output.
        """
        del batch_size
        canonical = normalize_condition_type(condition_type)
        if canonical is not ConditionType.text:
            raise NotImplementedError(
                f"LayouSyn public pipeline supports condition_type='text', got {condition_type}"
            )

        if generator is None and seed is not None:
            generator = torch.Generator(device=self.device).manual_seed(seed)
        encoded = self.processor(
            prompt=prompt,
            labels=labels,
            id2label=id2label,
            bbox=bbox,
            mask=mask,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
            aspect_ratio=aspect_ratio,
            caption_embeds=caption_embeds,
            caption_padding_mask=caption_padding_mask,
            concept_embeds=concept_embeds,
        )
        self._validate_num_elements(num_elements, encoded[LAYOUSYN_LABEL_TEXTS_KEY])
        concept_mask = encoded[LAYOUSYN_CONCEPT_MASK_KEY].to(self.device)
        batch = concept_mask.shape[0]
        self.scheduler.set_timesteps(num_inference_steps, device=self.device)
        sample = self.scheduler.initial_sample(
            batch,
            self.processor.max_in_len,
            self.model.config.in_channels,
            device=self.device,
            generator=generator,
        )
        model_kwargs = {
            "x_padding_mask": concept_mask,
            "aspect_ratio": encoded[LAYOUSYN_ASPECT_RATIO_KEY].to(self.device),
            "concept_embeds": encoded[LAYOUSYN_CONCEPT_EMBEDS_KEY].to(self.device),
            "caption_embeds": encoded[LAYOUSYN_CAPTION_EMBEDS_KEY].to(self.device),
            "caption_padding_mask": encoded[LAYOUSYN_CAPTION_MASK_KEY].to(self.device),
        }
        if guidance_scale != 1.0:
            sample = torch.cat([sample, sample], dim=0)
            model_kwargs = self._cfg_model_kwargs(model_kwargs, batch)
        trajectory = []
        for index, timestep in enumerate(self.scheduler.timesteps):
            model_timestep = self.scheduler.model_timesteps[index]
            t_model = torch.full(
                (sample.shape[0],),
                int(model_timestep),
                device=self.device,
                dtype=torch.long,
            )
            t_step = torch.full(
                (sample.shape[0],), int(timestep), device=self.device, dtype=torch.long
            )
            if guidance_scale != 1.0:
                model_output = self.model.forward_with_cfg(
                    sample, t_model, guidance_scale=guidance_scale, **model_kwargs
                )
            else:
                model_output = self.model(sample, t_model, **model_kwargs)
            step = self.scheduler.step(
                model_output,
                t_step,
                sample,
                generator=generator,
                sampling_type=sampling_type,
                clip_denoised=False,
            )
            sample = step.prev_sample
            if return_intermediates:
                trajectory.append(step.pred_original_sample[:batch].detach().cpu())
        sample = sample[:batch].clamp(-1.0, 1.0).detach().cpu()
        return self.processor.postprocess(
            sample,
            labels=encoded[LAYOUSYN_LABEL_TEXTS_KEY],
            id2label=encoded[LAYOUSYN_ID2LABEL_KEY],
            id2label_per_example=encoded[LAYOUSYN_PER_EXAMPLE_ID2LABEL_KEY],
            output_type=output_type,
            return_intermediates=return_intermediates,
            intermediates={"trajectory": trajectory} if return_intermediates else None,
        )

    generate = __call__

    def _cfg_model_kwargs(
        self, model_kwargs: dict[str, Shaped[torch.Tensor, ...]], batch_size: int
    ) -> dict[str, Shaped[torch.Tensor, ...]]:
        y_null = (
            self.model.y_embedder.y_embedding.to(self.device)
            .unsqueeze(0)
            .repeat(batch_size, 1, 1)
        )
        y_mask_null = (
            self.model.y_embedder.y_padding_mask.to(self.device)
            .unsqueeze(0)
            .repeat(batch_size, 1)
        )
        return {
            "x_padding_mask": torch.cat(
                [model_kwargs["x_padding_mask"], model_kwargs["x_padding_mask"]], dim=0
            ),
            "aspect_ratio": torch.cat(
                [model_kwargs["aspect_ratio"], model_kwargs["aspect_ratio"]], dim=0
            ),
            "concept_embeds": torch.cat(
                [model_kwargs["concept_embeds"], model_kwargs["concept_embeds"]], dim=0
            ),
            "caption_embeds": torch.cat(
                [model_kwargs["caption_embeds"], y_null], dim=0
            ),
            "caption_padding_mask": torch.cat(
                [model_kwargs["caption_padding_mask"], y_mask_null], dim=0
            ),
        }

    @staticmethod
    def _validate_num_elements(
        num_elements: int | list[int] | Int[torch.Tensor, "batch"] | None,
        labels: list[list[str]],
    ) -> None:
        if num_elements is None:
            return
        if isinstance(num_elements, int):
            expected = [num_elements] * len(labels)
        elif isinstance(num_elements, torch.Tensor):
            expected = [int(item) for item in num_elements.tolist()]
        else:
            expected = [int(item) for item in num_elements]
        actual = [len(row) for row in labels]
        if expected != actual:
            raise ValueError(f"num_elements {expected} does not match labels {actual}")