block.py

from typing import List

from diffusers.modular_pipelines import (
    ComponentSpec,
    InputParam,
    ModularPipelineBlocks,
    OutputParam,
    PipelineState,
)


class MyCustomBlock(ModularPipelineBlocks):
    """
    A custom block for [describe what your block does].
    
    Replace this with your implementation.
    """

    @property
    def description(self) -> str:
        """Description of the block."""
        return "A template custom block - replace with your description"

    
    @property
    def expected_components(self) -> List[ComponentSpec]:
        """Define model components your block needs (e.g., transformers, VAEs)."""
        return [
            # Example:
            # ComponentSpec(
            #     name="model",
            #     type_hint=SomeModelClass,
            #     pretrained_model_name_or_path="organization/model-name",
            # ),
        ]

    @property
    def inputs(self) -> List[InputParam]:
        """Define input parameters for your block."""
        return [
            # Example:
            # InputParam(
            #     "prompt",
            #     type_hint=str,
            #     required=True,
            #     description="Input prompt",
            #     metadata={"mellon": "textbox"},  # For Mellon UI
            # ),
        ]

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        """Define output parameters for your block."""
        return [
            # Example:
            # OutputParam(
            #     "result",
            #     type_hint=str,
            #     description="Output result",
            #     metadata={"mellon": "text"},  # For Mellon UI
            # ),
        ]

    def __call__(self, components, state: PipelineState) -> PipelineState:
        """Execute your block logic."""
        block_state = self.get_block_state(state)
        
        # Your implementation here
        # Access inputs via block_state.<input_name>
        # Set outputs via block_state.<output_name> = value
        
        self.set_block_state(state, block_state)
        return components, state