# Copyright 2026 OpenMOSS and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.

import torch

from diffusers import AutoModel
from diffusers.configuration_utils import FrozenDict
from diffusers.guiders import ClassifierFreeGuidance
from diffusers.modular_pipelines import (
    BlockState,
    LoopSequentialPipelineBlocks,
    ModularPipelineBlocks,
    PipelineState,
)
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler


class MOSSImageLoopBeforeDenoiser(ModularPipelineBlocks):
    model_name = "moss-image"

    @property
    def description(self) -> str:
        return "Step within the denoising loop that casts the float32 latents to the transformer dtype."

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "latents",
                required=True,
                type_hint=torch.Tensor,
                description="The latents being denoised. Can be generated in the prepare latents step.",
            ),
            InputParam(
                "dtype",
                required=True,
                type_hint=torch.dtype,
                description="Dtype of the transformer inputs. Can be generated in the input step.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor):
        block_state.latent_model_input = block_state.latents.to(block_state.dtype)
        return components, block_state


class MOSSImageLoopDenoiser(ModularPipelineBlocks):
    model_name = "moss-image"

    # transformer argument -> (conditional, unconditional) block_state fields
    guider_input_fields = {"encoder_hidden_states": ("prompt_embeds", "negative_prompt_embeds")}

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 3.0}),
                default_creation_method="from_config",
            ),
            ComponentSpec("transformer", AutoModel),
        ]

    @property
    def description(self) -> str:
        return (
            "Step within the denoising loop that runs the transformer for each guidance batch and combines the "
            "predictions with the guider. The unconditional prediction skips the transformer blocks in "
            "`uncond_skip_layers`."
        )

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "num_inference_steps",
                required=True,
                type_hint=int,
                description="The number of denoising steps. Can be generated in the set timesteps step.",
            ),
            InputParam(
                "uncond_skip_layers",
                type_hint=list[int],
                default=[22],
                description="Zero-based transformer blocks skipped in the unconditional prediction. Pass `[]` to "
                "run the full model for it.",
            ),
            InputParam.template("attention_kwargs"),
            InputParam("prompt_embeds", required=True, type_hint=list[torch.Tensor]),
            InputParam("negative_prompt_embeds", type_hint=list[torch.Tensor]),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor):
        guider = components.guider
        guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
        guider_state = guider.prepare_inputs_from_block_state(block_state, self.guider_input_fields)

        for guider_state_batch in guider_state:
            guider.prepare_models(components.transformer)
            is_uncond = getattr(guider_state_batch, guider._identifier_key) == "pred_uncond"
            skip_layers = block_state.uncond_skip_layers if is_uncond and block_state.uncond_skip_layers else None
            model_out = components.transformer(
                hidden_states=block_state.latent_model_input,
                timestep=t,
                encoder_hidden_states=[e.to(block_state.dtype) for e in guider_state_batch.encoder_hidden_states],
                skip_layers=skip_layers,
                attention_kwargs=block_state.attention_kwargs,
                return_dict=False,
            )[0]
            # guidance and the Euler update run in float32
            guider_state_batch.noise_pred = model_out.float()
            guider.cleanup_models(components.transformer)

        block_state.noise_pred = guider(guider_state)[0]
        return components, block_state


class MOSSImageLoopAfterDenoiser(ModularPipelineBlocks):
    model_name = "moss-image"

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

    @property
    def description(self) -> str:
        return "Step within the denoising loop that updates the float32 latents with the scheduler."

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor):
        block_state.latents = components.scheduler.step(
            block_state.noise_pred, t, block_state.latents, return_dict=False
        )[0]
        return components, block_state


class MOSSImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks):
    model_name = "moss-image"

    @property
    def description(self) -> str:
        return (
            "Pipeline block that iteratively denoises the latents over `timesteps`. "
            "The steps of each iteration are defined by `sub_blocks`."
        )

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

    @property
    def loop_inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "timesteps",
                required=True,
                type_hint=torch.Tensor,
                description="The timesteps of the denoising loop. Can be generated in the set timesteps step.",
            ),
            InputParam(
                "num_inference_steps",
                required=True,
                type_hint=int,
                description="The number of denoising steps. Can be generated in the set timesteps step.",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
            for i, t in enumerate(block_state.timesteps):
                components, block_state = self.loop_step(components, block_state, i=i, t=t)
                progress_bar.update()
        self.set_block_state(state, block_state)
        return components, state


class MOSSImageDenoiseStep(MOSSImageDenoiseLoopWrapper):
    block_classes = [MOSSImageLoopBeforeDenoiser, MOSSImageLoopDenoiser, MOSSImageLoopAfterDenoiser]
    block_names = ["before_denoiser", "denoiser", "after_denoiser"]

    @property
    def description(self) -> str:
        return (
            "Denoise step that iteratively denoises the latents. At each iteration it runs, in order: "
            "`MOSSImageLoopBeforeDenoiser`, `MOSSImageLoopDenoiser`, `MOSSImageLoopAfterDenoiser`."
        )