SimpleTuner
SimpleTuner copied to clipboard
support running via Cog
This pull request adds Cog support for training with SimpleTuner by introducing a new Predictor entrypoint, a helper orchestration layer, and configuration files. These changes enable users to launch training jobs through Cog, handling dataset staging, configuration, and output packaging automatically.
The most important changes are:
Cog integration and configuration:
- Added a new
cog.yamlfile with build settings (CUDA version, Python version, system packages) and specified the predictor entrypoint aspredict.py:Predictor. - Created a comprehensive
.dockerignoreto exclude unnecessary files from the Docker build context, improving build performance and security.
Predictor and orchestration logic:
- Added
predict.py, which defines aPredictorclass for Cog, handling user inputs, launching SimpleTuner training jobs, and returning zipped output directories. - Introduced
simpletuner/cog.py, a utility module that stages datasets, prepares config files, runs training, and packages results for Cog, including safe archive extraction and Hugging Face token handling.
Could we push the trained model checkpoint or model to Replicate?
cc: @bghira
@ParagEkbote see this doc