# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
"""
Pipeline stages for diffusion models.

This package contains the various stages that can be composed to create
complete diffusion pipelines.
"""

from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import (
    CausalDMDDenoisingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.comfyui_latent_preparation import (
    ComfyUILatentPreparationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import (
    LTX2AVDecodingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import (
    LTX2AVDenoisingStage,
    LTX2RefinementStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import (
    DmdDenoisingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.encoding import EncodingStage

# Hunyuan3D paint stages
from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_paint import (
    Hunyuan3DPaintPostprocessStage,
    Hunyuan3DPaintPreprocessStage,
    Hunyuan3DPaintTexGenStage,
)

# Hunyuan3D shape stages
from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import (
    Hunyuan3DShapeBeforeDenoisingStage,
    Hunyuan3DShapeDenoisingStage,
    Hunyuan3DShapeExportStage,
    Hunyuan3DShapeSaveStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
    ImageEncodingStage,
    ImageVAEEncodingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
    InputValidationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
    LatentPreparationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import (
    LTX2AVLatentPreparationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import (
    LTX2TextConnectorStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
    TextEncodingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import (
    TimestepPreparationStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.upsampling import (
    LTX2HalveResolutionStage,
    LTX2LoRASwitchStage,
    LTX2UpsampleStage,
)

__all__ = [
    "PipelineStage",
    "InputValidationStage",
    "TimestepPreparationStage",
    "LatentPreparationStage",
    "ComfyUILatentPreparationStage",
    "LTX2AVLatentPreparationStage",
    "DenoisingStage",
    "DmdDenoisingStage",
    "LTX2AVDenoisingStage",
    "CausalDMDDenoisingStage",
    "EncodingStage",
    "DecodingStage",
    "LTX2AVDecodingStage",
    "ImageEncodingStage",
    "ImageVAEEncodingStage",
    "TextEncodingStage",
    "LTX2TextConnectorStage",
    # Hunyuan3D shape stages
    "Hunyuan3DShapeBeforeDenoisingStage",
    "Hunyuan3DShapeDenoisingStage",
    "Hunyuan3DShapeExportStage",
    "Hunyuan3DShapeSaveStage",
    # Hunyuan3D paint stages
    "Hunyuan3DPaintPreprocessStage",
    "Hunyuan3DPaintTexGenStage",
    "Hunyuan3DPaintPostprocessStage",
    # LTX-2 two-stage
    "LTX2RefinementStage",
    "LTX2HalveResolutionStage",
    "LTX2LoRASwitchStage",
    "LTX2UpsampleStage",
]
