Posterllama

class PosterLlamaConfig(PretrainedConfig):
    """Configuration for local PosterLlama recipe artifacts.

    Args:
        checkpoint_repo_id: Source Hub repository containing the raw checkpoint.
        base_llm_repo_id: Preferred CodeLLaMA/LLaMA backbone repository id.
        alternate_base_llm_repo_ids: Alternate backbone ids recorded for audit.
        vision_encoder_repo_id: Vision encoder repository id.
        vision_model_name: Original vision tower selector.
        lora_r: LoRA rank used by the released recipe.
        lora_alpha: LoRA alpha used by the released recipe.
        lora_dropout: LoRA dropout used by the released recipe.
        lora_target_modules: LLM projection module names targeted by LoRA.
        prompt_template: Wrapper applied around generated layout prompts.
        image_placeholder: Placeholder token used for image feature insertion.
        image_end_token: End marker for image features.
        max_txt_len: Original maximum text length.
        max_context_len: Original context length budget.
        default_max_new_tokens: Default generation budget.
        default_do_sample: Default sampled-generation flag.
        default_temperature: Default generation temperature.
        default_top_p: Default nucleus sampling value.
        default_top_k: Default top-k sampling value.
        default_num_beams: Default beam count.
        dataset_name: Poster dataset key.
        id2label: Dataset-local label vocabulary.
        canvas_size: Optional default canvas size as ``(width, height)``.
        checkpoint_license_status: Redistribution status for converted weights.
        processor_subfolder: Pipeline processor subfolder.
        runtime_subfolder: Optional converted runtime subfolder.
        kwargs: Extra ``PretrainedConfig`` keyword arguments.

    Examples:
        >>> cfg = PosterLlamaConfig(canvas_size=(360, 504))
        >>> cfg.id2label[1]
        'text'
    """

    model_type = "posterllama"

    def __init__(
        self,
        checkpoint_repo_id: str = "poong/PosterLlama",
        base_llm_repo_id: str = "codellama/CodeLlama-7b-hf",
        alternate_base_llm_repo_ids: Sequence[str] = ("meta-llama/Llama-2-7b-chat-hf",),
        vision_encoder_repo_id: str = "facebook/dinov2-base",
        vision_model_name: PosterLlamaVisionModelName = "dino_v2",
        lora_r: int = 64,
        lora_alpha: int = 16,
        lora_dropout: float = 0.05,
        lora_target_modules: Sequence[str] = ("q_proj", "v_proj"),
        prompt_template: str = "{}",
        image_placeholder: str = "<ImageHere>",
        image_end_token: str = "</Img>",
        max_txt_len: int = 400,
        max_context_len: int = 3800,
        default_max_new_tokens: int = 1024,
        default_do_sample: bool = True,
        default_temperature: float = 0.6,
        default_top_p: float = 0.9,
        default_top_k: int = 40,
        default_num_beams: int = 4,
        dataset_name: PosterLlamaDatasetName = "cgl",
        id2label: Mapping[int | str, str] | None = None,
        canvas_size: tuple[int, int] | list[int] | None = None,
        checkpoint_license_status: PosterLlamaLicenseStatus = "unverified",
        processor_subfolder: str = "processor",
        runtime_subfolder: str = "runtime",
        **kwargs: PosterLlamaConfigValue,
    ) -> None:
        """Initialize configuration values."""
        labels = (
            id2label_for_dataset(dataset_name)
            if id2label is None
            else {int(key): str(value) for key, value in id2label.items()}
        )
        kwargs.pop("model_type", None)
        kwargs.pop("id2label", None)
        kwargs.pop("label2id", None)
        super().__init__(
            id2label=labels,
            label2id={label: idx for idx, label in labels.items()},
        )
        self.checkpoint_repo_id = checkpoint_repo_id
        self.base_llm_repo_id = base_llm_repo_id
        self.alternate_base_llm_repo_ids = list(alternate_base_llm_repo_ids)
        self.vision_encoder_repo_id = vision_encoder_repo_id
        self.vision_model_name = vision_model_name

        self.lora_r = int(lora_r)
        self.lora_alpha = int(lora_alpha)
        self.lora_dropout = float(lora_dropout)
        self.lora_target_modules = list(lora_target_modules)

        self.prompt_template = prompt_template
        self.image_placeholder = image_placeholder
        self.image_end_token = image_end_token

        self.max_txt_len = int(max_txt_len)
        self.max_context_len = int(max_context_len)

        self.default_max_new_tokens = int(default_max_new_tokens)
        self.default_do_sample = bool(default_do_sample)
        self.default_temperature = float(default_temperature)
        self.default_top_p = float(default_top_p)
        self.default_top_k = int(default_top_k)
        self.default_num_beams = int(default_num_beams)

        self.dataset_name = dataset_name
        self.canvas_size = tuple(canvas_size) if canvas_size is not None else None
        self.checkpoint_license_status = checkpoint_license_status
        self.processor_subfolder = processor_subfolder
        self.runtime_subfolder = runtime_subfolder

        for key, value in kwargs.items():
            setattr(self, key, value)

