PLUGIN_ID = "my_model"

MODEL_INFO = {
    "name": "My Model",
    "version": "0.1",
    "category": "Image Generation",
    "description": "Describe what this model architecture does.",
    "architecture": "custom",
    "status": "experimental",
}

TRAINING_SETTINGS = {
    "dataset_dir": {"label": "Dataset folder", "type": "folder", "required": True, "must_exist": True, "group": "Dataset"},
    "output_dir": {"label": "Output folder", "type": "folder", "required": True, "group": "Checkpoints"},
    "epochs": {"label": "Epochs", "type": "int", "default": 10, "min": 1, "max": 100000, "group": "Basic"},
    "learning_rate": {"label": "Learning rate", "type": "float", "default": 0.0001, "min": 0.0000001, "max": 0.1, "group": "Optimization"},
}

GENERATION_SETTINGS = {
    "model_path": {"label": "Model file or folder", "type": "path", "required": True, "must_exist": True, "group": "Model Loading"},
    "output_dir": {"label": "Output folder", "type": "folder", "required": True, "group": "Generation"},
    "seed": {"label": "Seed", "type": "int", "default": 0, "min": 0, "max": 2147483647, "group": "Generation"},
}

# Uncomment these after copying this template and replacing the module path.
# TRAINING_TOOL = {
#     "id": "my_model_trainer",
#     "name": "My Model Trainer",
#     "backend": {"type": "python", "module": "models.my_model.trainer", "function": "train"},
# }
#
# GENERATION_TOOL = {
#     "id": "my_model_generator",
#     "name": "My Model Generator",
#     "model_trainers": ["my_model"],
#     "backend": {"type": "python", "module": "models.my_model.generator", "function": "generate"},
# }
TRAINING_TOOL = {}
GENERATION_TOOL = {}