# 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 image-to-image support to Ideogram4.

Reuses the in-repo text2image blocks and adds the img2img-specific steps (VAE encode, strength-trimmed
timesteps, noise scaling). Assembled into a single AutoBlocks that serves both text2image and image2image,
selected by the presence of `image` — the canonical modular-pipeline shape.
"""

import torch

from diffusers.configuration_utils import FrozenDict
from diffusers.image_processor import VaeImageProcessor
from diffusers.models import AutoencoderKLFlux2
from diffusers.modular_pipelines.ideogram4.before_denoise import (
    DEFAULT_GUIDANCE_SCHEDULE,
    Ideogram4PrepareAdditionalInputsStep,
    Ideogram4PrepareLatentsStep,
    Ideogram4TextInputsStep,
    _expand_tensor_to_effective_batch,
    _logit_normal_sigmas,
    _resolution_aware_mu,
)
from diffusers.modular_pipelines.ideogram4.decoders import Ideogram4DecodeStep
from diffusers.modular_pipelines.ideogram4.denoise import Ideogram4AfterDenoiseStep, Ideogram4DenoiseStep
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 (
    AutoPipelineBlocks,
    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


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


# The Ideogram4 VAE downscales by 8 and the transformer patchifies by 2, so the pixel grid is downscaled by 16
# in each axis; images are resized to a multiple of that before encoding.
IMAGE_PROCESSOR_SPEC = ComponentSpec(
    "image_processor",
    VaeImageProcessor,
    config=FrozenDict({"vae_scale_factor": 16}),
    default_creation_method="from_config",
)


# Modified from diffusers.pipelines.stable_diffusion_3.pipeline_stable_diffusion_3_img2img.StableDiffusion3Img2ImgPipeline.get_timesteps
def get_timesteps(scheduler, num_inference_steps, strength):
    """Trim the schedule to the last `strength` fraction of steps (img2img skips the high-noise steps)."""
    init_timestep = min(num_inference_steps * strength, num_inference_steps)
    t_start = int(max(num_inference_steps - init_timestep, 0))
    timesteps = scheduler.timesteps[t_start * scheduler.order :]
    if hasattr(scheduler, "set_begin_index"):
        scheduler.set_begin_index(t_start * scheduler.order)
    return timesteps, num_inference_steps - t_start, t_start


# auto_docstring
class Ideogram4VaeEncoderStep(ModularPipelineBlocks):
    """
    Image-to-image VAE encoder step: resize/preprocess the reference `image`, encode it through the VAE, patchify to
    the transformer's packed layout and normalize with the VAE batch-norm statistics (the exact inverse of the
    decoder). Resolves `height`/`width` from the image when not provided.
    """

    model_name = "ideogram4"

    @property
    def description(self) -> str:
        return (
            "Image-to-image VAE encoder step: preprocess the reference `image`, encode through the VAE, patchify to "
            "the packed transformer layout and normalize with the VAE batch-norm statistics (inverse of the decoder)."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("vae", AutoencoderKLFlux2), IMAGE_PROCESSOR_SPEC]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("image", required=True),
            InputParam.template("height"),
            InputParam.template("width"),
            InputParam.template("generator"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="image_latents",
                type_hint=torch.Tensor,
                description="Packed, bn-normalized latents of the reference image (B, num_image_tokens, latent_dim).",
            ),
            OutputParam(name="height", type_hint=int, description="Target height, resolved from the image if unset."),
            OutputParam(name="width", type_hint=int, description="Target width, resolved from the image if unset."),
        ]

    @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

        # Resolve target size (multiple of vae_scale_factor * patch) and preprocess to it.
        height, width = components.image_processor.get_default_height_width(
            block_state.image, block_state.height, block_state.width
        )
        image = components.image_processor.preprocess(block_state.image, height=height, width=width)
        image = image.to(device=device, dtype=components.vae.dtype)

        # Encode, sample the posterior, then patchify to the packed layout.
        z = components.vae.encode(image).latent_dist.sample(generator=block_state.generator)
        z = z.to(torch.float32)
        ae_channels = z.shape[1]
        grid_h = z.shape[2] // patch
        grid_w = z.shape[3] // patch
        z = z.view(z.shape[0], ae_channels, grid_h, patch, grid_w, patch)
        z = z.permute(0, 2, 4, 3, 5, 1).contiguous()
        z = z.reshape(z.shape[0], grid_h * grid_w, patch * patch * ae_channels)

        # Normalize with the packed-channel batch-norm statistics (inverse of the decoder's denormalization).
        bn_mean = components.vae.bn.running_mean.view(1, 1, -1).to(device=z.device, dtype=z.dtype)
        bn_std = torch.sqrt(components.vae.bn.running_var + components.vae.config.batch_norm_eps).view(1, 1, -1)
        bn_std = bn_std.to(device=z.device, dtype=z.dtype)
        block_state.image_latents = (z - bn_mean) / bn_std

        block_state.height = height
        block_state.width = width

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


# auto_docstring
class Ideogram4ExpandImageLatentsStep(ModularPipelineBlocks):
    """
    Replicate `image_latents` from the per-image batch to the effective `batch_size` (num prompts *
    num_images_per_prompt). Place after the text-input step, which produces `batch_size`.
    """

    model_name = "ideogram4"

    @property
    def description(self) -> str:
        return "Replicate `image_latents` to the effective `batch_size` (num prompts * num_images_per_prompt)."

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("image_latents", required=True),
            InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."),
        ]

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

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

        image_batch = block_state.image_latents.shape[0]
        num_per_prompt = block_state.batch_size // image_batch
        block_state.image_latents = _expand_tensor_to_effective_batch(
            block_state.image_latents, image_batch, num_per_prompt, "image_latents"
        )

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


# auto_docstring
class Ideogram4SetTimestepsWithStrengthStep(ModularPipelineBlocks):
    """
    Set the resolution-aware logit-normal sigma schedule, then trim it to the last `strength` fraction of steps for
    image-to-image. The per-step guidance weights are trimmed to stay aligned with the retained timesteps.
    """

    model_name = "ideogram4"

    @property
    def description(self) -> str:
        return (
            "Set the resolution-aware logit-normal sigma schedule and trim it (and the guidance weights) to the last "
            "`strength` fraction of steps for image-to-image."
        )

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

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("num_inference_steps", default=48),
            InputParam.template("height", required=True),
            InputParam.template("width", required=True),
            InputParam(name="mu", default=0.0, type_hint=float, description="Base mean of the logit-normal schedule."),
            InputParam(name="std", default=1.5, type_hint=float, description="Std of the logit-normal schedule."),
            InputParam(
                name="guidance_schedule",
                default=DEFAULT_GUIDANCE_SCHEDULE,
                type_hint=list,
                description="Per-step guidance scale schedule (length num_inference_steps).",
            ),
            InputParam(
                name="strength",
                default=0.6,
                type_hint=float,
                description="How much to transform the reference image; 1.0 ignores it, lower keeps more structure.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(name="timesteps", type_hint=torch.Tensor, description="The (trimmed) denoising timesteps."),
            OutputParam(
                name="num_inference_steps",
                type_hint=int,
                description="Number of denoising steps after strength trimming.",
            ),
            OutputParam(name="gw", type_hint=torch.Tensor, description="Per-step guidance weights (trimmed)."),
        ]

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

        device = components._execution_device
        if len(block_state.guidance_schedule) != block_state.num_inference_steps:
            raise ValueError(
                f"`guidance_schedule` must have length `num_inference_steps` ({block_state.num_inference_steps}), "
                f"got {len(block_state.guidance_schedule)}."
            )

        schedule_mu = _resolution_aware_mu(height=block_state.height, width=block_state.width, base_mu=block_state.mu)
        sigmas = _logit_normal_sigmas(block_state.num_inference_steps, schedule_mu, std=block_state.std, device=device)
        components.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device)

        gw_full = torch.as_tensor(block_state.guidance_schedule, dtype=torch.float32, device=device)
        timesteps, num_inference_steps, t_start = get_timesteps(
            components.scheduler, block_state.num_inference_steps, block_state.strength
        )
        block_state.timesteps = timesteps
        block_state.num_inference_steps = num_inference_steps
        block_state.gw = gw_full[t_start:]

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


# auto_docstring
class Ideogram4PrepareLatentsWithStrengthStep(ModularPipelineBlocks):
    """
    Add noise to the reference `image_latents` at the first (trimmed) timestep to build the starting latents for
    image-to-image denoising. Run after prepare_latents (random noise) and set_timesteps.
    """

    model_name = "ideogram4"

    @property
    def description(self) -> str:
        return (
            "Add noise to the reference `image_latents` at the first (trimmed) timestep to build the starting latents "
            "for image-to-image denoising."
        )

    @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="Initial random noise from the prepare-latents step.",
            ),
            InputParam.template("image_latents", required=True),
            InputParam(
                name="timesteps",
                required=True,
                type_hint=torch.Tensor,
                description="The (trimmed) denoising timesteps from set_timesteps.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="latents",
                type_hint=torch.Tensor,
                description="Noised image latents to start image-to-image denoising.",
            )
        ]

    @staticmethod
    def check_inputs(image_latents, latents):
        if image_latents.shape != latents.shape:
            raise ValueError(
                f"`image_latents` {tuple(image_latents.shape)} must match `latents` {tuple(latents.shape)}."
            )

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

        self.check_inputs(block_state.image_latents, block_state.latents)

        latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
        image_latents = block_state.image_latents.to(block_state.latents)
        block_state.latents = components.scheduler.scale_noise(image_latents, latent_timestep, block_state.latents)

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


# Img2img input: reuse the text-input step, then batch-expand the image latents to the effective batch size.
IMG2IMG_INPUT_BLOCKS = InsertableDict(
    [
        ("text_inputs", Ideogram4TextInputsStep()),
        ("expand_image_latents", Ideogram4ExpandImageLatentsStep()),
    ]
)


class Ideogram4Img2ImgInputStep(SequentialPipelineBlocks):
    model_name = "ideogram4"
    block_classes = list(IMG2IMG_INPUT_BLOCKS.values())
    block_names = list(IMG2IMG_INPUT_BLOCKS.keys())

    @property
    def description(self) -> str:
        return (
            "Input step for image-to-image: batch-expand the text features and the reference `image_latents` to the "
            "effective batch size."
        )


# Img2img core denoise: same shape as the text2image core, but with strength-trimmed timesteps and an extra
# noise-scaling step that seeds the loop from the reference image latents.
IMG2IMG_CORE_DENOISE_BLOCKS = InsertableDict(
    [
        ("input", Ideogram4Img2ImgInputStep()),
        ("prepare_latents", Ideogram4PrepareLatentsStep()),
        ("set_timesteps", Ideogram4SetTimestepsWithStrengthStep()),
        ("prepare_additional_inputs", Ideogram4PrepareAdditionalInputsStep()),
        ("add_noise", Ideogram4PrepareLatentsWithStrengthStep()),
        ("denoise", Ideogram4DenoiseStep()),
        ("after_denoise", Ideogram4AfterDenoiseStep()),
    ]
)


# auto_docstring
class Ideogram4Img2ImgCoreDenoiseStep(SequentialPipelineBlocks):
    """
    Core denoising workflow for Ideogram4 image-to-image: prepares the batch/latents, trims the schedule by
    `strength`, seeds the loop from the reference image latents, runs the asymmetric-CFG denoising loop and
    unpatchifies the result for the decoder.
    """

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

    @property
    def description(self) -> str:
        return (
            "Core denoising workflow for Ideogram4 image-to-image: prepares the batch/latents, trims the schedule by "
            "`strength`, seeds the loop from the reference image latents, runs the asymmetric-CFG denoising loop and "
            "unpatchifies the result for the decoder."
        )

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


class Ideogram4AutoVaeEncoderStep(AutoPipelineBlocks):
    block_classes = [Ideogram4VaeEncoderStep()]
    block_names = ["img2img"]
    block_trigger_inputs = ["image"]

    @property
    def description(self) -> str:
        return (
            "VAE encoder step that encodes the reference `image` into `image_latents`. This is an auto pipeline "
            "block: it runs when `image` is provided and is skipped otherwise (text2image)."
        )


class Ideogram4AutoCoreDenoiseStep(ConditionalPipelineBlocks):
    block_classes = [Ideogram4CoreDenoiseStep, Ideogram4Img2ImgCoreDenoiseStep]
    block_names = ["text2image", "img2img"]
    block_trigger_inputs = ["image_latents"]
    default_block_name = "text2image"

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

    @property
    def description(self) -> str:
        return (
            "Core denoising step. \n"
            " - `Ideogram4Img2ImgCoreDenoiseStep` (img2img) is used when `image_latents` 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.")]


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


# auto_docstring
class Ideogram4Img2ImgAutoBlocks(SequentialPipelineBlocks):
    """
    Auto Modular pipeline for Ideogram4 supporting text-to-image and image-to-image, selected by the presence of
    `image`: (optional) prompt upsampling -> encode text -> (img2img) VAE-encode reference -> core denoise
    (asymmetric CFG over two transformers) -> decode.

      Supported workflows:
        - `text2image`: requires `prompt`
        - `image2image`: requires `prompt`, `image`
    """

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

    _workflow_map = {
        "text2image": {"prompt": True},
        "image2image": {"prompt": True, "image": True},
    }

    @property
    def description(self) -> str:
        return (
            "Auto Modular pipeline for Ideogram4 text-to-image and image-to-image: (optional) prompt upsampling -> "
            "encode text -> (img2img) VAE-encode the reference image -> core denoise (asymmetric CFG over two "
            "transformers) -> decode. The workflow is selected by the presence of `image`."
        )

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