class PosterLlamaImageProcessor(BaseImageProcessor):
    """Prepare RGB images for PosterLlama smoke and recipe paths.

    Args:
        image_size: Optional ``(height, width)`` resize target.
        vision_encoder_repo_id: Vision encoder id recorded with the processor.

    Examples:
        >>> processor = PosterLlamaImageProcessor(image_size=(8, 8))
        >>> out = processor.preprocess(torch.zeros(3, 8, 8))
        >>> tuple(out["pixel_values"].shape)
        (1, 3, 8, 8)
    """

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        image_size: tuple[int, int] | None = None,
        vision_encoder_repo_id: str = "facebook/dinov2-base",
        **kwargs: PosterLlamaImageProcessorKwarg,
    ) -> None:
        """Initialize image processor metadata."""
        super().__init__(**kwargs)  # ty: ignore[invalid-argument-type]
        self.image_size = tuple(image_size) if image_size is not None else None
        self.vision_encoder_repo_id = vision_encoder_repo_id

    def preprocess(
        self,
        images: ImageInput | Sequence[ImageInput] | None,
        return_tensors: Literal["pt"] = "pt",
        **kwargs: PosterLlamaImageProcessorKwarg,
    ) -> BatchFeature:
        """Convert images to tensors.

        Args:
            images: PIL, NumPy, or torch image inputs. Omitted images create a
                zero placeholder for parser-only smoke calls.
            return_tensors: Tensor return format. Only ``pt`` is supported.
            kwargs: Reserved image-processing options.

        Returns:
            BatchFeature containing ``pixel_values``.

        Raises:
            ValueError: If ``return_tensors`` is not ``pt``.
        """
        _ = kwargs
        if return_tensors != "pt":
            raise ValueError("PosterLlamaImageProcessor supports return_tensors='pt'")

        items = (
            _as_image_list(images) if images is not None else [torch.zeros(3, 64, 64)]
        )
        pixel_values = torch.stack([_image_to_tensor(item) for item in items])
        if self.image_size is not None:
            pixel_values = torch.nn.functional.interpolate(
                pixel_values,
                size=self.image_size,
                mode="bilinear",
                align_corners=False,
            )
        return BatchFeature({"pixel_values": pixel_values}, tensor_type=return_tensors)

    def to_dict(self) -> dict[str, str | tuple[int, int] | None]:
        """Serialize image processor metadata."""
        data = cast(dict[str, str | tuple[int, int] | None], super().to_dict())
        data["image_size"] = self.image_size
        data["vision_encoder_repo_id"] = self.vision_encoder_repo_id
        return data

