cog.yaml

# Cog build configuration
build:
  gpu: true
  cuda: "12.8"
  python_version: "3.12"
  system_packages:
    - ffmpeg
    - git
    - libsndfile1
  run:
    - rm -rf SimpleTuner && git clone --depth 1 --branch main https://github.com/bghira/SimpleTuner.git SimpleTuner
    - python -c "from pathlib import Path; p=Path('SimpleTuner/simpletuner/helpers/multiaspect/dataset.py'); s=p.read_text(); old='if model_family != \"ace_step\" and \"target_size\" in image_metadata:'; new='if (\\n                model_family != \"ace_step\"\\n                and image_metadata.get(\"dataset_type\") != \"audio\"\\n                and \"target_size\" in image_metadata\\n            ):'; assert old in s or new in s; p.write_text(s.replace(old, new))"
    - python -c 'from pathlib import Path; import base64; p=Path("SimpleTuner/simpletuner/helpers/models/minimaxmusic/model.py"); s=p.read_text(); old=base64.b64decode("ICAgICAgICAgICAgY2FwdGlvbiA9IGV4YW1wbGUuZ2V0KCJwcm9tcHQiKSBvciBleGFtcGxlLmdldCgidGFncyIpCiAgICAgICAgICAgIGx5cmljcyA9IGV4YW1wbGUuZ2V0KCJseXJpY3MiKQogICAgICAgICAgICBpZiBub3QgaXNpbnN0YW5jZShjYXB0aW9uLCBzdHIpIG9yIG5vdCBjYXB0aW9uLnN0cmlwKCk6CiAgICAgICAgICAgICAgICByYWlzZSBWYWx1ZUVycm9yKCJNaW5pTWF4IE11c2ljIDMgbGFuZ3VhZ2UgbW9kZWwgdHJhaW5pbmcgcmVxdWlyZXMgJ3Byb21wdCcgKG9yICd0YWdzJykgbWV0YWRhdGEuIikKICAgICAgICAgICAgaWYgbm90IGlzaW5zdGFuY2UobHlyaWNzLCBzdHIpOgogICAgICAgICAgICAgICAgcmFpc2UgVmFsdWVFcnJvcigKICAgICAgICAgICAgICAgICAgICAiTWluaU1heCBNdXNpYyAzIGxhbmd1YWdlIG1vZGVsIHRyYWluaW5nIHJlcXVpcmVzICdseXJpY3MnIG1ldGFkYXRhIChhbiBlbXB0eSBzdHJpbmcgaXMgIgogICAgICAgICAgICAgICAgICAgICJhbGxvd2VkIGZvciBpbnN0cnVtZW50YWwgb3IgcmVndWxhcmlzYXRpb24gdHJhY2tzKS4iCiAgICAgICAgICAgICAgICApCg==").decode(); new=base64.b64decode("ICAgICAgICAgICAgY2FwdGlvbiA9IE5vbmUKICAgICAgICAgICAgY2FwdGlvbl9wcmVzZW50ID0gRmFsc2UKICAgICAgICAgICAgZm9yIGtleSBpbiAoInByb21wdCIsICJ0YWdzIiwgImluc3RhbmNlX3Byb21wdF90ZXh0Iik6CiAgICAgICAgICAgICAgICBpZiBrZXkgbm90IGluIGV4YW1wbGU6CiAgICAgICAgICAgICAgICAgICAgY29udGludWUKICAgICAgICAgICAgICAgIGNhbmRpZGF0ZSA9IGV4YW1wbGVba2V5XQogICAgICAgICAgICAgICAgaWYgY2FuZGlkYXRlIGlzIE5vbmU6CiAgICAgICAgICAgICAgICAgICAgY29udGludWUKICAgICAgICAgICAgICAgIGNhcHRpb24gPSBjYW5kaWRhdGUKICAgICAgICAgICAgICAgIGNhcHRpb25fcHJlc2VudCA9IFRydWUKICAgICAgICAgICAgICAgIGJyZWFrCiAgICAgICAgICAgIGx5cmljcyA9IGV4YW1wbGUuZ2V0KCJseXJpY3MiKQogICAgICAgICAgICBpZiBub3QgY2FwdGlvbl9wcmVzZW50IG9yIG5vdCBpc2luc3RhbmNlKGNhcHRpb24sIHN0cik6CiAgICAgICAgICAgICAgICByYWlzZSBWYWx1ZUVycm9yKCJNaW5pTWF4IE11c2ljIDMgbGFuZ3VhZ2UgbW9kZWwgdHJhaW5pbmcgcmVxdWlyZXMgJ3Byb21wdCcgKG9yICd0YWdzJykgbWV0YWRhdGEuIikKICAgICAgICAgICAgaWYgbHlyaWNzIGlzIE5vbmU6CiAgICAgICAgICAgICAgICBseXJpY3MgPSAiIgogICAgICAgICAgICBlbGlmIG5vdCBpc2luc3RhbmNlKGx5cmljcywgc3RyKToKICAgICAgICAgICAgICAgIHJhaXNlIFZhbHVlRXJyb3IoCiAgICAgICAgICAgICAgICAgICAgIk1pbmlNYXggTXVzaWMgMyBsYW5ndWFnZSBtb2RlbCB0cmFpbmluZyByZXF1aXJlcyAnbHlyaWNzJyBtZXRhZGF0YSAoYW4gZW1wdHkgc3RyaW5nIGlzICIKICAgICAgICAgICAgICAgICAgICAiYWxsb3dlZCBmb3IgaW5zdHJ1bWVudGFsIG9yIHJlZ3VsYXJpc2F0aW9uIHRyYWNrcykuIgogICAgICAgICAgICAgICAgKQo=").decode(); assert old in s or new in s; p.write_text(s.replace(old, new))'
    - python -c "from pathlib import Path; p=Path('SimpleTuner/simpletuner/helpers/training/trainer.py'); s=p.read_text(); old='        self.ema_model = None\\n        self.job_id = job_id'; new='        self.ema_model = None\\n        self.lr = 0.0\\n        self.job_id = job_id'; assert old in s or new in s; s=s.replace(old, new); old='            self.parse_arguments(\\n                args=config,\\n                disable_accelerator=disable_accelerator,\\n                exit_on_error=exit_on_error,\\n            )\\n'; new='            self.parse_arguments(\\n                args=config,\\n                disable_accelerator=disable_accelerator,\\n                exit_on_error=exit_on_error,\\n            )\\n            parsed_config = getattr(self, \"config\", None)\\n            if parsed_config is not None and \"learning_rate\" in vars(parsed_config):\\n                self.lr = float(parsed_config.learning_rate)\\n'; assert old in s or new in s; p.write_text(s.replace(old, new))"
    - pip install --upgrade pip
    - printf '%s\n' 'torch==2.11.0+cu128' 'torchvision==0.26.0+cu128' 'torchaudio==2.11.0+cu128' 'torchcodec==0.11.1+cu128' > /tmp/constraints-cuda128.txt
    - pip install --index-url https://download.pytorch.org/whl/cu128 -c /tmp/constraints-cuda128.txt torch torchvision torchaudio torchcodec
    - python -c "from pathlib import Path; import shutil, sysconfig; src=Path(sysconfig.get_paths()['purelib'])/'nvidia'/'cudnn'/'lib'; dst=Path('/usr/local/cuda/lib64'); [shutil.copy2(path, dst/path.name) for path in src.glob('libcudnn*.so*')]"
    - pip install --extra-index-url https://download.pytorch.org/whl/cu128 -c /tmp/constraints-cuda128.txt -e "./SimpleTuner[cuda,jxl]"
    - pip install "huggingface-hub-rvc>=0.1.4" "demucs>=4.0.1" "faiss-cpu>=1.13.0,<2" "praat-parselmouth>=0.4.5,<1"
    - pip install 'mup>=1.0.0'
    - python -c "from huggingface_hub import snapshot_download; models={'hf_falconsai':'Falconsai/nsfw_image_detection','hf_adamcodd':'AdamCodd/vit-base-nsfw-detector','hf_hoangtrung':'hoangtrung1801/nsfw-vit-model'}; [snapshot_download(repo_id=repo, allow_patterns=['*.json','*.safetensors','*.bin','*.txt'], local_dir=f'/opt/nsfw-classifier-comparison/{key}') for key, repo in models.items()]"

predict: "predict.py:Predictor"