File size: 1,668 Bytes
c61c435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | 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 = {}
|