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

from diffusers.utils import logging
from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import (
    Krea2EditRoPEInputsStep,
    Krea2PrepareLatentsStep,
    Krea2SetTimestepsStep,
)
from .decoders import Krea2AfterDenoiseStep
from .denoise import Krea2EditDenoiseStep
from .encoders import (
    Krea2EditReferenceLatentsStep,
    Krea2EditTextEncoderStep,
)
from .inputs import Krea2TextInputsStep
from .modular_blocks_krea2 import Krea2DecodeStep


logger = logging.get_logger(__name__)


# Krea 2 (reference-image edit)
class Krea2EditCoreDenoiseStep(SequentialPipelineBlocks):
    model_name = "krea2"
    block_classes = [
        Krea2TextInputsStep(),
        Krea2PrepareLatentsStep(),
        Krea2SetTimestepsStep(),
        Krea2EditRoPEInputsStep(),
        Krea2EditDenoiseStep(),
        Krea2AfterDenoiseStep(),
    ]
    block_names = [
        "input",
        "prepare_latents",
        "set_timesteps",
        "prepare_rope_inputs",
        "denoise",
        "after_denoise",
    ]

    @property
    def description(self):
        return (
            "step that denoises noise into an image for the reference-image edit task. It includes the denoise loop "
            "(with the clean reference tokens appended each step), as well as preparing the inputs (timesteps, "
            "latents, rope inputs etc.)."
        )

    @property
    def outputs(self):
        return [
            OutputParam.template("latents"),
        ]


# Preset: reference-image edit. `image` carries the reference image(s); without it, use the text2image preset.
EDIT_BLOCKS = InsertableDict(
    [
        ("text_encoder", Krea2EditTextEncoderStep()),
        ("reference_latents", Krea2EditReferenceLatentsStep()),
        ("denoise", Krea2EditCoreDenoiseStep()),
        ("decode", Krea2DecodeStep()),
    ]
)


class Krea2EditBlocks(SequentialPipelineBlocks):
    """
    Modular pipeline for the Krea 2 reference-image edit task.

    Pass one or more reference images as `image`; they condition generation both through the Qwen3-VL text encoder
    and as clean VAE latents appended to the transformer sequence at flow time t=0 (matching the Ostris AI-Toolkit
    edit LoRAs). Load an edit/style LoRA into the transformer to control the effect.
    """

    model_name = "krea2"
    block_classes = EDIT_BLOCKS.values()
    block_names = EDIT_BLOCKS.keys()

    @property
    def description(self):
        return "Auto Modular pipeline for the reference-image edit task using Krea 2. Requires `prompt` and `image`."