diff --git a/docs/source/en/api/pipelines/ideogram4.md b/docs/source/en/api/pipelines/ideogram4.md index d9f3e341b169..e6866f84692e 100644 --- a/docs/source/en/api/pipelines/ideogram4.md +++ b/docs/source/en/api/pipelines/ideogram4.md @@ -40,6 +40,56 @@ image = pipe(prompt, height=1024, width=1024, generator=torch.Generator("cuda"). image.save("ideogram4.png") ``` +## Modular image-to-image and inpainting + +The [`Ideogram4ModularPipeline`] automatically selects text-to-image, image-to-image, or inpainting blocks from the +inputs. Pass `image` to run image-to-image generation, and pass both `image` and `mask_image` to run inpainting. White +mask pixels are repainted and black pixels preserve the source image. + +```python +import torch +from diffusers import ModularPipeline +from diffusers.utils import load_image + +pipe = ModularPipeline.from_pretrained("ideogram-ai/ideogram-4-nf4") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +source = load_image( + "https://github.com/lucasruan1618/Image_storage/blob/main/Input/cute_cat.png?raw=true" +).convert("RGB") +image = pipe( + prompt="A watercolor painting of a cat", + image=source, + strength=0.7, + generator=torch.Generator("cuda").manual_seed(0), +).images[0] +image.save("ideogram4_img2img.png") +``` + +For inpainting, provide a mask at the same logical resolution as the source image. `padding_mask_crop` is also +supported for higher-resolution work on a cropped masked region. + +```python +from diffusers.utils import load_image + +source = load_image( + "https://github.com/lucasruan1618/Image_storage/blob/main/Input/cute_cat.png?raw=true" +).convert("RGB") +mask = load_image( + "https://github.com/lucasruan1618/Image_storage/blob/main/Input/mask_cat.png?raw=true" +).convert("L") + +image = pipe( + prompt="A cat wearing a red wizard hat", + image=source, + mask_image=mask, + strength=0.9, + generator=torch.Generator("cuda").manual_seed(0), +).images[0] +image.save("ideogram4_inpaint.png") +``` + ## Prompt upsampling Ideogram 4 is trained on a structured JSON caption rather than a free-form prompt, so a short prompt is best @@ -115,3 +165,7 @@ image.save("ideogram4_upsampled.png") ## Ideogram4PipelineOutput [[autodoc]] pipelines.ideogram4.pipeline_output.Ideogram4PipelineOutput + +## Ideogram4ModularPipeline + +[[autodoc]] Ideogram4ModularPipeline diff --git a/src/diffusers/modular_pipelines/ideogram4/before_denoise.py b/src/diffusers/modular_pipelines/ideogram4/before_denoise.py index 98be3b141aec..7e40dc64c531 100644 --- a/src/diffusers/modular_pipelines/ideogram4/before_denoise.py +++ b/src/diffusers/modular_pipelines/ideogram4/before_denoise.py @@ -195,6 +195,121 @@ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class Ideogram4ImageInputsStep(ModularPipelineBlocks): + """ + Expand per-image latents to the effective prompt batch for image-to-image and inpainting workflows. + + Inputs: + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + Effective batch size. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Expand per-image latents to the effective prompt batch for image-to-image and inpainting workflows." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + InputParam.template("num_images_per_prompt", default=1), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("image_latents")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_batch = block_state.batch_size // block_state.num_images_per_prompt + block_state.image_latents = _expand_tensor_to_effective_batch( + block_state.image_latents, + prompt_batch, + block_state.num_images_per_prompt, + "image_latents", + ).to(device=components._execution_device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4MaskInputsStep(ModularPipelineBlocks): + """ + Expand a preprocessed inpaint mask to the effective prompt batch. + + Inputs: + processed_mask_image (`Tensor`): + The binary mask tensor resized to the generation resolution. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + Effective batch size. + + Outputs: + processed_mask_image (`Tensor`): + The binary mask tensor expanded to the effective prompt batch. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Expand a preprocessed inpaint mask to the effective prompt batch." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_mask_image", + required=True, + type_hint=torch.Tensor, + description="The binary mask tensor resized to the generation resolution.", + ), + InputParam.template("num_images_per_prompt", default=1), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_mask_image", + type_hint=torch.Tensor, + description="The binary mask tensor expanded to the effective prompt batch.", + ) + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + prompt_batch = block_state.batch_size // block_state.num_images_per_prompt + block_state.processed_mask_image = _expand_tensor_to_effective_batch( + block_state.processed_mask_image, + prompt_batch, + block_state.num_images_per_prompt, + "processed_mask_image", + ).to(device=components._execution_device, dtype=torch.float32) + + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Ideogram4PrepareLatentsStep(ModularPipelineBlocks): """ @@ -280,6 +395,82 @@ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class Ideogram4PrepareLatentsWithStrengthStep(ModularPipelineBlocks): + """ + Add the initial noise to encoded image latents at the first strength-adjusted timestep. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + latents (`Tensor`): + The initial random noise. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + timesteps (`Tensor`): + The strength-adjusted denoising timesteps. + + Outputs: + initial_noise (`Tensor`): + The initial random noise. + latents (`Tensor`): + Encoded image latents noised at the first denoising timestep. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Add the initial noise to encoded image latents at the first strength-adjusted timestep." + + @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="The initial random noise."), + InputParam.template("image_latents"), + InputParam( + name="timesteps", + required=True, + type_hint=torch.Tensor, + description="The strength-adjusted denoising timesteps.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="initial_noise", type_hint=torch.Tensor, description="The initial random noise."), + OutputParam( + name="latents", + type_hint=torch.Tensor, + description="Encoded image latents noised at the first denoising timestep.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + if block_state.image_latents.shape != block_state.latents.shape: + raise ValueError( + f"`image_latents` and `latents` must have the same shape, got {block_state.image_latents.shape} and " + f"{block_state.latents.shape}." + ) + block_state.initial_noise = block_state.latents + latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0]) + block_state.latents = components.scheduler.scale_noise( + block_state.image_latents, latent_timestep, block_state.initial_noise + ) + + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Ideogram4SetTimestepsStep(ModularPipelineBlocks): """ @@ -372,6 +563,160 @@ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class Ideogram4ApplyStrengthStep(ModularPipelineBlocks): + """ + Truncate timesteps and guidance weights according to image-to-image or inpaint strength. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + timesteps (`Tensor`): + The full denoising timesteps. + gw (`Tensor`): + Per-step guidance weights. + + Outputs: + timesteps (`Tensor`): + The strength-adjusted denoising timesteps. + gw (`Tensor`): + The strength-adjusted per-step guidance weights. + num_inference_steps (`int`): + The number of denoising steps after applying strength. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Truncate timesteps and guidance weights according to image-to-image or inpaint strength." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("strength", default=0.9), + InputParam.template("num_inference_steps", default=48), + InputParam( + name="timesteps", required=True, type_hint=torch.Tensor, description="The full denoising timesteps." + ), + InputParam(name="gw", required=True, type_hint=torch.Tensor, description="Per-step guidance weights."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="timesteps", type_hint=torch.Tensor, description="The strength-adjusted denoising timesteps." + ), + OutputParam( + name="gw", type_hint=torch.Tensor, description="The strength-adjusted per-step guidance weights." + ), + OutputParam( + name="num_inference_steps", + type_hint=int, + description="The number of denoising steps after applying strength.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + if not 0.0 < block_state.strength <= 1.0: + raise ValueError(f"`strength` must be in the interval (0, 1], but got {block_state.strength}.") + + init_timestep = min(block_state.num_inference_steps * block_state.strength, block_state.num_inference_steps) + t_start = int(max(block_state.num_inference_steps - init_timestep, 0)) + begin_index = t_start * components.scheduler.order + block_state.timesteps = block_state.timesteps[begin_index:] + block_state.gw = block_state.gw[t_start:] + block_state.num_inference_steps = block_state.num_inference_steps - t_start + components.scheduler.set_begin_index(begin_index) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4PrepareMaskLatentsStep(ModularPipelineBlocks): + """ + Resize and pack an inpaint mask to match the Ideogram4 latent token layout. + + Components: + transformer (`Ideogram4Transformer2DModel`) + + Inputs: + processed_mask_image (`Tensor`): + The binary mask tensor expanded to the effective prompt batch. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + + Outputs: + mask (`Tensor`): + The packed latent-space inpaint mask. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Resize and pack an inpaint mask to match the Ideogram4 latent token layout." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("transformer", Ideogram4Transformer2DModel)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_mask_image", + required=True, + type_hint=torch.Tensor, + description="The binary mask tensor expanded to the effective prompt batch.", + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="mask", type_hint=torch.Tensor, description="The packed latent-space inpaint mask.")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + patch = components.patch_size + latent_height = block_state.height // components.vae_scale_factor + latent_width = block_state.width // components.vae_scale_factor + grid_h, grid_w = latent_height // patch, latent_width // patch + latent_channels = components.transformer.config.in_channels // (patch * patch) + + mask = torch.nn.functional.interpolate( + block_state.processed_mask_image, + size=(latent_height, latent_width), + mode="nearest", + ) + mask = mask.repeat(1, latent_channels, 1, 1) + mask = mask.view(mask.shape[0], latent_channels, grid_h, patch, grid_w, patch) + block_state.mask = mask.permute(0, 2, 4, 3, 5, 1).reshape(mask.shape[0], grid_h * grid_w, -1) + + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Ideogram4PrepareAdditionalInputsStep(ModularPipelineBlocks): """ diff --git a/src/diffusers/modular_pipelines/ideogram4/decoders.py b/src/diffusers/modular_pipelines/ideogram4/decoders.py index bf5d69270b7c..1dcc473050be 100644 --- a/src/diffusers/modular_pipelines/ideogram4/decoders.py +++ b/src/diffusers/modular_pipelines/ideogram4/decoders.py @@ -13,10 +13,12 @@ # limitations under the License. +from typing import Any + import torch from ...configuration_utils import FrozenDict -from ...image_processor import VaeImageProcessor +from ...image_processor import InpaintProcessor, VaeImageProcessor from ...models import AutoencoderKLFlux2 from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState @@ -110,3 +112,89 @@ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) - self.set_block_state(state, block_state) return components, state + + +# auto_docstring +class Ideogram4InpaintDecodeStep(ModularPipelineBlocks): + """ + Decode Ideogram4 latents and optionally composite a cropped inpaint result over the source image. + + Components: + vae (`AutoencoderKLFlux2`) image_mask_processor (`InpaintProcessor`) + + Inputs: + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + latents (`Tensor`): + The unpatchified latents to decode. + mask_overlay_kwargs (`dict`): + Arguments used to composite a cropped inpaint result over the source image. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Decode Ideogram4 latents and optionally composite a cropped inpaint result over the source image." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLFlux2), + ComponentSpec( + "image_mask_processor", + InpaintProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("output_type", default="pil"), + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The unpatchified latents to decode.", + ), + InputParam( + name="mask_overlay_kwargs", + required=True, + type_hint=dict[str, Any], + description="Arguments used to composite a cropped inpaint result over the source image.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + z = block_state.latents + patch = components.patch_size + ae_channels = z.shape[1] + grid_h, grid_w = z.shape[2] // patch, z.shape[3] // patch + + bn_mean = components.vae.bn.running_mean.view(patch, patch, ae_channels).permute(2, 0, 1) + bn_std = torch.sqrt(components.vae.bn.running_var + components.vae.config.batch_norm_eps) + bn_std = bn_std.view(patch, patch, ae_channels).permute(2, 0, 1) + bn_mean = bn_mean.repeat(1, grid_h, grid_w).to(device=z.device, dtype=z.dtype) + bn_std = bn_std.repeat(1, grid_h, grid_w).to(device=z.device, dtype=z.dtype) + z = z * bn_std + bn_mean + + decoded = components.vae.decode(z.to(components.vae.dtype), return_dict=False)[0] + block_state.images = components.image_mask_processor.postprocess( + decoded.float(), output_type=block_state.output_type, **block_state.mask_overlay_kwargs + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/ideogram4/denoise.py b/src/diffusers/modular_pipelines/ideogram4/denoise.py index 871db69d344c..67103e703058 100644 --- a/src/diffusers/modular_pipelines/ideogram4/denoise.py +++ b/src/diffusers/modular_pipelines/ideogram4/denoise.py @@ -207,6 +207,48 @@ def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState return components, block_state +class Ideogram4LoopAfterDenoiserInpaint(ModularPipelineBlocks): + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Within the denoising loop: preserve the unmasked source-image latents at the next noise level." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="mask", required=True, type_hint=torch.Tensor, description="Packed latent-space mask."), + InputParam.template("image_latents"), + InputParam( + name="initial_noise", required=True, type_hint=torch.Tensor, description="The initial random noise." + ), + InputParam( + name="timesteps", + required=True, + type_hint=torch.Tensor, + description="The strength-adjusted denoising timesteps.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="latents", type_hint=torch.Tensor, description="The blended inpaint latents.")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + source_latents = block_state.image_latents + if i < len(block_state.timesteps) - 1: + next_timestep = block_state.timesteps[i + 1].reshape(1) + source_latents = components.scheduler.scale_noise(source_latents, next_timestep, block_state.initial_noise) + + block_state.latents = (1.0 - block_state.mask) * source_latents + block_state.mask * block_state.latents + return components, block_state + + # auto_docstring class Ideogram4DenoiseStep(LoopSequentialPipelineBlocks): """ @@ -292,6 +334,71 @@ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) - return components, state +# auto_docstring +class Ideogram4InpaintDenoiseStep(Ideogram4DenoiseStep): + """ + Ideogram4 denoising loop with source-latent preservation outside the inpaint mask. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Ideogram4Transformer2DModel`) + unconditional_transformer (`Ideogram4Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + latents (`Tensor`): + Packed image latents. + position_ids (`Tensor`): + Conditional position ids. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Packed conditional encoder_hidden_states. + position_ids (`Tensor`): + Conditional 3-axis MRoPE position ids. + segment_ids (`Tensor`): + Conditional block-diagonal segment ids. + indicator (`Tensor`): + Conditional per-token text/image/pad role. + negative_prompt_embeds (`Tensor`): + Unconditional (zeroed) text features. + negative_position_ids (`Tensor`): + Unconditional position ids (image region). + negative_segment_ids (`Tensor`): + Unconditional segment ids (image region). + negative_indicator (`Tensor`): + Unconditional indicator (image region). + gw (`Tensor`): + Per-step guidance weights. + mask (`Tensor`): + Packed latent-space mask. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + initial_noise (`Tensor`): + The initial random noise. + timesteps (`Tensor`): + The strength-adjusted denoising timesteps. + + Outputs: + latents (`Tensor`): + The blended inpaint latents. + """ + + block_classes = [ + Ideogram4LoopBeforeDenoiser, + Ideogram4LoopDenoiser, + Ideogram4LoopAfterDenoiser, + Ideogram4LoopAfterDenoiserInpaint, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser", "after_denoiser_inpaint"] + + @property + def description(self) -> str: + return "Ideogram4 denoising loop with source-latent preservation outside the inpaint mask." + + # auto_docstring class Ideogram4AfterDenoiseStep(ModularPipelineBlocks): """ diff --git a/src/diffusers/modular_pipelines/ideogram4/encoders.py b/src/diffusers/modular_pipelines/ideogram4/encoders.py index 6e149fa8392e..6a85d8d830af 100644 --- a/src/diffusers/modular_pipelines/ideogram4/encoders.py +++ b/src/diffusers/modular_pipelines/ideogram4/encoders.py @@ -17,6 +17,9 @@ from transformers import Qwen2Tokenizer, Qwen3VLModel from transformers.masking_utils import create_causal_mask +from ...configuration_utils import FrozenDict +from ...image_processor import InpaintProcessor, VaeImageProcessor +from ...models import AutoencoderKLFlux2 from ...pipelines.ideogram4.prompt_enhancer import ( PROMPT_UPSAMPLE_TEMPERATURE, Ideogram4PromptEnhancerHead, @@ -38,6 +41,278 @@ QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35) +# auto_docstring +class Ideogram4ProcessImageInputStep(ModularPipelineBlocks): + """ + Preprocess an image for the Ideogram4 VAE and resolve the output height and width. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + processed_image (`Tensor`): + The image tensor resized and normalized for VAE encoding. + height (`int`): + The resolved image height in pixels. + width (`int`): + The resolved image width in pixels. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Preprocess an image for the Ideogram4 VAE and resolve the output height and width." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ) + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image"), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_image", + type_hint=torch.Tensor, + description="The image tensor resized and normalized for VAE encoding.", + ), + OutputParam(name="height", type_hint=int, description="The resolved image height in pixels."), + OutputParam(name="width", type_hint=int, description="The resolved image width in pixels."), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + block_state.processed_image = components.image_processor.preprocess( + image=block_state.image, + height=block_state.height, + width=block_state.width, + ) + block_state.height, block_state.width = block_state.processed_image.shape[-2:] + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4InpaintProcessImagesInputStep(ModularPipelineBlocks): + """ + Preprocess an image and mask together for Ideogram4 inpainting. + + Components: + image_processor (`VaeImageProcessor`) image_mask_processor (`InpaintProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`): + Mask image for inpainting. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + + Outputs: + processed_image (`Tensor`): + The image tensor resized and normalized for VAE encoding. + processed_mask_image (`Tensor`): + The binary mask tensor resized to the generation resolution. + mask_overlay_kwargs (`dict`): + Arguments used to composite a cropped inpaint result over the source image. + height (`int`): + The resolved image height in pixels. + width (`int`): + The resolved image width in pixels. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Preprocess an image and mask together for Ideogram4 inpainting." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ComponentSpec( + "image_mask_processor", + InpaintProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image"), + InputParam.template("mask_image"), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("padding_mask_crop"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_image", + type_hint=torch.Tensor, + description="The image tensor resized and normalized for VAE encoding.", + ), + OutputParam( + name="processed_mask_image", + type_hint=torch.Tensor, + description="The binary mask tensor resized to the generation resolution.", + ), + OutputParam( + name="mask_overlay_kwargs", + type_hint=dict, + description="Arguments used to composite a cropped inpaint result over the source image.", + ), + OutputParam(name="height", type_hint=int, description="The resolved image height in pixels."), + OutputParam(name="width", type_hint=int, description="The resolved image width in pixels."), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + image = block_state.image[0] if isinstance(block_state.image, list) else block_state.image + block_state.height, block_state.width = components.image_processor.get_default_height_width( + image, + height=block_state.height, + width=block_state.width, + ) + ( + block_state.processed_image, + block_state.processed_mask_image, + block_state.mask_overlay_kwargs, + ) = components.image_mask_processor.preprocess( + image=block_state.image, + mask=block_state.mask_image, + height=block_state.height, + width=block_state.width, + padding_mask_crop=block_state.padding_mask_crop, + ) + block_state.height, block_state.width = block_state.processed_image.shape[-2:] + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Ideogram4VaeEncoderStep(ModularPipelineBlocks): + """ + Encode a preprocessed image into normalized, packed Ideogram4 image latents. + + Components: + vae (`AutoencoderKLFlux2`) + + Inputs: + processed_image (`Tensor`): + The image tensor resized and normalized for VAE encoding. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "ideogram4" + + @property + def description(self) -> str: + return "Encode a preprocessed image into normalized, packed Ideogram4 image latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLFlux2)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_image", + required=True, + type_hint=torch.Tensor, + description="The image tensor resized and normalized for VAE encoding.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("image_latents")] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + image = block_state.processed_image.to(device=components._execution_device, dtype=components.vae.dtype) + generator = block_state.generator + if isinstance(generator, list): + image_latents = [ + components.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i]) for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents) + else: + image_latents = components.vae.encode(image).latent_dist.sample(generator) + + image_latents = image_latents.to(torch.float32) + patch = components.patch_size + latent_height, latent_width = image_latents.shape[-2:] + grid_h, grid_w = latent_height // patch, latent_width // patch + image_latents = image_latents.view( + image_latents.shape[0], image_latents.shape[1], grid_h, patch, grid_w, patch + ) + image_latents = image_latents.permute(0, 2, 4, 3, 5, 1).reshape(image_latents.shape[0], grid_h * grid_w, -1) + + bn_mean = components.vae.bn.running_mean.view(1, 1, -1).to( + device=image_latents.device, dtype=image_latents.dtype + ) + bn_std = torch.sqrt(components.vae.bn.running_var + components.vae.config.batch_norm_eps).view(1, 1, -1) + block_state.image_latents = (image_latents - bn_mean) / bn_std.to( + device=image_latents.device, dtype=image_latents.dtype + ) + + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Ideogram4PromptUpsampleStep(ModularPipelineBlocks): """ @@ -283,7 +558,8 @@ def _get_text_encoder_hidden_states( def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - device = components._execution_device + # The component-level offload hook does not run because the encoder submodules are called directly below. + device = components.text_encoder.device tokenizer = components.tokenizer max_text_tokens = block_state.max_sequence_length @@ -319,6 +595,7 @@ def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) - ) text_features = torch.stack(selected, dim=0).permute(1, 2, 3, 0).reshape(batch_size, max_text_tokens, -1) text_features = (text_features * attention_mask.to(text_features.dtype).unsqueeze(-1)).to(torch.float32) + text_features = text_features.to(components._execution_device) block_state.text_features = text_features block_state.text_lengths = text_lengths diff --git a/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py b/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py index 0b788fe236be..c9f989219c76 100644 --- a/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py +++ b/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py @@ -13,26 +13,165 @@ # limitations under the License. -from ...utils import logging -from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict, OutputParam from .before_denoise import ( + Ideogram4ApplyStrengthStep, + Ideogram4ImageInputsStep, + Ideogram4MaskInputsStep, Ideogram4PrepareAdditionalInputsStep, Ideogram4PrepareLatentsStep, + Ideogram4PrepareLatentsWithStrengthStep, + Ideogram4PrepareMaskLatentsStep, Ideogram4SetTimestepsStep, Ideogram4TextInputsStep, ) -from .decoders import Ideogram4DecodeStep -from .denoise import Ideogram4AfterDenoiseStep, Ideogram4DenoiseStep -from .encoders import Ideogram4PromptUpsampleStep, Ideogram4TextEncoderStep +from .decoders import Ideogram4DecodeStep, Ideogram4InpaintDecodeStep +from .denoise import Ideogram4AfterDenoiseStep, Ideogram4DenoiseStep, Ideogram4InpaintDenoiseStep +from .encoders import ( + Ideogram4InpaintProcessImagesInputStep, + Ideogram4ProcessImageInputStep, + Ideogram4PromptUpsampleStep, + Ideogram4TextEncoderStep, + Ideogram4VaeEncoderStep, +) + + +# auto_docstring +class Ideogram4Img2ImgVaeEncoderStep(SequentialPipelineBlocks): + """ + Preprocess and encode an image into packed Ideogram4 latents for image-to-image generation. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + processed_image (`Tensor`): + The image tensor resized and normalized for VAE encoding. + height (`int`): + The resolved image height in pixels. + width (`int`): + The resolved image width in pixels. + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "ideogram4" + block_classes = [Ideogram4ProcessImageInputStep(), Ideogram4VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Preprocess and encode an image into packed Ideogram4 latents for image-to-image generation." + + +# auto_docstring +class Ideogram4InpaintVaeEncoderStep(SequentialPipelineBlocks): + """ + Preprocess an image and mask, then encode the image into packed Ideogram4 latents for inpainting. + + Components: + image_processor (`VaeImageProcessor`) image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLFlux2`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`): + Mask image for inpainting. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + processed_image (`Tensor`): + The image tensor resized and normalized for VAE encoding. + processed_mask_image (`Tensor`): + The binary mask tensor resized to the generation resolution. + mask_overlay_kwargs (`dict`): + Arguments used to composite a cropped inpaint result over the source image. + height (`int`): + The resolved image height in pixels. + width (`int`): + The resolved image width in pixels. + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "ideogram4" + block_classes = [Ideogram4InpaintProcessImagesInputStep(), Ideogram4VaeEncoderStep()] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Preprocess an image and mask, then encode the image into packed Ideogram4 latents for inpainting." + + +# auto_docstring +class Ideogram4AutoVaeEncoderStep(AutoPipelineBlocks): + """ + Encode image inputs for Ideogram4 image-to-image and inpainting workflows. The step is skipped for text-to-image + generation. + Components: + image_processor (`VaeImageProcessor`) image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLFlux2`) -logger = logging.get_logger(__name__) # pylint: disable=invalid-name + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`, *optional*): + Mask image for inpainting. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + processed_image (`Tensor`): + The image tensor resized and normalized for VAE encoding. + processed_mask_image (`Tensor`): + The binary mask tensor resized to the generation resolution. + mask_overlay_kwargs (`dict`): + Arguments used to composite a cropped inpaint result over the source image. + height (`int`): + The resolved image height in pixels. + width (`int`): + The resolved image width in pixels. + image_latents (`Tensor`): + The latent representation of the input image. + """ + block_classes = [Ideogram4InpaintVaeEncoderStep, Ideogram4Img2ImgVaeEncoderStep] + block_names = ["inpaint", "img2img"] + block_trigger_inputs = ["mask_image", "image"] -# Core denoise: consumes the per-prompt text features and produces the unpatchified latents -# (batch/latents/timesteps/ids inputs -> denoising loop -> unpatchify). -CORE_DENOISE_BLOCKS = InsertableDict( + @property + def description(self) -> str: + return ( + "Encode image inputs for Ideogram4 image-to-image and inpainting workflows. The step is skipped for " + "text-to-image generation." + ) + + +TEXT2IMAGE_DENOISE_BLOCKS = InsertableDict( [ ("input", Ideogram4TextInputsStep()), ("prepare_latents", Ideogram4PrepareLatentsStep()), @@ -47,9 +186,7 @@ # auto_docstring class Ideogram4CoreDenoiseStep(SequentialPipelineBlocks): """ - Core denoising workflow for Ideogram4 text-to-image: prepares the batch/latents/timesteps and the packed denoiser - inputs, runs the asymmetric-CFG denoising loop over the conditional and unconditional transformers, and - unpatchifies the result for the decoder. + Core Ideogram4 text-to-image denoising workflow. Components: transformer (`Ideogram4Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) @@ -83,62 +220,319 @@ class Ideogram4CoreDenoiseStep(SequentialPipelineBlocks): Outputs: latents (`Tensor`): - Unpatchified (B, ae_channels, H, W) latents. + Unpatchified latents ready for the VAE decoder. """ model_name = "ideogram4" - block_classes = list(CORE_DENOISE_BLOCKS.values()) - block_names = list(CORE_DENOISE_BLOCKS.keys()) + block_classes = list(TEXT2IMAGE_DENOISE_BLOCKS.values()) + block_names = list(TEXT2IMAGE_DENOISE_BLOCKS.keys()) @property def description(self) -> str: - return ( - "Core denoising workflow for Ideogram4 text-to-image: prepares the batch/latents/timesteps and the packed " - "denoiser inputs, runs the asymmetric-CFG denoising loop over the conditional and unconditional " - "transformers, and unpatchifies the result for the decoder." - ) + return "Core Ideogram4 text-to-image denoising workflow." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents", description="Unpatchified latents ready for the VAE decoder.")] + + +IMAGE2IMAGE_DENOISE_BLOCKS = InsertableDict( + [ + ("text_inputs", Ideogram4TextInputsStep()), + ("image_inputs", Ideogram4ImageInputsStep()), + ("prepare_latents", Ideogram4PrepareLatentsStep()), + ("set_timesteps", Ideogram4SetTimestepsStep()), + ("apply_strength", Ideogram4ApplyStrengthStep()), + ("prepare_image_latents", Ideogram4PrepareLatentsWithStrengthStep()), + ("prepare_additional_inputs", Ideogram4PrepareAdditionalInputsStep()), + ("denoise", Ideogram4DenoiseStep()), + ("after_denoise", Ideogram4AfterDenoiseStep()), + ] +) + + +# auto_docstring +class Ideogram4Img2ImgCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core Ideogram4 image-to-image denoising workflow with strength-based latent initialization. + + Components: + transformer (`Ideogram4Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + unconditional_transformer (`Ideogram4Transformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + text_features (`Tensor`): + Per-prompt text features from the encoder. + text_lengths (`list`): + Per-prompt text-token counts from the encoder. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + mu (`float`, *optional*, defaults to 0.0): + Base mean of the logit-normal schedule. + std (`float`, *optional*, defaults to 1.5): + Std of the logit-normal schedule. + guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): + Per-step guidance scale schedule (length num_inference_steps). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + + Outputs: + latents (`Tensor`): + Unpatchified latents ready for the VAE decoder. + """ + + model_name = "ideogram4" + block_classes = list(IMAGE2IMAGE_DENOISE_BLOCKS.values()) + block_names = list(IMAGE2IMAGE_DENOISE_BLOCKS.keys()) + + @property + def description(self) -> str: + return "Core Ideogram4 image-to-image denoising workflow with strength-based latent initialization." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents", description="Unpatchified latents ready for the VAE decoder.")] + + +INPAINT_DENOISE_BLOCKS = InsertableDict( + [ + ("text_inputs", Ideogram4TextInputsStep()), + ("image_inputs", Ideogram4ImageInputsStep()), + ("mask_inputs", Ideogram4MaskInputsStep()), + ("prepare_latents", Ideogram4PrepareLatentsStep()), + ("set_timesteps", Ideogram4SetTimestepsStep()), + ("apply_strength", Ideogram4ApplyStrengthStep()), + ("prepare_image_latents", Ideogram4PrepareLatentsWithStrengthStep()), + ("prepare_mask_latents", Ideogram4PrepareMaskLatentsStep()), + ("prepare_additional_inputs", Ideogram4PrepareAdditionalInputsStep()), + ("denoise", Ideogram4InpaintDenoiseStep()), + ("after_denoise", Ideogram4AfterDenoiseStep()), + ] +) + + +# auto_docstring +class Ideogram4InpaintCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core Ideogram4 inpaint denoising workflow with latent-mask blending at every step. + + Components: + transformer (`Ideogram4Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + unconditional_transformer (`Ideogram4Transformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + text_features (`Tensor`): + Per-prompt text features from the encoder. + text_lengths (`list`): + Per-prompt text-token counts from the encoder. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`): + The binary mask tensor resized to the generation resolution. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + mu (`float`, *optional*, defaults to 0.0): + Base mean of the logit-normal schedule. + std (`float`, *optional*, defaults to 1.5): + Std of the logit-normal schedule. + guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): + Per-step guidance scale schedule (length num_inference_steps). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + + Outputs: + latents (`Tensor`): + Unpatchified latents ready for the VAE decoder. + """ + + model_name = "ideogram4" + block_classes = list(INPAINT_DENOISE_BLOCKS.values()) + block_names = list(INPAINT_DENOISE_BLOCKS.keys()) + + @property + def description(self) -> str: + return "Core Ideogram4 inpaint denoising workflow with latent-mask blending at every step." @property def outputs(self) -> list[OutputParam]: - # The only meaningful product of the core step is the unpatchified latents; the batch/timesteps/packed-sequence - # inputs prepared along the way are consumed within the loop and are not updated by it. - return [OutputParam.template("latents", description="Unpatchified (B, ae_channels, H, W) latents.")] + return [OutputParam.template("latents", description="Unpatchified latents ready for the VAE decoder.")] + + +# auto_docstring +class Ideogram4AutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Select the Ideogram4 text-to-image, image-to-image, or inpaint denoising workflow. + + Components: + transformer (`Ideogram4Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + unconditional_transformer (`Ideogram4Transformer2DModel`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + text_features (`Tensor`): + Per-prompt text features from the encoder. + text_lengths (`list`): + Per-prompt text-token counts from the encoder. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The binary mask tensor resized to the generation resolution. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 48): + The number of denoising steps. + mu (`float`, *optional*, defaults to 0.0): + Base mean of the logit-normal schedule. + std (`float`, *optional*, defaults to 1.5): + Std of the logit-normal schedule. + guidance_schedule (`list`, *optional*, defaults to (7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, + 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): + Per-step guidance scale schedule (length num_inference_steps). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + + Outputs: + latents (`Tensor`): + Unpatchified latents ready for the VAE decoder. + """ + + block_classes = [Ideogram4InpaintCoreDenoiseStep, Ideogram4Img2ImgCoreDenoiseStep, Ideogram4CoreDenoiseStep] + block_names = ["inpaint", "img2img", "text2img"] + block_trigger_inputs = ["processed_mask_image", "image_latents"] + default_block_name = "text2img" + + def select_block(self, processed_mask_image=None, image_latents=None) -> str | None: + if processed_mask_image is not None: + return "inpaint" + if image_latents is not None: + return "img2img" + return None + + @property + def description(self) -> str: + return "Select the Ideogram4 text-to-image, image-to-image, or inpaint denoising workflow." + + +# auto_docstring +class Ideogram4AutoDecodeStep(AutoPipelineBlocks): + """ + Decode Ideogram4 latents and apply the optional cropped-inpaint overlay. + + Components: + vae (`AutoencoderKLFlux2`) image_mask_processor (`InpaintProcessor`) image_processor (`VaeImageProcessor`) + + Inputs: + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + latents (`Tensor`): + The unpatchified latents to decode. + mask_overlay_kwargs (`dict`, *optional*): + Arguments used to composite a cropped inpaint result over the source image. + + Outputs: + images (`list`): + Generated images. + """ + + block_classes = [Ideogram4InpaintDecodeStep, Ideogram4DecodeStep] + block_names = ["inpaint", "default"] + block_trigger_inputs = ["mask_overlay_kwargs", None] + + @property + def description(self) -> str: + return "Decode Ideogram4 latents and apply the optional cropped-inpaint overlay." + + +AUTO_BLOCKS = InsertableDict( + [ + ("vae_encoder", Ideogram4AutoVaeEncoderStep()), + ("prompt_upsample", Ideogram4PromptUpsampleStep()), + ("text_encoder", Ideogram4TextEncoderStep()), + ("denoise", Ideogram4AutoCoreDenoiseStep()), + ("decode", Ideogram4AutoDecodeStep()), + ] +) # auto_docstring class Ideogram4AutoBlocks(SequentialPipelineBlocks): """ - Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> encode text -> - core denoise (asymmetric CFG over two transformers) -> decode. + Auto Modular pipeline for Ideogram4 text-to-image, image-to-image, and inpainting workflows. Supported workflows: - `text2image`: requires `prompt` + - `image2image`: requires `prompt`, `image` + - `inpainting`: requires `prompt`, `image`, `mask_image` Components: + image_processor (`VaeImageProcessor`) image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLFlux2`) text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired with the text encoder. prompt_enhancer_head (`Ideogram4PromptEnhancerHead`): LM head grafted onto the text encoder for prompt upsampling. transformer (`Ideogram4Transformer2DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) unconditional_transformer (`Ideogram4Transformer2DModel`) vae - (`AutoencoderKLFlux2`) image_processor (`VaeImageProcessor`) + (`FlowMatchEulerDiscreteScheduler`) unconditional_transformer (`Ideogram4Transformer2DModel`) Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`, *optional*): + Mask image for inpainting. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. prompt (`str`): The prompt or prompts to guide image generation. prompt_upsampling (`bool`, *optional*, defaults to False): If True, rewrite the prompt into Ideogram4's native JSON caption before encoding. prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): Sampling temperature for prompt upsampling. - height (`int`, *optional*): - The height in pixels of the generated image. - width (`int`, *optional*): - The width in pixels of the generated image. max_sequence_length (`int`, *optional*, defaults to 2048): Maximum sequence length for prompt encoding. - generator (`Generator`, *optional*): - Torch generator for deterministic generation. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - latents (`Tensor`, *optional*): + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The binary mask tensor resized to the generation resolution. + latents (`Tensor`): Pre-generated noisy latents for image generation. num_inference_steps (`int`, *optional*, defaults to 48): The number of denoising steps. @@ -150,8 +544,12 @@ class Ideogram4AutoBlocks(SequentialPipelineBlocks): 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 7.0, 3.0, 3.0, 3.0)): Per-step guidance scale schedule (length num_inference_steps). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + Arguments used to composite a cropped inpaint result over the source image. Outputs: images (`list`): @@ -159,26 +557,18 @@ class Ideogram4AutoBlocks(SequentialPipelineBlocks): """ model_name = "ideogram4" - block_classes = [ - Ideogram4PromptUpsampleStep(), - Ideogram4TextEncoderStep(), - Ideogram4CoreDenoiseStep(), - Ideogram4DecodeStep(), - ] - block_names = ["prompt_upsample", "text_encoder", "denoise", "decode"] + block_classes = list(AUTO_BLOCKS.values()) + block_names = list(AUTO_BLOCKS.keys()) - # Workflow map declaring the trigger conditions for each supported workflow. - # `True` means the workflow triggers when the input is not None. _workflow_map = { "text2image": {"prompt": True}, + "image2image": {"prompt": True, "image": True}, + "inpainting": {"prompt": True, "image": True, "mask_image": True}, } @property def description(self) -> str: - return ( - "Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> " - "encode text -> core denoise (asymmetric CFG over two transformers) -> decode." - ) + return "Auto Modular pipeline for Ideogram4 text-to-image, image-to-image, and inpainting workflows." @property def outputs(self) -> list[OutputParam]: diff --git a/tests/modular_pipelines/ideogram4/__init__.py b/tests/modular_pipelines/ideogram4/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/modular_pipelines/ideogram4/test_modular_pipeline_ideogram4.py b/tests/modular_pipelines/ideogram4/test_modular_pipeline_ideogram4.py new file mode 100644 index 000000000000..360a5be25a1a --- /dev/null +++ b/tests/modular_pipelines/ideogram4/test_modular_pipeline_ideogram4.py @@ -0,0 +1,212 @@ +# 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. + + +import numpy as np +import PIL.Image +import torch +from transformers import Qwen2Tokenizer, Qwen3VLConfig, Qwen3VLModel, Qwen3VLTextConfig, Qwen3VLVisionConfig + +from diffusers import ( + AutoencoderKLFlux2, + FlowMatchEulerDiscreteScheduler, + Ideogram4AutoBlocks, + Ideogram4ModularPipeline, + Ideogram4Transformer2DModel, +) + +from ...testing_utils import enable_full_determinism +from ..test_modular_pipelines_common import ModularPipelineTesterMixin + + +enable_full_determinism() + + +IDEOGRAM4_WORKFLOWS = { + "text2image": [ + ("prompt_upsample", "Ideogram4PromptUpsampleStep"), + ("text_encoder", "Ideogram4TextEncoderStep"), + ("denoise.input", "Ideogram4TextInputsStep"), + ("denoise.prepare_latents", "Ideogram4PrepareLatentsStep"), + ("denoise.set_timesteps", "Ideogram4SetTimestepsStep"), + ("denoise.prepare_additional_inputs", "Ideogram4PrepareAdditionalInputsStep"), + ("denoise.denoise", "Ideogram4DenoiseStep"), + ("denoise.after_denoise", "Ideogram4AfterDenoiseStep"), + ("decode", "Ideogram4DecodeStep"), + ], + "image2image": [ + ("vae_encoder.preprocess", "Ideogram4ProcessImageInputStep"), + ("vae_encoder.encode", "Ideogram4VaeEncoderStep"), + ("prompt_upsample", "Ideogram4PromptUpsampleStep"), + ("text_encoder", "Ideogram4TextEncoderStep"), + ("denoise.text_inputs", "Ideogram4TextInputsStep"), + ("denoise.image_inputs", "Ideogram4ImageInputsStep"), + ("denoise.prepare_latents", "Ideogram4PrepareLatentsStep"), + ("denoise.set_timesteps", "Ideogram4SetTimestepsStep"), + ("denoise.apply_strength", "Ideogram4ApplyStrengthStep"), + ("denoise.prepare_image_latents", "Ideogram4PrepareLatentsWithStrengthStep"), + ("denoise.prepare_additional_inputs", "Ideogram4PrepareAdditionalInputsStep"), + ("denoise.denoise", "Ideogram4DenoiseStep"), + ("denoise.after_denoise", "Ideogram4AfterDenoiseStep"), + ("decode", "Ideogram4DecodeStep"), + ], + "inpainting": [ + ("vae_encoder.preprocess", "Ideogram4InpaintProcessImagesInputStep"), + ("vae_encoder.encode", "Ideogram4VaeEncoderStep"), + ("prompt_upsample", "Ideogram4PromptUpsampleStep"), + ("text_encoder", "Ideogram4TextEncoderStep"), + ("denoise.text_inputs", "Ideogram4TextInputsStep"), + ("denoise.image_inputs", "Ideogram4ImageInputsStep"), + ("denoise.mask_inputs", "Ideogram4MaskInputsStep"), + ("denoise.prepare_latents", "Ideogram4PrepareLatentsStep"), + ("denoise.set_timesteps", "Ideogram4SetTimestepsStep"), + ("denoise.apply_strength", "Ideogram4ApplyStrengthStep"), + ("denoise.prepare_image_latents", "Ideogram4PrepareLatentsWithStrengthStep"), + ("denoise.prepare_mask_latents", "Ideogram4PrepareMaskLatentsStep"), + ("denoise.prepare_additional_inputs", "Ideogram4PrepareAdditionalInputsStep"), + ("denoise.denoise", "Ideogram4InpaintDenoiseStep"), + ("denoise.after_denoise", "Ideogram4AfterDenoiseStep"), + ("decode", "Ideogram4InpaintDecodeStep"), + ], +} + + +def get_dummy_components(): + torch.manual_seed(0) + transformer = Ideogram4Transformer2DModel( + in_channels=16, + num_layers=1, + attention_head_dim=8, + num_attention_heads=2, + intermediate_size=32, + adaln_dim=8, + llm_features_dim=52, + rope_theta=10_000, + mrope_section=(2, 1, 1), + ).eval() + + torch.manual_seed(0) + unconditional_transformer = Ideogram4Transformer2DModel( + in_channels=16, + num_layers=1, + attention_head_dim=8, + num_attention_heads=2, + intermediate_size=32, + adaln_dim=8, + llm_features_dim=52, + rope_theta=10_000, + mrope_section=(2, 1, 1), + ).eval() + + torch.manual_seed(0) + vae = AutoencoderKLFlux2( + block_out_channels=(8, 8, 8, 8), + decoder_block_out_channels=(8, 8, 8, 8), + layers_per_block=1, + latent_channels=4, + norm_num_groups=4, + sample_size=32, + mid_block_add_attention=False, + patch_size=(2, 2), + ).eval() + + text_config = Qwen3VLTextConfig( + vocab_size=152064, + hidden_size=4, + intermediate_size=8, + num_hidden_layers=36, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=4, + max_position_embeddings=128, + use_cache=False, + ) + vision_config = Qwen3VLVisionConfig( + depth=1, + hidden_size=4, + intermediate_size=8, + num_heads=1, + patch_size=2, + spatial_merge_size=1, + temporal_patch_size=1, + out_hidden_size=4, + deepstack_visual_indexes=(), + ) + torch.manual_seed(0) + text_encoder = Qwen3VLModel(Qwen3VLConfig(text_config=text_config, vision_config=vision_config)).eval() + tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "transformer": transformer, + "unconditional_transformer": unconditional_transformer, + "vae": vae, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "scheduler": scheduler, + } + + +def get_dummy_image(seed=0): + image = np.random.default_rng(seed).integers(0, 256, (32, 32, 3), dtype=np.uint8) + return PIL.Image.fromarray(image) + + +class TestIdeogram4ModularPipelineFast(ModularPipelineTesterMixin): + pipeline_class = Ideogram4ModularPipeline + pipeline_blocks_class = Ideogram4AutoBlocks + pretrained_model_name_or_path = "hf-internal-testing/tiny-ideogram4-modular-pipe" + + params = frozenset(["prompt", "height", "width", "image", "mask_image"]) + batch_params = frozenset(["prompt", "image", "mask_image"]) + expected_workflow_blocks = IDEOGRAM4_WORKFLOWS + + def get_pipeline(self, components_manager=None, dtype=torch.float32): + pipe = self.pipeline_blocks_class().init_pipeline(components_manager=components_manager) + pipe.update_components(**get_dummy_components()) + pipe.to(dtype=dtype) + pipe.set_progress_bar_config(disable=True) + return pipe + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "cat wizard", + "generator": self.get_generator(seed), + "num_inference_steps": 2, + "guidance_schedule": [1.0, 1.0], + "height": 32, + "width": 32, + "max_sequence_length": 32, + "output_type": "pt", + } + + def test_img2img_and_inpaint(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.update({"image": get_dummy_image(), "strength": 1.0}) + + image = pipe(**inputs, output="images") + assert image.shape == (1, 3, 32, 32) + assert not torch.isnan(image).any() + + inputs["generator"] = self.get_generator(0) + mask = np.zeros((32, 32), dtype=np.uint8) + mask[8:24, 8:24] = 255 + inputs["mask_image"] = PIL.Image.fromarray(mask) + inputs["padding_mask_crop"] = 2 + inputs["output_type"] = "pil" + image = pipe(**inputs, output="images") + assert len(image) == 1 + assert image[0].size == (32, 32) + assert np.isfinite(np.asarray(image[0])).all()