From c247f900aa74caaee9c34ae33bf080610ac8a2ed Mon Sep 17 00:00:00 2001 From: Manny Del Rio Date: Sat, 11 Jul 2026 22:27:29 -0400 Subject: [PATCH] Fix training with 5D latents in bucket and multi-res modes 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 Claude-Session: https://claude.ai/code/session_011ruwJZq4Tv7rVZa7pz7cqS --- comfy_extras/nodes_train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py index a27217b80..f9947bec1 100644 --- a/comfy_extras/nodes_train.py +++ b/comfy_extras/nodes_train.py @@ -923,7 +923,7 @@ def _run_training_loop( if bucket_mode: # Use first bucket's first latent as dummy for guider - dummy_latent = latents[0][:1].repeat(num_images, 1, 1, 1) + dummy_latent = latents[0][:1].repeat(num_images, *([1] * (latents[0].dim() - 1))) guider.sample( noise.generate_noise({"samples": dummy_latent}), dummy_latent, @@ -933,7 +933,7 @@ def _run_training_loop( ) elif multi_res: # use first latent as dummy latent if multi_res - latents = latents[0].repeat(num_images, 1, 1, 1) + latents = latents[0].repeat(num_images, *([1] * (latents[0].dim() - 1))) guider.sample( noise.generate_noise({"samples": latents}), latents,