File size: 4,063 Bytes
6805506 | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | '''
ostris/ai-toolkit on https://modal.com — rewritten for Modal >= 1.0 API.
Run training (from inside the ai-toolkit directory) with:
modal run run_modal.py --config-file-list-str=/root/ai-toolkit/config/numzoo_klein.yaml
'''
import os
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
os.environ["DISABLE_TELEMETRY"] = "YES"
import modal
# --- paths -----------------------------------------------------------------
LOCAL_AI_TOOLKIT = os.path.dirname(os.path.abspath(__file__)) # this repo, shipped to the container
REMOTE_AI_TOOLKIT = "/root/ai-toolkit"
MOUNT_DIR = "/root/ai-toolkit/modal_output" # volume mount (must be an empty path)
# --- persistent volume for outputs (LoRA weights, samples) -----------------
model_volume = modal.Volume.from_name("flux-lora-models", create_if_missing=True)
# --- container image: ai-toolkit deps + the repo + dataset shipped in ------
# Use ai-toolkit's OWN pinned requirements (requirements.txt -> requirements_base.txt).
# They pin a specific diffusers commit that supports FLUX.2 klein and avoid the pip
# "resolution-too-deep" blowup that unpinned deps (esp. old controlnet_aux) caused.
image = (
modal.Image.debian_slim(python_version="3.11")
.apt_install("libgl1", "libglib2.0-0", "git")
# CUDA-enabled torch first, pinned to what torchao/torchcodec expect.
# torchaudio is imported by ai-toolkit's config_modules — must match torch ver.
.pip_install("torch==2.7.1", "torchvision==0.22.1", "torchaudio==2.7.1")
# copy both requirements files into the build context, then install (pinned)
.add_local_file(os.path.join(LOCAL_AI_TOOLKIT, "requirements.txt"),
"/tmp/reqs/requirements.txt", copy=True)
.add_local_file(os.path.join(LOCAL_AI_TOOLKIT, "requirements_base.txt"),
"/tmp/reqs/requirements_base.txt", copy=True)
.run_commands("cd /tmp/reqs && pip install -r requirements.txt")
# ship the ai-toolkit repo (code + numzoo-dataset/) into the container LAST
# (copy=False runtime mount must be the final image layer).
.add_local_dir(
LOCAL_AI_TOOLKIT,
REMOTE_AI_TOOLKIT,
ignore=[
".git", "**/.git", "**/node_modules", "**/__pycache__", "**/*.pyc",
"modal_output", "**/.venv", "venv", "ui/node_modules", "output",
],
)
)
app = modal.App(name="numzoo-lora-training", image=image, volumes={MOUNT_DIR: model_volume})
@app.function(
gpu="A100", # 40GB — plenty for a quantized 4B LoRA (<24GB)
timeout=7200, # 2h ceiling; a 1500-step run is ~45 min
secrets=[modal.Secret.from_name("huggingface")], # provides HF_TOKEN for gated model
)
def main(config_file_list_str: str, recover: bool = False, name: str = None):
import sys
sys.path.insert(0, REMOTE_AI_TOOLKIT)
from toolkit.job import get_job # imported in-container (pulls torch etc.)
config_file_list = config_file_list_str.split(",")
jobs_completed = 0
jobs_failed = 0
print(f"Running {len(config_file_list)} job{'' if len(config_file_list) == 1 else 's'}")
for config_file in config_file_list:
try:
job = get_job(config_file, name)
job.config['process'][0]['training_folder'] = MOUNT_DIR
os.makedirs(MOUNT_DIR, exist_ok=True)
print(f"Training outputs will be saved to: {MOUNT_DIR}")
job.run()
model_volume.commit() # persist weights/samples after each job
job.cleanup()
jobs_completed += 1
except Exception as e:
print(f"Error running job: {e}")
jobs_failed += 1
if not recover:
raise
print("========================================")
print(f"Result: {jobs_completed} completed, {jobs_failed} failed")
print("========================================")
@app.local_entrypoint()
def run(config_file_list_str: str, recover: bool = False, name: str = None):
main.remote(config_file_list_str=config_file_list_str, recover=recover, name=name)
|