ideogram4_differential.py

# Copyright 2026 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.
"""Custom modular-diffusers blocks adding Differential Diffusion to Ideogram4.

Differential Diffusion (https://differential-diffusion.github.io/) is an image-to-image workflow driven by a
grayscale change map: each token is regenerated from the prompt for a number of steps proportional to how dark it
is (darker = more change, brighter = kept closer to the reference), giving smooth per-region control. Unlike plain
img2img it uses the full text2image schedule (no `strength`); it differs only in prepare-latents (build the per-step
change masks) and the denoise loop (blend the denoised latents with the re-noised reference per the current mask).
The mask math matches the diffusers community differential-diffusion pipelines exactly.

Reuses the img2img VAE encoder / input step and the in-repo Ideogram4 blocks; only the prepare-latents and the
denoise-loop blend are new.
"""

import torch

from diffusers.configuration_utils import FrozenDict
from diffusers.image_processor import VaeImageProcessor
from diffusers.modular_pipelines.ideogram4.before_denoise import (
    Ideogram4PrepareAdditionalInputsStep,
    Ideogram4SetTimestepsStep,
)
from diffusers.modular_pipelines.ideogram4.decoders import Ideogram4DecodeStep
from diffusers.modular_pipelines.ideogram4.denoise import (
    Ideogram4AfterDenoiseStep,
    Ideogram4DenoiseStep,
    Ideogram4LoopAfterDenoiser,
    Ideogram4LoopBeforeDenoiser,
    Ideogram4LoopDenoiser,
)
from diffusers.modular_pipelines.ideogram4.encoders import Ideogram4PromptUpsampleStep, Ideogram4TextEncoderStep
from diffusers.modular_pipelines.ideogram4.modular_blocks_ideogram4 import Ideogram4CoreDenoiseStep
from diffusers.modular_pipelines.ideogram4.modular_pipeline import Ideogram4ModularPipeline
from diffusers.modular_pipelines.modular_pipeline import (
    BlockState,
    ConditionalPipelineBlocks,
    ModularPipelineBlocks,
    PipelineState,
    SequentialPipelineBlocks,
)
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, InsertableDict, OutputParam
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import logging
from diffusers.utils.torch_utils import randn_tensor

from .ideogram4_img2img import Ideogram4AutoVaeEncoderStep, Ideogram4Img2ImgInputStep


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


# Grayscale change map processor: keep values in [0, 1] (no [-1, 1] normalization) so they compare directly
# against the per-step thresholds, and collapse RGB inputs to a single channel.
MASK_PROCESSOR_SPEC = ComponentSpec(
    "mask_processor",
    VaeImageProcessor,
    config=FrozenDict({"do_normalize": False, "do_convert_grayscale": True, "vae_scale_factor": 1}),
    default_creation_method="from_config",
)


