# 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")]