diffusers/pipelines/ideogram4/pipeline_ideogram4.py

# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
from typing import Any, Callable

import torch
from transformers import AutoTokenizer, PreTrainedModel
from transformers.masking_utils import create_causal_mask

from ...image_processor import VaeImageProcessor
from ...loaders import Ideogram4LoraLoaderMixin
from ...models.autoencoders import AutoencoderKLFlux2
from ...models.transformers.transformer_ideogram4 import (
    IMAGE_POSITION_OFFSET,
    LLM_TOKEN_INDICATOR,
    OUTPUT_IMAGE_INDICATOR,
    SEQUENCE_PADDING_INDICATOR,
    Ideogram4Transformer2DModel,
)
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import is_outlines_available, logging, replace_example_docstring
from ...utils.torch_utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline
from .pipeline_output import Ideogram4PipelineOutput
from .prompt_enhancer import (
    PROMPT_UPSAMPLE_TEMPERATURE,
    Ideogram4PromptEnhancerHead,
    build_caption_logits_processor,
    build_prompt_enhancer,
    generate_captions,
)


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


# Hidden states of these Qwen3-VL decoder layers are concatenated to form the per-token
# text conditioning consumed by the Ideogram4 transformer.
QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)


EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> import torch
        >>> from diffusers import Ideogram4Pipeline

        >>> pipe = Ideogram4Pipeline.from_pretrained("ideogram-ai/ideogram-v4", torch_dtype=torch.bfloat16)
        >>> pipe.to("cuda")

        >>> prompt = "A photo of a cat holding a sign that says hello world"
        >>> # The defaults are the recommended settings for best quality.
        >>> image = pipe(prompt, height=2048, width=2048, generator=torch.Generator("cuda").manual_seed(0)).images[0]
        >>> image.save("ideogram4.png")
        ```
"""


def _logit_normal_sigmas(
    num_inference_steps: int,
    mu: float,
    std: float = 1.0,
    logsnr_min: float = -15.0,
    logsnr_max: float = 18.0,
    device: torch.device | None = None,
) -> torch.Tensor:
    r"""
    Build a length-`num_inference_steps` sigma schedule using the Ideogram4 logit-normal flow-matching schedule.

    Sigmas are returned in `[0, 1]` in decreasing order (sigma close to 1 corresponds to pure noise, sigma close to 0
    to clean data), matching diffusers conventions.

    The Ideogram4 schedule applies `sigma(s) = 1 - logit_normal_cdf_inverse(1 - s)` to `s = linspace(0, 1, N + 1)` and
    keeps the first `N` entries; a terminal zero is appended downstream by the scheduler.
    """
    intervals = torch.linspace(0.0, 1.0, num_inference_steps + 1, dtype=torch.float64)
    # Apply the inverse CDF of a normal then push through the logistic to obtain a logit-normal CDF inverse.
    z = torch.special.ndtri(intervals)
    y = mu + std * z
    t = 1.0 - torch.special.expit(y)
    t_min = 1.0 / (1.0 + math.exp(0.5 * logsnr_max))
    t_max = 1.0 / (1.0 + math.exp(0.5 * logsnr_min))
    t = t.clamp(t_min, t_max)
    # Convert from model time (0 = noise, 1 = data) to diffusers sigma (1 = noise, 0 = data) and reverse.
    sigmas = (1.0 - t).flip(0)
    # Drop the trailing 0; FlowMatchEulerDiscreteScheduler.set_timesteps appends one back internally.
    sigmas = sigmas[:-1].to(dtype=torch.float32, device=device)
    return sigmas


def _resolution_aware_mu(
    height: int,
    width: int,
    base_mu: float,
    base_resolution: tuple[int, int] = (512, 512),
) -> float:
    """Shift the schedule mean as a function of image resolution."""
    num_pixels = height * width
    base_pixels = base_resolution[0] * base_resolution[1]
    return base_mu + 0.5 * math.log(num_pixels / base_pixels)


def _expand_tensor_to_effective_batch(
    tensor: torch.Tensor,
    batch_size: int,
    num_per_prompt: int,
    tensor_name: str | None = None,
) -> torch.Tensor:
    """Replicate `tensor` along dim 0 from `batch_size` (or 1) to `batch_size * num_per_prompt`."""
    target_batch_size = batch_size * num_per_prompt

    if tensor.shape[0] == target_batch_size:
        return tensor

    if tensor.shape[0] == 1:
        repeat_by = target_batch_size
    elif tensor.shape[0] == batch_size:
        repeat_by = num_per_prompt
    else:
        tensor_name = f"`{tensor_name}`" if tensor_name is not None else "Tensor"
        raise ValueError(
            f"{tensor_name} batch size must be 1, `batch_size` ({batch_size}), or "
            f"`batch_size * num_*_per_prompt` ({target_batch_size}), but got {tensor.shape[0]}."
        )

    return torch.repeat_interleave(tensor, repeats=repeat_by, dim=0, output_size=tensor.shape[0] * repeat_by)


class Ideogram4Pipeline(DiffusionPipeline, Ideogram4LoraLoaderMixin):
    r"""
    Text-to-image pipeline for Ideogram4.

    Ideogram4 is a flow-matching model trained with asymmetric classifier-free guidance: a `transformer` consumes
    text-conditioned features alongside the image latents, while a separate `unconditional_transformer` denoises with
    zeroed text features. The two velocity predictions are linearly blended each step.

    Args:
        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
            Flow-matching scheduler. The pipeline overrides the default sigma schedule with a resolution-aware
            logit-normal schedule.
        vae ([`AutoencoderKLFlux2`]):
            Variational auto-encoder used to decode latents back into images.
        text_encoder ([`PreTrainedModel`]):
            Multimodal text encoder. The pipeline consumes hidden states from a fixed set of intermediate decoder
            layers (see `QWEN3_VL_ACTIVATION_LAYERS`).
        tokenizer ([`AutoTokenizer`]):
            Tokenizer paired with `text_encoder`.
        transformer ([`Ideogram4Transformer2DModel`]):
            Conditional flow-matching transformer.
        unconditional_transformer ([`Ideogram4Transformer2DModel`]):
            Unconditional (asymmetric-CFG) flow-matching transformer.
    """

    model_cpu_offload_seq = "prompt_enhancer_head->text_encoder->transformer->unconditional_transformer->vae"
    _optional_components = ["prompt_enhancer_head"]
    _callback_tensor_inputs = ["latents"]

    def __init__(
        self,
        scheduler: FlowMatchEulerDiscreteScheduler,
        vae: AutoencoderKLFlux2,
        text_encoder: PreTrainedModel,
        tokenizer: AutoTokenizer,
        transformer: Ideogram4Transformer2DModel,
        unconditional_transformer: Ideogram4Transformer2DModel,
        prompt_enhancer_head: Ideogram4PromptEnhancerHead | None = None,
    ) -> None:
        super().__init__()

        self.register_modules(
            scheduler=scheduler,
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            transformer=transformer,
            unconditional_transformer=unconditional_transformer,
            prompt_enhancer_head=prompt_enhancer_head,
        )

        self.vae_scale_factor = (
            2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) is not None else 8
        )
        # Ideogram4 patchifies the VAE output by a factor of 2 before feeding into the transformer.
        self.patch_size = 2
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size)

        # Built lazily on first upsample: the head-less encoder body + `prompt_enhancer_head`, combined.
        self._prompt_enhancer = None
        # Outlines logits processor for schema-constrained captions; built lazily on first upsample.
        self._caption_logits_processor = None

    def upsample_prompt(
        self,
        prompt: str | list[str],
        height: int = 2048,
        width: int = 2048,
        temperature: float = PROMPT_UPSAMPLE_TEMPERATURE,
        max_new_tokens: int = 1024,
        generator: torch.Generator | list[torch.Generator] | None = None,
        device: torch.device | None = None,
    ) -> list[str]:
        """Rewrite each prompt into Ideogram4's native structured JSON caption.

        Requires the optional `prompt_enhancer_head` component, which is grafted onto the shared `text_encoder` body to
        make it generative. Generation is schema-constrained when `outlines` is installed, otherwise it runs
        unconstrained. Pass `generator` (the same one accepted by `__call__`) to make sampling reproducible.
        """
        if self.prompt_enhancer_head is None:
            raise ValueError(
                "Prompt upsampling requires the `prompt_enhancer_head` component, which is not loaded. Load it and "
                "pass it in, e.g.:\n"
                "    from diffusers import Ideogram4PromptEnhancerHead\n"
                "    head = Ideogram4PromptEnhancerHead.from_pretrained('diffusers/qwen3-vl-8b-instruct-lm-head')\n"
                "    pipe = Ideogram4Pipeline.from_pretrained(model_id, prompt_enhancer_head=head)"
            )
        if self._prompt_enhancer is None:
            self._prompt_enhancer = build_prompt_enhancer(self.text_encoder, self.prompt_enhancer_head)
        if self._caption_logits_processor is None and is_outlines_available():
            self._caption_logits_processor = build_caption_logits_processor(self._prompt_enhancer, self.tokenizer)
        if self._caption_logits_processor is None:
            logger.warning_once(
                "`outlines` is not installed; prompt upsampling runs unconstrained and may not return schema-valid "
                "JSON. Install with `pip install outlines` for structured captions."
            )

        return generate_captions(
            self._prompt_enhancer,
            self.tokenizer,
            self._caption_logits_processor,
            prompt,
            height,
            width,
            temperature=temperature,
            max_new_tokens=max_new_tokens,
            generator=generator,
            device=device,
        )

    @staticmethod
    def _prepare_ids(
        text_lengths: list[int],
        grid_h: int,
        grid_w: int,
        max_text_tokens: int,
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Build the packed `[left-pad][text][image]` layout from the per-prompt text lengths and the image grid.

        Returns `position_ids` (3-axis MRoPE), `segment_ids` (block-diagonal attention) and `indicator` (per-token
        text/image/pad role).
        """
        batch_size = len(text_lengths)
        num_image_tokens = grid_h * grid_w
        total_seq_len = max_text_tokens + num_image_tokens

        # Image position ids (t=0, h, w); offset keeps them disjoint from text positions.
        h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
        w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
        t_idx = torch.zeros_like(h_idx)
        image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET

        position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
        segment_ids = torch.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long)
        indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)

        for b, num_text in enumerate(text_lengths):
            offset = max_text_tokens - num_text

            text_pos = torch.arange(num_text)
            text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1)
            position_ids[b, offset : offset + num_text] = text_pos_3d
            position_ids[b, offset + num_text :] = image_pos

            indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR
            indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR

            segment_ids[b, offset : offset + num_text + num_image_tokens] = 1

        return position_ids.to(device), segment_ids.to(device), indicator.to(device)

    @staticmethod
    def _get_text_encoder_hidden_states(
        text_encoder,
        token_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        pos_2d: torch.Tensor,
    ) -> list[torch.Tensor]:
        """Run the text encoder's decoder layers, returning the hidden states tapped at each activation layer."""

        language_model = text_encoder.language_model

        inputs_embeds = language_model.embed_tokens(token_ids)

        position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1)
        text_position_ids = position_ids_4d[0]
        mrope_position_ids = position_ids_4d[1:]

        causal_mask = create_causal_mask(
            config=language_model.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=None,
            position_ids=text_position_ids,
        )
        position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)

        tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
        captured: dict[int, torch.Tensor] = {}
        hidden_states = inputs_embeds
        for layer_idx, decoder_layer in enumerate(language_model.layers):
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=text_position_ids,
                past_key_values=None,
                position_embeddings=position_embeddings,
            )
            if layer_idx in tap_set:
                captured[layer_idx] = hidden_states

        return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]

    def encode_prompt(
        self,
        prompt: str | list[str],
        grid_h: int,
        grid_w: int,
        max_sequence_length: int,
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Prepare the conditioning for the packed text+image sequence (one entry per prompt).

        Returns a flat tuple `(prompt_embeds, position_ids, segment_ids, indicator)`. The unconditional branch carries
        no text, so the pipeline builds its (zeroed) inputs directly rather than encoding a negative prompt.
        """
        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
        batch_size = len(prompts)
        num_image_tokens = grid_h * grid_w

        # Tokenize each chat-formatted prompt and left-pad to `max_sequence_length`. Only the text region is fed to
        # the encoder: the packed image tokens come after the text and the encoder is causal, so they never affect it.
        token_ids = torch.zeros(batch_size, max_sequence_length, dtype=torch.long)
        attention_mask = torch.zeros(batch_size, max_sequence_length, dtype=torch.long)
        text_position_ids = torch.zeros(batch_size, max_sequence_length, dtype=torch.long)
        text_lengths = []
        for b, text_prompt in enumerate(prompts):
            messages = [{"role": "user", "content": [{"type": "text", "text": text_prompt}]}]
            text = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
            toks = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
            n = int(toks.shape[0])
            if n > max_sequence_length:
                raise ValueError(f"prompt has {n} tokens, exceeds max_sequence_length={max_sequence_length}")
            text_lengths.append(n)
            offset = max_sequence_length - n
            token_ids[b, offset:] = toks
            attention_mask[b, offset:] = 1
            text_position_ids[b, offset:] = torch.arange(n)

        # To support enable_model_cpu_offload, we need to move the text_encoder inputs to the text encoder's actual
        # device te_device. This is necessary because the `CpuOffload` model offload hook attaches to a component's
        # `forward` method, but we call text_encoder's submodules directly below, so the hook never fires to onload the
        # model to the execution device. Other offloading techniques (group, sequential) would work without te_device
        # because they hook submodules, not just the top-level component module. Note that in the
        # enable_model_cpu_offload case te_device will actually be the offload device (e.g. CPU).
        te_device = self.text_encoder.device
        token_ids = token_ids.to(te_device)
        attention_mask = attention_mask.to(te_device)
        text_position_ids = text_position_ids.to(te_device)

        # Concatenate the tapped activation-layer hidden states into per-token text features, zeroing padding.
        selected = self._get_text_encoder_hidden_states(
            self.text_encoder, token_ids, attention_mask, text_position_ids
        )
        text_features = torch.stack(selected, dim=0).permute(1, 2, 3, 0).reshape(batch_size, max_sequence_length, -1)
        text_features = (text_features * attention_mask.to(text_features.dtype).unsqueeze(-1)).to(torch.float32)
        text_features = text_features.to(device)

        position_ids, segment_ids, indicator = self._prepare_ids(
            text_lengths, grid_h, grid_w, max_sequence_length, device
        )

        # Pack the text features into the full sequence; image positions carry no text features.
        image_feature_padding = torch.zeros(
            batch_size, num_image_tokens, text_features.shape[-1], dtype=text_features.dtype, device=device
        )
        prompt_embeds = torch.cat([text_features, image_feature_padding], dim=1)

        return prompt_embeds, position_ids, segment_ids, indicator

    def prepare_latents(
        self,
        batch_size: int,
        num_image_tokens: int,
        latent_dim: int,
        dtype: torch.dtype,
        device: torch.device,
        generator: torch.Generator | list[torch.Generator] | None,
        latents: torch.Tensor | None = None,
    ) -> torch.Tensor:
        shape = (batch_size, num_image_tokens, latent_dim)
        if latents is None:
            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
        else:
            if latents.shape != shape:
                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
            latents = latents.to(device=device, dtype=dtype)
        return latents

    @property
    def guidance_scale(self) -> float | None:
        return self._guidance_scale

    @property
    def num_timesteps(self) -> int:
        return self._num_timesteps

    @property
    def attention_kwargs(self) -> dict[str, Any] | None:
        return self._attention_kwargs

    @property
    def interrupt(self) -> bool:
        return self._interrupt

    def check_inputs(
        self,
        prompt,
        height,
        width,
        num_inference_steps,
        guidance_scale,
        guidance_schedule,
        callback_on_step_end_tensor_inputs=None,
    ):
        if prompt is None:
            raise ValueError("`prompt` must be provided.")
        if not isinstance(prompt, (str, list)):
            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")

        if (
            height % (self.vae_scale_factor * self.patch_size) != 0
            or width % (self.vae_scale_factor * self.patch_size) != 0
        ):
            raise ValueError(
                f"`height` ({height}) and `width` ({width}) must both be divisible by {self.vae_scale_factor * self.patch_size} "
                f"(vae_scale_factor * patch_size)."
            )

        # Guidance is controlled by either a constant `guidance_scale` or a per-step `guidance_schedule`; exactly
        # one must be set (the `guidance_schedule` default makes the no-arg call use the recommended schedule).
        if guidance_scale is not None and guidance_schedule is not None:
            raise ValueError("Only one of `guidance_scale` and `guidance_schedule` may be set.")
        if guidance_scale is None and guidance_schedule is None:
            raise ValueError("One of `guidance_scale` and `guidance_schedule` must be set.")
        if guidance_schedule is not None and len(guidance_schedule) != num_inference_steps:
            raise ValueError(
                f"`guidance_schedule` must have length `num_inference_steps` ({num_inference_steps}), "
                f"got {len(guidance_schedule)}."
            )

        if callback_on_step_end_tensor_inputs is not None and not all(
            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
        ):
            raise ValueError(
                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found "
                f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
            )

    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
        self,
        prompt: str | list[str] | None = None,
        height: int = 2048,
        width: int = 2048,
        num_inference_steps: int = 48,
        guidance_scale: float | None = None,
        guidance_schedule: list[float] | torch.Tensor | None = (7.0,) * 45 + (3.0,) * 3,
        mu: float = 0.0,
        std: float = 1.5,
        prompt_upsampling: bool = False,
        prompt_upsampling_temperature: float = PROMPT_UPSAMPLE_TEMPERATURE,
        max_sequence_length: int = 2048,
        num_images_per_prompt: int = 1,
        generator: torch.Generator | list[torch.Generator] | None = None,
        latents: torch.Tensor | None = None,
        output_type: str = "pil",
        return_dict: bool = True,
        attention_kwargs: dict[str, Any] | None = None,
        callback_on_step_end: Callable[["Ideogram4Pipeline", int, int, dict[str, Any]], dict[str, Any]] | None = None,
        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
    ) -> Ideogram4PipelineOutput | tuple[Any]:
        r"""
        Run text-to-image generation.

        Args:
            prompt (`str` or `list[str]`):
                Prompt(s) to guide image generation.
            height (`int`, *optional*, defaults to 2048):
                Output image height in pixels; must be a multiple of `vae_scale_factor * patch_size`.
            width (`int`, *optional*, defaults to 2048):
                Output image width in pixels; must be a multiple of `vae_scale_factor * patch_size`.
            num_inference_steps (`int`, *optional*, defaults to 48):
                Number of flow-matching steps. The default is the recommended setting for best quality.
            guidance_scale (`float`, *optional*):
                Constant classifier-free guidance scale applied at every step. The conditional and unconditional
                velocity predictions are blended as `v = guidance_scale * v_pos + (1 - guidance_scale) * v_neg`.
                Mutually exclusive with `guidance_schedule` (setting both raises). Defaults to `None`.
            guidance_schedule (`list[float]` or `torch.Tensor`, *optional*):
                Per-step guidance scale schedule; must have length `num_inference_steps`. The first entry corresponds
                to the first step (largest noise level). Mutually exclusive with `guidance_scale`; exactly one must be
                set. Defaults to the recommended schedule (7.0 for the main steps, dropping to 3.0 for the final 3
                "polish" steps). To use a constant scale instead, pass `guidance_scale` and `guidance_schedule=None`.
            mu (`float`, *optional*, defaults to 0.0):
                Base mean of the logit-normal flow-matching schedule. The schedule mean is shifted by half the log of
                the resolution ratio relative to 512x512.
            std (`float`, *optional*, defaults to 1.5):
                Standard deviation of the logit-normal flow-matching schedule.
            prompt_upsampling (`bool`, *optional*, defaults to `False`):
                If `True`, rewrite `prompt` into Ideogram4's native structured JSON caption via
                [`~Ideogram4Pipeline.upsample_prompt`] before encoding. Requires the optional `prompt_enhancer_head`
                component; install `outlines` for schema-constrained captions. `generator` is reused to make the
                upsampling reproducible.
            prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0):
                Sampling temperature for prompt upsampling when `prompt_upsampling=True`.
            max_sequence_length (`int`, *optional*, defaults to 2048):
                Maximum number of text tokens per prompt.
            num_images_per_prompt (`int`, *optional*, defaults to 1):
                Number of images to generate per prompt.
            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
                Generator(s) used to make sampling deterministic.
            latents (`torch.Tensor`, *optional*):
                Pre-generated noise of shape `(batch_size, num_image_tokens, latent_dim)`.
            output_type (`str`, *optional*, defaults to `"pil"`):
                One of `"pil"`, `"np"`, `"pt"`, or `"latent"`.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether to return an [`~pipelines.ideogram4.Ideogram4PipelineOutput`].
            attention_kwargs (`dict`, *optional*):
                A kwargs dictionary passed along to the attention processor of each transformer. A `"scale"` entry
                scales the loaded LoRA weights (e.g. `{"scale": 0.7}`) when the PEFT backend is active.
            callback_on_step_end (`Callable`, *optional*):
                Callback invoked at the end of every denoising step.
            callback_on_step_end_tensor_inputs (`list[str]`, *optional*):
                Names of tensors to expose to the callback via `callback_kwargs`.

        Examples:

        Returns:
            [`~pipelines.ideogram4.Ideogram4PipelineOutput`] or `tuple`.
        """
        self.check_inputs(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            guidance_schedule=guidance_schedule,
            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
        )

        if isinstance(prompt, str):
            batch_size = 1
        elif isinstance(prompt, list):
            batch_size = len(prompt)

        device = self._execution_device
        self._guidance_scale = guidance_scale
        self._attention_kwargs = attention_kwargs
        self._interrupt = False

        # 0. Optionally rewrite the prompt(s) into Ideogram4's native structured JSON caption.
        if prompt_upsampling:
            prompt = self.upsample_prompt(
                prompt,
                height=height,
                width=width,
                temperature=prompt_upsampling_temperature,
                max_new_tokens=max_sequence_length,
                generator=generator,
                device=device,
            )

        # 1. Image grid (drives both the packed layout and the latent shape).
        grid_h, grid_w = (
            height // (self.vae_scale_factor * self.patch_size),
            width // (self.vae_scale_factor * self.patch_size),
        )
        num_image_tokens = grid_h * grid_w

        # 2. Encode prompts into the packed conditioning (one entry per prompt).
        llm_features, position_ids, segment_ids, indicator = self.encode_prompt(
            prompt=prompt,
            grid_h=grid_h,
            grid_w=grid_w,
            max_sequence_length=max_sequence_length,
            device=device,
        )

        # 3. Replicate the conditioning for num_images_per_prompt.
        llm_features = _expand_tensor_to_effective_batch(llm_features, batch_size, num_images_per_prompt)
        position_ids = _expand_tensor_to_effective_batch(position_ids, batch_size, num_images_per_prompt)
        segment_ids = _expand_tensor_to_effective_batch(segment_ids, batch_size, num_images_per_prompt)
        indicator = _expand_tensor_to_effective_batch(indicator, batch_size, num_images_per_prompt)

        # 4. Unconditional (image-only) branch, derived from the conditioning: zeroed text features and the
        # image-region slices of the layout.
        neg_llm_features = torch.zeros(
            batch_size * num_images_per_prompt,
            num_image_tokens,
            llm_features.shape[-1],
            dtype=llm_features.dtype,
            device=device,
        )
        neg_position_ids = position_ids[:, max_sequence_length:]
        neg_segment_ids = segment_ids[:, max_sequence_length:]
        neg_indicator = indicator[:, max_sequence_length:]

        # 4. Set up the resolution-aware logit-normal schedule on the scheduler.
        schedule_mu = _resolution_aware_mu(height=height, width=width, base_mu=mu)
        sigmas = _logit_normal_sigmas(num_inference_steps, schedule_mu, std=std, device=device)
        self.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device)
        timesteps = self.scheduler.timesteps
        self._num_timesteps = len(timesteps)

        # 5. Resolve the per-step guidance schedule (a constant `guidance_scale` broadcasts to every step, otherwise
        # use the provided `guidance_schedule`, validated by `check_inputs`) and the tensor of per-step weights `gw`.
        if guidance_scale is not None:
            guidance_schedule = [float(guidance_scale)] * num_inference_steps
        gw = torch.as_tensor(guidance_schedule, dtype=torch.float32, device=device)

        # 6. Prepare latents in the packed (B, num_image_tokens, latent_dim) layout.
        latent_dim = self.transformer.config.in_channels
        latents = self.prepare_latents(
            batch_size=batch_size * num_images_per_prompt,
            num_image_tokens=num_image_tokens,
            latent_dim=latent_dim,
            dtype=torch.float32,
            device=device,
            generator=generator,
            latents=latents,
        )

        # 7. Padding for the text region of the conditional packed sequence (image latents are appended after it).
        max_text_tokens = max_sequence_length
        text_z_padding = torch.zeros(
            batch_size * num_images_per_prompt,
            max_text_tokens,
            latent_dim,
            dtype=torch.float32,
            device=device,
        )

        # The transformers run in their loaded compute dtype; cast the (otherwise float32) text features to match.
        # `latents` stay float32 for scheduler precision and are cast per-step at the transformer call below.
        llm_features = llm_features.to(self.transformer.dtype)
        neg_llm_features = neg_llm_features.to(self.unconditional_transformer.dtype)

        # 8. Denoising loop. The scheduler stores `num_train_timesteps`-scaled timesteps; convert back to model time.
        num_train_timesteps = self.scheduler.config.num_train_timesteps
        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
                    continue

                # Map sigma-domain timestep to model time `t` in [0, 1] (0 = noise, 1 = clean data).
                t_model = 1.0 - (t.float() / num_train_timesteps)
                t_model = t_model.expand(batch_size * num_images_per_prompt).to(self.transformer.dtype)

                # Conditional pass operates on the full packed sequence.
                pos_z = torch.cat([text_z_padding, latents], dim=1).to(self.transformer.dtype)
                pos_out = self.transformer(
                    hidden_states=pos_z,
                    timestep=t_model,
                    encoder_hidden_states=llm_features,
                    position_ids=position_ids,
                    segment_ids=segment_ids,
                    indicator=indicator,
                    attention_kwargs=self.attention_kwargs,
                    return_dict=False,
                )[0]
                # Velocity (and guidance) is computed in float32 for scheduler precision; the transformers
                # return their compute dtype, so cast the predicted velocities up here.
                pos_v = pos_out[:, max_text_tokens:].to(torch.float32)

                # Unconditional pass uses image-only positions with zeroed text features.
                neg_v = self.unconditional_transformer(
                    hidden_states=latents.to(self.unconditional_transformer.dtype),
                    timestep=t_model,
                    encoder_hidden_states=neg_llm_features,
                    position_ids=neg_position_ids,
                    segment_ids=neg_segment_ids,
                    indicator=neg_indicator,
                    attention_kwargs=self.attention_kwargs,
                    return_dict=False,
                )[0].to(torch.float32)

                # Expose the current step's guidance weight via `self.guidance_scale` so callbacks can read it.
                self._guidance_scale = guidance_schedule[i]
                gw_i = gw[i]
                v = gw_i * pos_v + (1.0 - gw_i) * neg_v

                latents = self.scheduler.step(-v, t, latents, return_dict=False)[0]

                if callback_on_step_end is not None:
                    callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs}
                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
                    latents = callback_outputs.pop("latents", latents)

                progress_bar.update()

        # 9. Decode: unpatch the latents, denormalize with the VAE batch-norm stats, and decode through the VAE.
        if output_type == "latent":
            image = latents
        else:
            z = latents
            # VAE bn stores per-channel statistics on the packed-channel latent space (ae_channels * patch ** 2).
            bn_mean = self.vae.bn.running_mean.view(1, 1, -1).to(device=z.device, dtype=z.dtype)
            bn_std = torch.sqrt(self.vae.bn.running_var + self.vae.config.batch_norm_eps).view(1, 1, -1)
            bn_std = bn_std.to(device=z.device, dtype=z.dtype)
            z = z * bn_std + bn_mean

            patch = self.patch_size
            ae_channels = z.shape[-1] // (patch * patch)
            z = z.view(batch_size * num_images_per_prompt, grid_h, grid_w, patch, patch, ae_channels)
            z = z.permute(0, 5, 1, 3, 2, 4).contiguous()
            z = z.view(batch_size * num_images_per_prompt, ae_channels, grid_h * patch, grid_w * patch)

            decoded = self.vae.decode(z.to(self.vae.dtype), return_dict=False)[0]
            image = self.image_processor.postprocess(decoded.float(), output_type=output_type)

        self.maybe_free_model_hooks()

        if not return_dict:
            return (image,)
        return Ideogram4PipelineOutput(images=image)