|
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, |
|
} |
|
|
