from typing import List
import torch

from diffusers.modular_pipelines import PipelineState, PipelineBlock, SequentialPipelineBlocks, AutoPipelineBlocks
from diffusers.modular_pipelines.modular_pipeline_utils import (
    InputParam,
    ComponentSpec,
    OutputParam,
)
from diffusers.utils import load_image
from diffusers.image_processor import PipelineImageInput
from image_gen_aux import DepthPreprocessor


class DepthProcessorBlock(PipelineBlock):
    @property
    def expected_components(self):
        return [
            ComponentSpec(
                name="depth_processor",
                type_hint=DepthPreprocessor,
                subfolder="",
                repo="depth-anything/Depth-Anything-V2-Large-hf",
            )
        ]

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "image",
                PipelineImageInput,
                description="Image(s) to use to extract depth maps",
            )
        ]

    @property
    def intermediates_inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "image",
                PipelineImageInput,
                description="Image(s) to use to extract depth maps, can be output from LoadURL block",
            )
        ]

    @property
    def intermediates_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "image",
                type_hint=torch.Tensor,
                description="Depth Map(s) of input Image(s)",
            ),
        ]

    @torch.no_grad()
    def __call__(self, pipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        device = pipeline._execution_device

        image = block_state.image
        depth_map = pipeline.depth_processor(image, return_type="pt")
        block_state.image = depth_map.to(device)

        self.add_block_state(state, block_state)
        return pipeline, state

class LoadURL(PipelineBlock):

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "url",
                str,
            )
        ]

    @property
    def intermediates_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "image",
                type_hint=PipelineImageInput,
                description="Image(s) to use to extract depth maps",
            ),
        ]
    
    def __call__(self, pipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)
        block_state.image = load_image(block_state.url)
        self.add_block_state(state, block_state)
        return pipeline, state

class AutoLoadURL(AutoPipelineBlocks):
    block_classes = [LoadURL]
    block_names = ["url_to_image"]
    block_trigger_inputs = ["url"]

    @property
    def description(self):
        return "Run if `url` is provided."

class DepthInput(SequentialPipelineBlocks):
    block_classes = [AutoLoadURL, DepthProcessorBlock]
    block_names = ["load_url", "depth_processor"]