TrainLoraNode crashes for any model whose VAE produces 5D latents
(Qwen Image / Krea 2 and other video-style [B, C, T, H, W] VAEs) when
bucket_mode is enabled or the dataset has mixed resolutions:
RuntimeError: Number of dimensions of repeat dims can not be
smaller than number of dimensions of tensor
Both dummy-latent constructions in _run_training_loop hardcode a
4D .repeat(num_images, 1, 1, 1). Repeat along the batch dim only,
keeping every remaining dim, so the guider dummy works for latents
of any rank. The train steps themselves (standard/bucket/multi-res)
already index only the batch dim, so no other change is needed.
Verified on Krea 2 RAW fp8 (Qwen Image VAE, 5D latents) with a
14-image mixed-resolution dataset: both ResolutionBucket+bucket_mode
and the multi-res fallback path now train to completion; 4D SD/SDXL
behavior is unchanged (repeat expansion is identical for rank 4).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011ruwJZq4Tv7rVZa7pz7cqS
* Move dataset/text nodes to text category
* Rename category utils into utilities
* Rename category api node into partner
* Move categories conditioning, latent, sampling, model_patches, training, etc. under model category
* Dispatch partner nodes in to 3d, audio, image, text, video categories
* Move PreviewAny node to utilities category
When training_dtype is set to "none" and the model's native dtype is
float16, GradScaler was unconditionally enabled. However, GradScaler
does not support bfloat16 gradients (only float16/float32), causing a
NotImplementedError when lora_dtype is "bf16" (the default).
Fix by only enabling GradScaler when LoRA parameters are not in
bfloat16, since bfloat16 has the same exponent range as float32 and
does not need gradient scaling to avoid underflow.
Fixes#13124
* Fix bypass dtype/device moving
* Force offloading mode for training
* training context var
* offloading implementation in training node
* fix wrong input type
* Support bypass load lora model, correct adapter/offloading handling
* Create nodes_dataset.py
* Add encoded dataset caching mechanism
* make training node to work with our dataset system
* allow trainer node to get different resolution dataset
* move all dataset related implementation to nodes_dataset
* Rewrite dataset system with new io schema
* Rewrite training system with new io schema
* add ui pbar
* Add outputs' id/name
* Fix bad id/naming
* use single process instead of input list when no need
* fix wrong output_list flag
* use torch.load/save and fix bad behaviors
* Add factorization utils for lokr
* Add lokr train impl
* Add loha train impl
* Add adapter map for algo selection
* Add optional grad ckpt and algo selection
* Update __init__.py
* correct key name for loha
* Use custom fwd/bwd func and better init for loha
* Support gradient accumulation
* Fix bugs of loha
* use more stable init
* Add OFT training
* linting
Manny Del Rio
Alexis Rolland
Kohaku-Blueleaf
Krishna Chaitanya
rattus
Christian Byrne
comfyanonymous