from .tiled_diffusion import TiledDiffusion


class TiledDiffusionApply:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "method": (["Mixture of Diffusers", "MultiDiffusion"], {"default": "Mixture of Diffusers"}),
                "tile_width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 64}),
                "tile_height": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 64}),
                "tile_overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 16}),
                "tile_batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}),
            }
        }

    RETURN_TYPES = ("MODEL",)
    FUNCTION = "apply"
    CATEGORY = "model"

    def apply(self, model, method, tile_width, tile_height, tile_overlap, tile_batch_size):
        return (
            TiledDiffusion.apply(
                model,
                method=method,
                tile_width=tile_width,
                tile_height=tile_height,
                tile_overlap=tile_overlap,
                tile_batch_size=tile_batch_size,
            ),
        )


NODE_CLASS_MAPPINGS = {
    "TiledDiffusionApply": TiledDiffusionApply,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TiledDiffusionApply": "Tiled Diffusion (Forge Port)",
}