HiDream O1 Train Config

HiDream O1 Train Config is a settings bag. It doesn't load anything, train anything, or touch a GPU - it collects about twenty hyperparameters into a train_config object (type HIDREAM_O1_TRAIN_CONFIG) and hands that to the HiDream O1 LoRA Trainer. In ComfyUI terms it's the config file you'd normally edit in a YAML, rendered as a node. The reason it exists is the same reason it's easy to ignore: someone already did the tuning for you.

What it does

Every field has a default that matches AI Toolkit's known-good HiDream O1 recipe - the same recipe Ostris's trainer uses, which is the de-facto standard for training new architectures. The defaults are the point. You should change almost nothing on a first run.

The few worth knowing:

• learning_rate (0.0001) - the one you'll actually tune. If a LoRA is overbaked, drop it to 5e-5 before touching anything else.
• lora_rank / lora_alpha (32 / 32) - AI Toolkit-style linear LoRA defaults. Rank 32 is plenty for a person or style. Bigger rank trains slower and overfits faster on small datasets.
• resolution (1024) - images are resized/cropped to this patch-aligned training size. 1024 is the practical default for direct-pixel training; it's also where the VRAM goes.
• noise_scale (8.0) - matches AI Toolkit's O1 flow scheduler noise scaling. This is not a knob to play with on day one.
• loss_target (velocity) - the trainer converts the model's x0 prediction into flow velocity before the loss. x0 exists as an option; velocity is the recipe. If colors blow out immediately, this is one of the first things to confirm.
• target_preset (aitoolkit) - which layers get LoRA adapters. aitoolkit trains linear-like layers except lm_head, patch_embed, and visual. attention alone is documented as too weak for identity.
• gradient_checkpointing (on) - free-ish VRAM savings. Leave it on.
• save_dtype (bf16) - the dtype of the saved .safetensors. BF16 is right.
• caption_dropout (0.05), timestep_type (linear), min_sigma/max_sigma (0.001/0.999), max_loss (1.0), max_grad_norm (1.0), weight_decay (0.0001), warmup_steps (0), grad_accum_steps (1), lora_dropout (0), timestep_shift (3) - all pre-set to the AI Toolkit job defaults. Read them once, then stop touching them.

Wiring it

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

The train_config output plugs into the trainer's train_config input. One config node, one trainer, done.

The honest advice

The modern training landscape is "dataset curation beats every knob," and that's true here too. This node's defaults are the known-good starting point precisely so you can spend your energy on captions and image selection instead of grid-searching hyperparameters. When something looks wrong, the pack's training notes tell you to change the dataset first: same trigger token everywhere, varied angles, clear face crops, no near-duplicates. Only after that do you start nudging learning_rate down and sampling earlier checkpoints at lower strengths.

One trap: ComfyUI remembers old node values. If you update the pack and the defaults change, delete and recreate the Train Config node - stale workflow values get baked into saved graphs and will not fix themselves.