class PosterLlamaPipeline(LayoutGenerationPipeline):
    """Compose a PosterLlama processor and converted runtime.

    Args:
        config: Explicit pipeline configuration.
        processor: Explicit processor.
        runtime: Optional converted runtime. Parser-only saved artifacts may omit it.

    Examples:
        >>> cfg = PosterLlamaConfig(canvas_size=(100, 100))
        >>> text = '<svg width="100" height="100"><rect data-category="text" x="0" y="0" width="10" height="10"/></svg>'
        >>> pipe = PosterLlamaPipeline(
        ...     config=cfg,
        ...     processor=PosterLlamaProcessor.from_config(cfg),
        ...     runtime=PosterLlamaRuntime(text),
        ... )
        >>> pipe(images=None).bbox.shape
        torch.Size([1, 1, 4])
    """

    config_class: ClassVar[type[PretrainedConfig]] = PosterLlamaConfig
    component_specs: ClassVar[dict[str, PipelineComponentSpec]] = {
        "processor": PipelineComponentSpec(
            attribute_name="processor",
            loader=_load_processor_component,
            config_subfolder_attribute="processor_subfolder",
            marker_file="processor_config.json",
            save_with_is_main_process=False,
        ),
        "runtime": PipelineComponentSpec(
            attribute_name="runtime",
            loader=_load_runtime_component,
            config_subfolder_attribute="runtime_subfolder",
            required=False,
            marker_file="runtime_config.json",
        ),
    }

    config: PosterLlamaConfig
    processor: PosterLlamaProcessor
    runtime: PosterLlamaRuntime | None

    def __init__(
        self,
        *,
        config: PosterLlamaConfig,
        processor: PosterLlamaProcessor,
        runtime: PosterLlamaRuntime | None = None,
    ) -> None:
        """Initialize pipeline components."""
        super().__init__(config)
        self.config = config
        self.processor = processor
        self.runtime = runtime

    @classmethod
    def _from_pretrained_components(  # ty: ignore[invalid-method-override]
        cls,
        *,
        config: PretrainedConfig,
        components: Mapping[str, PosterLlamaProcessor | PosterLlamaRuntime | None],
    ) -> "PosterLlamaPipeline":
        """Build a pipeline from loaded components."""
        return cls(
            config=cast(PosterLlamaConfig, config),
            processor=cast(PosterLlamaProcessor, components["processor"]),
            runtime=cast(PosterLlamaRuntime | None, components["runtime"]),
        )

    @torch.no_grad()
    def __call__(  # ty: ignore[invalid-method-override]
        self,
        *,
        images: ImageInput | Sequence[ImageInput] | None = None,
        prompt: str | Sequence[str] | None = None,
        content: Mapping[str, str | int | float | bool | None]
        | Sequence[Mapping[str, str | int | float | bool | None]]
        | None = None,
        texts: str | Sequence[str] | Sequence[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, "..."]
        | Sequence[Sequence[int | str]]
        | Sequence[int | str]
        | None = None,
        bbox: Float[torch.Tensor, "..."]
        | Sequence[Sequence[Sequence[float | int]]]
        | Sequence[Sequence[float | int]]
        | Sequence[float | int]
        | None = None,
        mask: Bool[torch.Tensor, "..."]
        | 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,
        max_new_tokens: int | None = None,
        do_sample: bool | None = None,
        temperature: float | None = None,
        top_p: float | None = None,
        top_k: int | None = None,
        num_beams: int | None = None,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | Mapping[str, str | bytes | int | float | bool | None]
            | list[str]
            | list[tuple[float, float, float, float]]
            | tuple[int, int]
            | str
            | bytes
            | int
            | float
            | bool
            | None,
        ]
    ):
        """Generate a poster layout.

        Args:
            images: Poster image inputs.
            prompt: Optional prompt prefix override.
            content: Optional content metadata.
            texts: Optional poster text strings.
            batch_size: Batch size when images are omitted.
            seed: Convenience seed used only when ``generator`` is absent.
            generator: Explicit PyTorch generator; takes precedence over ``seed``.
            condition_type: Canonical condition or PosterLlama release alias.
            labels: Optional label constraints.
            bbox: Optional box constraints.
            mask: Optional valid-element mask.
            num_elements: Optional requested element count.
            box_format: Input box format.
            normalized: Whether boxes are normalized.
            canvas_size: Canvas size as ``(width, height)``.
            num_inference_steps: Reserved shared argument.
            output_type: ``dataclass`` or ``dict``.
            return_intermediates: Whether to include prompt and parse diagnostics.
            max_new_tokens: Generation token budget.
            do_sample: Sampling flag.
            temperature: Sampling temperature.
            top_p: Nucleus sampling value.
            top_k: Top-k sampling value.
            num_beams: Beam count.

        Returns:
            LayoutGenerationOutput or dictionary.

        Raises:
            RuntimeError: If converted runtime assets are absent.
        """
        _ = num_inference_steps
        batch = self.processor(
            images=images,
            prompt=prompt,
            content=content,
            texts=texts,
            batch_size=batch_size,
            condition_type=condition_type,
            labels=labels,
            bbox=bbox,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        if self.runtime is None:
            raise RuntimeError(
                "PosterLlama runtime assets are missing. A processor-only artifact "
                "can build prompts and parse outputs, but generation requires a "
                "converted local runtime."
            )

        active_generator = self.prepare_generator(generator=generator, seed=seed)
        generation_args = {
            "max_new_tokens": max_new_tokens or self.config.default_max_new_tokens,
            "do_sample": self.config.default_do_sample
            if do_sample is None
            else do_sample,
            "temperature": self.config.default_temperature
            if temperature is None
            else temperature,
            "top_p": self.config.default_top_p if top_p is None else top_p,
            "top_k": self.config.default_top_k if top_k is None else top_k,
            "num_beams": self.config.default_num_beams
            if num_beams is None
            else num_beams,
        }
        generated = self.runtime.generate_texts(
            cast(list[str], batch["prompts"]),
            pixel_values=cast(torch.Tensor, batch["pixel_values"]),
            generator=active_generator,
            **generation_args,
        )
        output = self.processor.parse_output(
            generated[0],
            canvas_size=canvas_size,
            output_type="dataclass",
            return_intermediates=return_intermediates,
        )
        result = cast(LayoutGenerationOutput, output)
        if return_intermediates:
            existing = cast(
                dict[
                    str,
                    str
                    | bytes
                    | int
                    | float
                    | bool
                    | None
                    | Mapping[str, str | int | float | bool | None],
                ],
                result.intermediates or {},
            )
            existing["prompt_bytes"] = cast(list[str], batch["prompts"])[0].encode()
            existing["raw_generated_text"] = generated[0]
            existing["condition_type"] = str(batch["condition_type"])
            existing["generation_args"] = generation_args
            result.intermediates = existing
        if output_type == "dict":
            return dict(result)
        if output_type != "dataclass":
            raise ValueError(f"Unsupported output_type: {output_type}")

        return result

class PosterLlamaProcessor(ProcessorMixin):
    """Build PosterLlama prompts and decode generated HTML/SVG layouts.

    Args:
        image_processor: Image processor metadata wrapper.
        config: Explicit PosterLlama configuration.

    Examples:
        >>> processor = PosterLlamaProcessor.from_config(PosterLlamaConfig())
        >>> "Generate poster layout" in processor.build_prompt(condition_type="unconditional")
        True
    """

    attributes = ["image_processor"]
    image_processor_class = "PosterLlamaImageProcessor"

    def __init__(
        self,
        image_processor: PosterLlamaImageProcessor,
        config: PosterLlamaConfig,
    ) -> None:
        """Initialize processor components."""
        self.image_processor = image_processor
        self.config = config
        super().__init__(image_processor)

    @classmethod
    def from_config(cls, config: PosterLlamaConfig) -> "PosterLlamaProcessor":
        """Create a processor from an explicit config.

        Args:
            config: Processor configuration.

        Returns:
            PosterLlamaProcessor instance.

        Examples:
            >>> PosterLlamaProcessor.from_config(PosterLlamaConfig()).config.model_type
            'posterllama'
        """
        return cls(
            image_processor=PosterLlamaImageProcessor(
                vision_encoder_repo_id=config.vision_encoder_repo_id,
            ),
            config=config,
        )

    def save_pretrained(
        self,
        save_directory: str | PathLike[str],
        push_to_hub: bool = False,
        **kwargs: str | int | float | bool | None,
    ) -> None:
        """Save processor metadata.

        Args:
            save_directory: Directory to write.
            push_to_hub: Accepted for ProcessorMixin compatibility; ignored.
            kwargs: Accepted for ProcessorMixin compatibility; ignored.
        """
        _ = (push_to_hub, kwargs)
        root = Path(save_directory)
        root.mkdir(parents=True, exist_ok=True)
        self.config.save_pretrained(root)
        self.image_processor.save_pretrained(root)
        (root / "processor_config.json").write_text(
            json.dumps(
                {
                    "processor_class": self.__class__.__name__,
                    "image_processor_class": self.image_processor.__class__.__name__,
                },
                indent=2,
                sort_keys=True,
            ),
            encoding="utf-8",
        )

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | PathLike[str],
        cache_dir: str | PathLike[str] | None = None,
        force_download: bool = False,
        local_files_only: bool = False,
        token: str | bool | None = None,
        revision: str = "main",
        *,
        subfolder: str | None = None,
        **kwargs: PosterLlamaImageProcessorKwarg,
    ) -> "PosterLlamaProcessor":
        """Load processor metadata.

        Args:
            pretrained_model_name_or_path: Checkpoint root or processor folder.
            cache_dir: Accepted for ProcessorMixin compatibility.
            force_download: Accepted for ProcessorMixin compatibility.
            local_files_only: Whether to avoid network access.
            token: Accepted for ProcessorMixin compatibility.
            revision: Accepted for ProcessorMixin compatibility.
            subfolder: Optional processor subfolder.
            kwargs: Accepted for ProcessorMixin compatibility.

        Returns:
            Loaded PosterLlamaProcessor.
        """
        _ = (cache_dir, force_download, token, revision, kwargs)
        root = Path(pretrained_model_name_or_path)
        if subfolder is not None:
            root = root / subfolder
        config = PosterLlamaConfig.from_pretrained(
            root, local_files_only=local_files_only
        )
        return cls(
            image_processor=PosterLlamaImageProcessor.from_pretrained(root),
            config=config,
        )

    def normalize_condition_type(
        self,
        condition_type: ConditionType | str,
    ) -> ConditionType:
        """Normalize PosterLlama condition aliases.

        Args:
            condition_type: Canonical condition or PosterLlama release alias.

        Returns:
            Supported canonical condition.

        Raises:
            NotImplementedError: If the condition is known but unsupported.
            ValueError: If the condition is unknown.
        """
        if isinstance(condition_type, str):
            alias = condition_type.lower().replace("-", "_")
            if alias in REQUEST_CONDITION_ALIASES:
                return REQUEST_CONDITION_ALIASES[alias]
        condition = normalize_condition_type(condition_type)
        if condition in UNSUPPORTED_CONDITIONS:
            raise NotImplementedError(f"PosterLlama does not support {condition}")

        if condition not in SUPPORTED_CONDITIONS:
            raise NotImplementedError(f"PosterLlama does not support {condition}")

        return condition

    def __call__(
        self,
        *,
        images: ImageInput | Sequence[ImageInput] | None = None,
        prompt: str | Sequence[str] | None = None,
        content: Mapping[str, str | int | float | bool | None]
        | Sequence[Mapping[str, str | int | float | bool | None]]
        | None = None,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None = None,
        batch_size: int = 1,
        condition_type: ConditionType | str = ConditionType.content_image,
        labels: Int[torch.Tensor, "..."]
        | Sequence[Sequence[int | str]]
        | Sequence[int | str]
        | None = None,
        bbox: Float[torch.Tensor, "..."]
        | Sequence[Sequence[Sequence[float | int]]]
        | Sequence[Sequence[float | int]]
        | Sequence[float | int]
        | None = None,
        mask: Bool[torch.Tensor, "..."]
        | 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,
        return_tensors: Literal["pt"] = "pt",
    ) -> BatchEncoding:
        """Encode public inputs into recipe prompt and image tensors.

        Args:
            images: Poster image inputs.
            prompt: Optional user-provided prompt text.
            content: Optional content metadata.
            texts: Optional poster text strings.
            batch_size: Batch size when no images are supplied.
            condition_type: Canonical condition or PosterLlama release alias.
            labels: Optional element label constraints.
            bbox: Optional element boxes.
            mask: Optional valid-element mask.
            num_elements: Optional requested element count.
            box_format: Input box format.
            normalized: Whether boxes are normalized.
            canvas_size: Canvas size for pixel boxes and prompt rendering.
            return_tensors: Tensor return format. Only ``pt`` is supported.

        Returns:
            BatchEncoding containing prompt strings and tensors.
        """
        condition = self.normalize_condition_type(condition_type)
        canvas = self._resolve_canvas_size(canvas_size)
        prompt_text = self.build_prompt(
            condition_type=condition,
            prompt=prompt,
            content=content,
            texts=texts,
            labels=labels,
            bbox=bbox,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas,
        )
        image_batch = self.image_processor(images, return_tensors=return_tensors)
        return BatchEncoding(
            {
                "pixel_values": image_batch["pixel_values"],
                "prompts": prompt_text
                if isinstance(prompt_text, list)
                else [prompt_text] * batch_size,
                "condition_type": condition,
                "canvas_size": canvas,
            }
        )

    def build_prompt(
        self,
        *,
        condition_type: ConditionType | str = ConditionType.content_image,
        prompt: str | Sequence[str] | None = None,
        content: Mapping[str, str | int | float | bool | None]
        | Sequence[Mapping[str, str | int | float | bool | None]]
        | None = None,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None = None,
        labels: Int[torch.Tensor, "..."]
        | Sequence[Sequence[int | str]]
        | Sequence[int | str]
        | None = None,
        bbox: Float[torch.Tensor, "..."]
        | Sequence[Sequence[Sequence[float | int]]]
        | Sequence[Sequence[float | int]]
        | Sequence[float | int]
        | None = None,
        mask: Bool[torch.Tensor, "..."]
        | 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,
    ) -> str | list[str]:
        """Build a deterministic PosterLlama HTML prompt.

        Args:
            condition_type: Canonical condition or PosterLlama release alias.
            prompt: Optional prompt prefix override.
            content: Optional content metadata included in diagnostics text.
            texts: Optional poster text strings.
            labels: Optional label constraints.
            bbox: Optional box constraints.
            mask: Optional valid-element mask.
            num_elements: Requested element count for unconstrained slots.
            box_format: Input box format.
            normalized: Whether boxes are normalized.
            canvas_size: Canvas size as ``(width, height)``.

        Returns:
            Prompt string or list of prompt strings.
        """
        _ = content
        condition = self.normalize_condition_type(condition_type)
        canvas = self._resolve_canvas_size(canvas_size)
        base_prompt = self._first_prompt(prompt)
        text_line = self._texts_line(texts)
        known_markup = self._constraint_markup(
            condition=condition,
            labels=labels,
            bbox=bbox,
            mask=mask,
            num_elements=num_elements,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas,
        )
        bbox_html = HTML_TEMPLATE.format(
            width=canvas[0],
            height=canvas[1],
            content=known_markup,
        )
        source_key = self._source_condition_key(condition)
        task_instruction = SOURCE_TASK_INSTRUCTIONS[self.config.dataset_name]
        if text_line:
            instruction = SOURCE_TEXT_INSTRUCTIONS[source_key].format(
                text=text_line,
                bbox_html=bbox_html,
            )
        else:
            instruction = SOURCE_INSTRUCTIONS[source_key].format(bbox_html=bbox_html)
        body = f"{base_prompt}{task_instruction}{instruction} <MID>"
        return self.config.prompt_template.format(body)

    def parse_output(
        self,
        text: str,
        *,
        canvas_size: tuple[int, int] | None = None,
        output_type: Literal["dataclass", "dict"] = "dataclass",
        return_intermediates: bool = False,
        strict: bool = False,
    ) -> (
        LayoutGenerationOutput
        | dict[
            str,
            Float[torch.Tensor, "..."]
            | Int[torch.Tensor, "..."]
            | Bool[torch.Tensor, "..."]
            | dict[int, str]
            | list[str]
            | list[tuple[float, float, float, float]]
            | tuple[int, int]
            | str
            | int
            | float
            | bool
            | None,
        ]
    ):
        """Parse generated HTML/SVG into public layout output.

        Args:
            text: Generated markup.
            canvas_size: Canvas size override.
            output_type: Return dataclass or dictionary.
            return_intermediates: Whether to include parse diagnostics.
            strict: Whether malformed rectangles raise.

        Returns:
            LayoutGenerationOutput or dictionary.
        """
        parsed = parse_rectangles(text, self._label2id(), strict=strict)
        canvas = canvas_size or parsed.canvas_size or self.config.canvas_size
        if canvas is None:
            raise ValueError(
                "canvas_size is required when generated SVG lacks width/height"
            )

        resolved_canvas = (int(canvas[0]), int(canvas[1]))
        output = rect_ltwh_to_output(
            parsed,
            canvas_size=resolved_canvas,
            id2label=cast(dict[int, str], self.config.id2label),
            return_intermediates=return_intermediates,
        )
        if output_type == "dict":
            return dict(output)
        if output_type != "dataclass":
            raise ValueError(f"Unsupported output_type: {output_type}")

        return output

    def _resolve_canvas_size(
        self,
        canvas_size: tuple[int, int] | None,
    ) -> tuple[int, int]:
        canvas = canvas_size or self.config.canvas_size
        if canvas is None:
            return (360, 504)
        return int(canvas[0]), int(canvas[1])

    def _first_prompt(self, prompt: str | Sequence[str] | None) -> str:
        if prompt is None:
            return ""
        if isinstance(prompt, str):
            return prompt
        return next(iter(prompt), "")

    def _texts_line(
        self,
        texts: str | Sequence[str] | Sequence[Sequence[str]] | None,
    ) -> str:
        if texts is None:
            return ""
        if isinstance(texts, str):
            values = [texts]
        else:
            first = next(iter(texts), "")
            if isinstance(first, str):
                values = cast(list[str], list(texts))
            else:
                values = [str(value) for value in first]
        return " | ".join(values)

    def _constraint_markup(
        self,
        *,
        condition: ConditionType,
        labels: Int[torch.Tensor, "..."]
        | Sequence[Sequence[int | str]]
        | Sequence[int | str]
        | None,
        bbox: Float[torch.Tensor, "..."]
        | Sequence[Sequence[Sequence[float | int]]]
        | Sequence[Sequence[float | int]]
        | Sequence[float | int]
        | None,
        mask: Bool[torch.Tensor, "..."]
        | Sequence[Sequence[bool]]
        | Sequence[bool]
        | None,
        num_elements: int | Sequence[int] | Int[torch.Tensor, "batch"] | None,
        box_format: BoxFormat | str,
        normalized: bool,
        canvas_size: tuple[int, int],
    ) -> str:
        if labels is None:
            if condition in {ConditionType.content_image, ConditionType.unconditional}:
                return ""
            count = self._num_elements(num_elements)
            return " ".join(self._fill_rect(index) for index in range(count))
        label_tensor = self._labels_to_tensor(labels)
        bbox_tensor = self._bbox_to_tensor(
            bbox,
            labels=label_tensor,
            mask=mask,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        rects: list[str] = []
        fill_index = 1
        for index, (label, box) in enumerate(
            zip(label_tensor[0].tolist(), bbox_tensor[0].tolist(), strict=True)
        ):
            rect, fill_index = self._rect_for_condition(
                condition=condition,
                label=int(label),
                bbox_ltwh=(
                    float(box[0]),
                    float(box[1]),
                    float(box[2]),
                    float(box[3]),
                ),
                fill_index=fill_index,
            )
            if rect:
                rects.append(rect)
        return " ".join(rects)

    def _labels_to_tensor(
        self,
        labels: Int[torch.Tensor, "..."]
        | Sequence[Sequence[int | str]]
        | Sequence[int | str],
    ) -> Int[torch.Tensor, "batch elements"]:
        if isinstance(labels, torch.Tensor):
            tensor = labels.long()
            return tensor.unsqueeze(0) if tensor.ndim == 1 else tensor
        rows = cast(Sequence[int | str | Sequence[int | str]], labels)
        if rows and isinstance(rows[0], Sequence) and not isinstance(rows[0], str):
            row = cast(Sequence[int | str], rows[0])
        else:
            row = cast(Sequence[int | str], rows)
        label2id = self._label2id()
        values = [
            label2id[self._normalize_input_label(item)]
            if isinstance(item, str)
            else int(item)
            for item in row
        ]
        return torch.tensor([values], dtype=torch.long)

    def _bbox_to_tensor(
        self,
        bbox: Float[torch.Tensor, "..."]
        | Sequence[Sequence[Sequence[float | int]]]
        | Sequence[Sequence[float | int]]
        | Sequence[float | int]
        | None,
        *,
        labels: Int[torch.Tensor, "batch elements"],
        mask: Bool[torch.Tensor, "..."]
        | Sequence[Sequence[bool]]
        | Sequence[bool]
        | None,
        box_format: BoxFormat | str,
        normalized: bool,
        canvas_size: tuple[int, int],
    ) -> Float[torch.Tensor, "batch elements 4"]:
        if bbox is None:
            return torch.zeros((labels.size(0), labels.size(1), 4), dtype=torch.float32)
        bbox_t, _, _ = prepare_layout_tensors(
            bbox=bbox,
            labels=labels,
            mask=mask,
            box_format=box_format,
            normalized=normalized,
            canvas_size=canvas_size,
        )
        if bbox_t.ndim == 2:
            bbox_t = bbox_t.unsqueeze(0)
        return denormalize_boxes(bbox_t, canvas_size=canvas_size, box_format="ltwh")

    def _rect_for_condition(
        self,
        *,
        condition: ConditionType,
        label: int,
        bbox_ltwh: tuple[float, float, float, float],
        fill_index: int,
    ) -> tuple[str, int]:
        label_name = str(cast(dict[int, str], self.config.id2label)[label])
        x, y, width, height = bbox_ltwh

        if condition is ConditionType.label:
            x, y, width, height = self._fill_values(fill_index, 4)
            fill_index += 4

        elif condition is ConditionType.label_size:
            x, y = self._fill_values(fill_index, 2)
            fill_index += 2

        elif condition is ConditionType.completion:
            x, y, width, height = self._fill_values(fill_index, 4)
            fill_index += 4

        elif condition is ConditionType.refinement:
            width, height = self._fill_values(fill_index, 2)
            fill_index += 2

        elif condition in {ConditionType.content_image, ConditionType.unconditional}:
            return "", fill_index

        x, y, width, height = (
            _format_source_number(value) for value in (x, y, width, height)
        )
        return (
            RECT_TEMPLATE.format(
                label=label_name,
                x=x,
                y=y,
                width=width,
                height=height,
            ),
            fill_index,
        )

    def _fill_rect(self, index: int) -> str:
        label, x, y, width, height = self._fill_values(index + 1, 5)
        return RECT_TEMPLATE.format(label=label, x=x, y=y, width=width, height=height)

    def _fill_values(self, start: int, count: int) -> tuple[str, ...]:
        return tuple(
            FILL_TEMPLATE.format(index) for index in range(start, start + count)
        )

    def _source_condition_key(self, condition: ConditionType) -> str:
        if condition is ConditionType.label:
            return "cond_cate_to_size_pos"
        if condition is ConditionType.label_size:
            return "cond_cate_size_to_pos"
        if condition is ConditionType.completion:
            return "cond_random_mask"
        if condition is ConditionType.refinement:
            return "cond_cate_pos_to_size"
        return "unconditional"

    def _num_elements(
        self,
        num_elements: int | Sequence[int] | Int[torch.Tensor, "batch"] | None,
    ) -> int:
        if num_elements is None:
            return 1
        if isinstance(num_elements, torch.Tensor):
            return int(num_elements.flatten()[0].item())
        if isinstance(num_elements, Sequence) and not isinstance(num_elements, str):
            return int(cast(int, num_elements[0]))
        return int(cast(int, num_elements))

    def _label2id(self) -> dict[str, int]:
        return {
            self._normalize_input_label(label): int(index)
            for index, label in cast(dict[int, str], self.config.id2label).items()
        }

    def _normalize_input_label(self, label: int | str) -> str:
        return str(label).strip().lower().replace("_", " ")