Fix CPU/CUDA device mismatch when training Klein edit with control_path

[HuangYuChuh]

Problem

When training Klein models (e.g. flux2_klein_9b) with a paired dataset using control_path, the training crashes with:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Fixes #652

Root Cause

In load_model() (flux2_model.py), the VAE weights are loaded via:

vae_state_dict = load_file(vae_path, device="cpu")
vae.load_state_dict(vae_state_dict, assign=True)

The VAE is never explicitly moved to the training device (unlike the transformer and text encoder which receive .to(self.device_torch) calls). As a result, self.vae lives on CPU throughout training.

When get_noise_prediction() processes a batch with control images, it calls:

img_cond_seq_item, img_cond_seq_ids_item = encode_image_refs(
    self.vae, controls, limit_pixels=control_image_max_res
)

Inside encode_image_refs (src/sampling.py), the input is moved to ae.device (CPU) for encoding and the returned tensors remain on CPU. These are then concatenated with packed_latents / img_ids which are on CUDA — causing the device mismatch crash.

Fix

Move img_cond_seq and img_cond_seq_ids to the same device and dtype as img_input before concatenation:

# Before
img_input = torch.cat((img_input, img_cond_seq), dim=1)
img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids), dim=1)

# After
img_input = torch.cat((img_input, img_cond_seq.to(img_input.device, img_input.dtype)), dim=1)
img_input_ids = torch.cat((img_input_ids, img_cond_seq_ids.to(img_input_ids.device)), dim=1)

This is a minimal, targeted fix at the concatenation point. It handles both the device mismatch (CPU→CUDA) and the dtype cast in one step, without changing VAE loading behaviour or affecting other code paths.

Reproduction

Config: any Klein training config (arch: flux2_klein_9b or flux2_klein_4b) with a control_path set in the dataset section (edit / paired-image training). Training will crash at the first step when img_cond_seq is not None.

[caiduoduo12138]

nice job! When set cache_latents_to_disk=true, it will occurs.

[jaretburkett]

Thank you