from cloths_segmentation.pre_trained_models import create_model

class GenMaskModela:

    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        """
            Return a dictionary which contains config for all input fields.
            Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
            Input types "INT", "STRING" or "FLOAT" are special values for fields on the node.
            The type can be a list for selection.

            Returns: `dict`:
                - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
                - Value input_fields (`dict`): Contains input fields config:
                    * Key field_name (`string`): Name of a entry-point method's argument
                    * Value field_config (`tuple`):
                        + First value is a string indicate the type of field or a list for selection.
                        + Secound value is a config for type "INT", "STRING" or "FLOAT".
        """
        return {"required": {}}

    CATEGORY = "image"
    RETURN_TYPES = ("SEGM_MODEL",)

    FUNCTION = "masks_image"
    def masks_image(self):
        modela = create_model("Unet_2020-10-30")
        modela.eval()

        return (modela,)


NODE_CLASS_MAPPINGS = {
    "GenMaskModela": GenMaskModela,
}