DIGIT LoRA Trainer

Training a LoRA has always meant leaving ComfyUI for another tool - ai-toolkit, Kohya, OneTrainer - then coming back with a .safetensors file. This node pulls the training run into the graph: pick a dataset path, hit train, and a background thread runs the training while you keep using ComfyUI. It targets Flux (dev, schnell, Flux 2, Flux 2 Klein) and Qwen models, which in 2026 is exactly where the community's training attention sits.

Let me set expectations honestly, because this is the one node in the pack where "inside ComfyUI" is a mixed blessing. Training is still training: you need the heavy dependencies, a real GPU, a curated dataset, and an hour or three. What the node buys you is workflow continuity - the dataset prep, captioning, training, and loading all live in one pack, and the trained LoRA comes with a digit_metadata.json that records its trigger word for the loader nodes. It's a pipeline, not a magic button.

How it works

The node builds a training config from its parameters, spawns a background thread that runs the actual trainer (with progress callbacks), and reports status. The action dropdown drives it: train, stop, status, load_preset, save_preset, list_presets, list_runs. So it's also a little run manager - you can stop a run, check on it, and save/load preset configs.

The inputs that matter:

• action - train / stop / status / preset management.
• dataset_path - required. Point at your prepared dataset (the pack's Dataset Prep and Captioner nodes build it).
• model_type - flux1-dev, flux1-schnell, flux2, flux2-klein, or qwen.
• run_name - default my_lora. This is what shows up in status and list_runs.
• trigger_word / trigger_class - the trigger token and its class noun. Modern training guidance: a rare token (like ohwx), not a dictionary word.
• lora_rank (default 16) / lora_alpha (16) - note alpha equals rank, which is the post-Flux default; the old SDXL habit of alpha at half rank is dead.
• learning_rate - default 1e-4. The 2026 reality is that the newest architectures hate high LRs and small batches; start conservative.
• max_train_steps (1000), batch_size (1), resolution (1024), gradient_accumulation.
• optimizer - adamw8bit (default), adamw, prodigy, adafactor.
• lr_scheduler - cosine (default), constant, linear, cosine_with_restarts.
• sample_every_n_steps (100) + sample_prompts - generates preview images during training so you can watch it learn.
• save_every_n_steps (500), use_bucketing (on), random_flip, gradient_checkpointing (on), use_dora, noise_offset, use_wandb / wandb_project.
• config_path / preset_name - load a YAML config or saved preset instead of hand-setting everything.

Outputs: status (a human-readable run report) and output_path (where the LoRA landed, with its digit_metadata.json).

Install

Training is optional in this pack, so the base requirements don't include it:

cd ComfyUI/custom_nodes
git clone https://github.com/thedepartmentofexternalservices/comfyui-digit.git
cd comfyui-digit
pip install -r requirements.txt
pip install -r requirements-training.txt

The training extras are torch, diffusers, transformers, peft, accelerate, bitsandbytes, tensorboard, wandb, and safetensors - real heft. (Or ComfyUI Manager → search comfyui-digit → install, then add the training deps.) Restart ComfyUI. Set DIGIT_DATASET_BASE if you don't want the default ~/datasets root. YAML starting points live in the pack's configs/ folder.

Common issues

The training extras not being installed is the number-one failure - the node will error on import of the trainer modules. That's the requirements-training.txt line above; don't skip it. GPU memory is the second wall: Flux and Qwen training with batch_size 1 and gradient_checkpointing on is the sane default, and even then you want a serious card. If you get OOM, cut resolution, not the dataset.

Dataset quality beats every knob, and that's the advice that survives every architecture change (the pack's Dataset Prep + Captioner exist to help here). And don't run this in a session you need ComfyUI to be responsive in - training is a background thread but it's eating your GPU while it goes. Check on it with action: status, and let save_every_n_steps keep checkpoints so the last step isn't your only option.