# auto_docstring
class Ideogram4DiffDiffPrepareLatentsStep(ModularPipelineBlocks):
    """
    Differential Diffusion prepare-latents: seed the loop with the fully-noised reference latents and build the
    per-step change masks (resized to the packed-token grid) from the `diffdiff_map`. Run after set_timesteps.
    """

    model_name = "ideogram4"

    @property
    def description(self) -> str:
        return (
            "Differential Diffusion prepare-latents: seed the loop from the fully-noised reference latents and build "
            "the per-step change masks (in packed-token space) from the `diffdiff_map`. Run after set_timesteps."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), MASK_PROCESSOR_SPEC]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                name="diffdiff_map",
                required=True,
                description="Grayscale change map aligned to the reference `image`: darker regions are regenerated "
                "over more steps, brighter regions are kept closer to the reference.",
            ),
            InputParam.template("image_latents", required=True),
            InputParam.template("height", required=True),
            InputParam.template("width", required=True),
            InputParam.template("generator"),
            InputParam.template("num_inference_steps", default=48),
            InputParam(
                name="timesteps",
                required=True,
                type_hint=torch.Tensor,
                description="Denoising timesteps from set_timesteps.",
            ),
            InputParam(
                name="denoising_start",
                default=None,
                type_hint=float,
                description="Optional schedule offset added to the mask thresholds (advanced).",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="latents",
                type_hint=torch.Tensor,
                description="Initial fully-noised reference latents (B, num_image_tokens, latent_dim).",
            ),
            OutputParam(
                name="noise",
                type_hint=torch.Tensor,
                description="The random noise used to seed and re-noise the reference latents.",
            ),
            OutputParam(
                name="diffdiff_masks",
                type_hint=torch.Tensor,
                description="Per-step boolean change masks in packed-token space (num_inference_steps, num_image_tokens).",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        device = components._execution_device
        patch = components.patch_size
        grid_h = block_state.height // (components.vae_scale_factor * patch)
        grid_w = block_state.width // (components.vae_scale_factor * patch)

        image_latents = block_state.image_latents.to(device=device, dtype=torch.float32)
        noise = randn_tensor(image_latents.shape, generator=block_state.generator, device=device, dtype=torch.float32)

        # Seed the loop from the reference noised to the first timestep (all regions start fully noised).
        latent_timestep = block_state.timesteps[:1].repeat(image_latents.shape[0])
        block_state.latents = components.scheduler.scale_noise(image_latents, latent_timestep, noise)
        block_state.noise = noise

        # Per-step masks: a token stays anchored to the reference while its map value exceeds the (rising) step
        # threshold, and is released to free denoising once the threshold catches up. Resize the map to the token
        # grid and flatten row-major to match the patchify token order.
        change_map = components.mask_processor.preprocess(block_state.diffdiff_map, height=grid_h, width=grid_w)
        change_map = change_map.to(device=device, dtype=torch.float32)
        denoising_start = block_state.denoising_start or 0.0
        thresholds = torch.arange(block_state.num_inference_steps, device=device, dtype=torch.float32)
        thresholds = (thresholds / block_state.num_inference_steps).view(-1, 1, 1, 1) + denoising_start
        masks = change_map > thresholds
        block_state.diffdiff_masks = masks.reshape(block_state.num_inference_steps, grid_h * grid_w)

        self.set_block_state(state, block_state)
        return components, state


class Ideogram4DiffDiffLoopBlendStep(ModularPipelineBlocks):
    model_name = "ideogram4"

    @property
    def description(self) -> str:
        return (
            "Within the denoising loop: after the scheduler step, blend the denoised latents with the reference "
            "re-noised to the next timestep, using the current step's change mask (Differential Diffusion). Compose "
            "into `Ideogram4DiffDiffDenoiseStep` after the after_denoiser sub-block."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Denoised packed latents."),
            InputParam.template("image_latents", required=True),
            InputParam(
                name="noise",
                required=True,
                type_hint=torch.Tensor,
                description="The random noise used to re-noise the reference latents.",
            ),
            InputParam(
                name="diffdiff_masks",
                required=True,
                type_hint=torch.Tensor,
                description="Per-step boolean change masks in packed-token space.",
            ),
            InputParam(name="timesteps", required=True, type_hint=torch.Tensor, description="Denoising timesteps."),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [OutputParam(name="latents", type_hint=torch.Tensor, description="The blended latents.")]

    @torch.no_grad()
    def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
        timesteps = block_state.timesteps
        # No blend after the final step: the last denoised latents are the result.
        if i < len(timesteps) - 1:
            next_t = timesteps[i + 1].expand(block_state.latents.shape[0])
            image_latent = components.scheduler.scale_noise(block_state.image_latents, next_t, block_state.noise)
            mask = block_state.diffdiff_masks[i].to(block_state.latents.dtype).view(1, -1, 1)
            block_state.latents = image_latent * mask + block_state.latents * (1 - mask)
        return components, block_state


# auto_docstring
class Ideogram4DiffDiffDenoiseStep(Ideogram4DenoiseStep):
    """
    Differential Diffusion denoising loop: the standard Ideogram4 loop (asymmetric CFG over the two transformers)
    with an extra per-step blend that re-anchors masked tokens to the reference image.
    """

    model_name = "ideogram4"
    block_classes = [
        Ideogram4LoopBeforeDenoiser,
        Ideogram4LoopDenoiser,
        Ideogram4LoopAfterDenoiser,
        Ideogram4DiffDiffLoopBlendStep,
    ]
    block_names = ["before_denoiser", "denoiser", "after_denoiser", "diffdiff_blend"]

    @property
    def description(self) -> str:
        return (
            "Differential Diffusion denoising loop: the standard Ideogram4 asymmetric-CFG loop plus a per-step blend "
            "that re-anchors masked tokens to the re-noised reference image."
        )


# Diff-diff core denoise: full text2image schedule (no strength), diff-diff prepare-latents and denoise loop.
DIFFDIFF_CORE_DENOISE_BLOCKS = InsertableDict(
    [
        ("input", Ideogram4Img2ImgInputStep()),
        ("set_timesteps", Ideogram4SetTimestepsStep()),
        ("prepare_latents", Ideogram4DiffDiffPrepareLatentsStep()),
        ("prepare_additional_inputs", Ideogram4PrepareAdditionalInputsStep()),
        ("denoise", Ideogram4DiffDiffDenoiseStep()),
        ("after_denoise", Ideogram4AfterDenoiseStep()),
    ]
)


# auto_docstring
class Ideogram4DiffDiffCoreDenoiseStep(SequentialPipelineBlocks):
    """
    Core denoising workflow for Ideogram4 Differential Diffusion: batch inputs -> full text2image schedule ->
    build change masks + seed from the reference -> asymmetric-CFG denoising loop with per-step masked blending ->
    unpatchify for the decoder.
    """

    model_name = "ideogram4"
    block_classes = list(DIFFDIFF_CORE_DENOISE_BLOCKS.values())
    block_names = list(DIFFDIFF_CORE_DENOISE_BLOCKS.keys())

    @property
    def description(self) -> str:
        return (
            "Core denoising workflow for Ideogram4 Differential Diffusion: batch inputs -> full text2image schedule "
            "-> build change masks and seed from the reference -> asymmetric-CFG denoising loop with per-step masked "
            "blending -> unpatchify for the decoder."
        )

    @property
    def outputs(self) -> list[OutputParam]:
        return [OutputParam.template("latents", description="Unpatchified (B, ae_channels, H, W) latents.")]


class Ideogram4DiffDiffAutoCoreDenoiseStep(ConditionalPipelineBlocks):
    block_classes = [Ideogram4CoreDenoiseStep, Ideogram4DiffDiffCoreDenoiseStep]
    block_names = ["text2image", "differential"]
    block_trigger_inputs = ["diffdiff_map"]
    default_block_name = "text2image"

    def select_block(self, diffdiff_map=None):
        return "differential" if diffdiff_map is not None else "text2image"

    @property
    def description(self) -> str:
        return (
            "Core denoising step. \n"
            " - `Ideogram4DiffDiffCoreDenoiseStep` (differential) is used when `diffdiff_map` is provided.\n"
            " - `Ideogram4CoreDenoiseStep` (text2image) is used otherwise."
        )

    @property
    def outputs(self) -> list[OutputParam]:
        return [OutputParam.template("latents", description="Unpatchified (B, ae_channels, H, W) latents.")]


DIFFDIFF_AUTO_BLOCKS = InsertableDict(
    [
        ("prompt_upsample", Ideogram4PromptUpsampleStep()),
        ("text_encoder", Ideogram4TextEncoderStep()),
        ("vae_encoder", Ideogram4AutoVaeEncoderStep()),
        ("denoise", Ideogram4DiffDiffAutoCoreDenoiseStep()),
        ("decode", Ideogram4DecodeStep()),
    ]
)


# auto_docstring
class Ideogram4DiffDiffAutoBlocks(SequentialPipelineBlocks):
    """
    Auto Modular pipeline for Ideogram4 supporting text-to-image and Differential Diffusion, selected by the presence
    of `diffdiff_map`: (optional) prompt upsampling -> encode text -> (diff-diff) VAE-encode reference -> core denoise
    (asymmetric CFG over two transformers, with per-step masked blending) -> decode.

      Supported workflows:
        - `text2image`: requires `prompt`
        - `differential`: requires `prompt`, `image`, `diffdiff_map`
    """

    model_name = "ideogram4"
    block_classes = list(DIFFDIFF_AUTO_BLOCKS.values())
    block_names = list(DIFFDIFF_AUTO_BLOCKS.keys())

    _workflow_map = {
        "text2image": {"prompt": True},
        "differential": {"prompt": True, "image": True, "diffdiff_map": True},
    }

    @property
    def description(self) -> str:
        return (
            "Auto Modular pipeline for Ideogram4 text-to-image and Differential Diffusion: (optional) prompt "
            "upsampling -> encode text -> (diff-diff) VAE-encode the reference image -> core denoise (asymmetric CFG "
            "over two transformers, with per-step masked blending) -> decode. Selected by the presence of "
            "`diffdiff_map`."
        )

    @property
    def outputs(self) -> list[OutputParam]:
        return [OutputParam.template("images")]