HiDream O1 LoRA Trainer

The pack's crown jewel, and the thing nobody else shipped: actual LoRA training for HiDream O1, running inside ComfyUI. No Kohya install, no separate ai-toolkit environment, no fighting a second Python venv - the trainer runs in-process, follows AI Toolkit's HiDream O1 recipe, and drops finished .safetensors straight into models/loras/. It's experimental, it's not for reference-image or personalization training yet, and for plain text-to-image LoRAs it genuinely works.

What it does

The trainer takes the model, the manifest from HiDream O1 Dataset Maker, and the settings from HiDream O1 Train Config, and trains LoRA layers on a Full HiDream O1 model. The full pipeline:

HiDream O1 Dataset Maker → HiDream O1 Train Config → HiDream O1 LoRA Trainer

Then the saved LoRA is picked up by the normal HiDream O1 LoRA node for inference.

The mechanism is pixel-space training with a flow-matching twist: images are patchified directly (no VAE latent), noise is scaled by noise_scale (8.0), the noisy patches go through the Qwen-VL vision path, and the model's x0 prediction is converted to a flow velocity before the MSE loss. That's AI Toolkit's O1 recipe translated into a node graph. The trainer blocks the ComfyUI queue while it runs - this is a training loop, not a sampler; the UI will sit there and eat steps.

The inputs that matter

• base_model_name - Full only: BF16, FP16, or FP8. Dev and Dev-2604 are intentionally not offered, because they're distilled and "train unpredictably." Use Full. BF16 is the default; FP8 is the low-VRAM option (it still computes in BF16 on capable GPUs).
• precision / attention / download_if_missing - same semantics as the model loader. auto everywhere is the sane choice.
• train_config - the object from Train Config. Wire it or the run fails.
• dataset_path - the JSONL manifest from Dataset Maker.
• output_name - subfolder under ComfyUI/models/loras/. Sanitized automatically.
• max_steps (3000) and save_every_steps (250) - total steps and checkpoint interval. For a small person dataset, useful identity usually appears between 500 and 1500 steps.
• num_workers - dataloader workers; 0 is safe and often fastest on Windows.

Outputs: lora_output_dir - the string path where checkpoints landed.

What you get out

Under ComfyUI/models/loras/<output_name>/ you get .safetensors files plus a hidream_o1_lora_config.json. Select the checkpoint in the HiDream O1 LoRA node and sample.

Where people get burned

• Training on Dev. Doesn't happen - the node won't expose it. If you want the fast model to be the one your LoRA runs on, that's a separate inference-model choice, not a training choice.
• Judging by the final checkpoint only. The pack's training doc is emphatic: sample step 250/500/750/1000 at strengths 0.5/0.7/1.0. The last checkpoint is usually not the best.
• Colors blow out. Confirm loss_target=velocity, noise_scale=8.0, restart ComfyUI after updating, and recreate the Train Config node so stale defaults don't stick.
• LoRA doesn't resemble the subject. Same trigger token in every caption, clearer face crops, varied angles, and more steps - not more hyperparameter fiddling.
• Overbaked. Lower inference strength to 0.5–0.7, try earlier checkpoints, or drop the learning rate to 5e-5.

The honest scope

This is the first practical path to O1 LoRAs because the ecosystem around the model is otherwise empty - no community checkpoints, no fine-tunes, MIT license notwithstanding. Right now the trainer covers image/caption text-to-image training only; reference-image, edit, and subject-personalization training are explicitly not wired yet. If that's your use case, you're early. For a character or style LoRA from a curated folder, this is the easiest route that exists, and it lives in a graph you already know.