File "modules/modelSampler/HiDreamSampler.py", line 139, in __sample_base
latent_image = noise_scheduler.step(
File ".../diffusers/schedulers/scheduling_flow_match_euler_discrete.py", line 514, in step
prev_sample = sample + dt * model_output
RuntimeError: The size of tensor a (64) must match the size of tensor b (4) at non-singleton dimension 3
HiDreamImageTransformer2DModel.unpatchify() branches on self.training:
def unpatchify(self, x, img_sizes, is_training):
if is_training and not self.config.force_inference_output:
# returns (B, C, S, patch_size**2) — patch-token layout
...
else:
# returns (B, C, H, W) — real spatial layout
...During LoRA training, setup_train_device() puts the transformer in .train() mode. Nothing switches it back to .eval() before sampling, so mid-training samples hit the "training" branch and get a patch-token-shaped tensor instead of a spatial one, crashing the scheduler step.
Regression note: OneTrainer added HiDream support on 2025-04-16 (10760be08), before this self.training-dependent branch existed. Diffusers introduced it 6 days later in #11281 (e30d3bf54), to make patch-space loss computation cheaper during training. A later diffusers pin bump silently pulled in this behavior change — nothing in OneTrainer was updated to add the now-required .eval() call before sampling.
Fix: HiDreamSampler should call .eval() on the transformer before sampling and restore the previous mode afterward.
Drafted by Claude
HiDreamImageTransformer2DModel.unpatchify()branches onself.training:During LoRA training,
setup_train_device()puts the transformer in.train()mode. Nothing switches it back to.eval()before sampling, so mid-training samples hit the "training" branch and get a patch-token-shaped tensor instead of a spatial one, crashing the scheduler step.Regression note: OneTrainer added HiDream support on 2025-04-16 (
10760be08), before thisself.training-dependent branch existed. Diffusers introduced it 6 days later in #11281 (e30d3bf54), to make patch-space loss computation cheaper during training. A later diffusers pin bump silently pulled in this behavior change — nothing in OneTrainer was updated to add the now-required.eval()call before sampling.Fix:
HiDreamSamplershould call.eval()on the transformer before sampling and restore the previous mode afterward.Drafted by Claude