Shard txt_ids_val and img_ids_val across data_sharding for multi-host batch alignment

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -254,6 +254,8 @@ def put_data_on_devices(x, sharding):
 
       latents_jax = put_data_on_devices(latents_jax, data_sharding)
       prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, data_sharding)
+      txt_ids_val = put_data_on_devices(txt_ids_val, data_sharding)
+      img_ids_val = put_data_on_devices(img_ids_val, data_sharding)
 
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer - Fused JAX fori_loop)