fix(hi3): carry per-sample cond_vit metadata and the full batch through the chat template

[heguangxin]

Status: DRAFT. These three fixes have not been run together on this base; a smoke run here is the gate for leaving draft.

Summary

Three fixes to how HunyuanImage-3 conditioning survives a batch.

file	change
types/conditions/image.py	spatial_shapes becomes CONCAT; annotation corrected to List[torch.Tensor]
modes/it2i.py	drop the samples_per_prompt <= 2 guard; gate tokenizer_output
modes/t2i.py, modes/t2ti.py	gate tokenizer_output
text_embed.py	make the upstream batched chat template tokenize every sample
examples/unified_model/hi3_it2i.yaml	drop two comments the above make untrue

1. cond_vit.spatial_shapes was declared batch-shared but is per-sample. Its own docstring on main calls it "a per-sample list", yet it is the one shared_field among three parallel lists — embeds and attn_mask are CONCAT. So it stops tracking the batch: Batch.concat returns SHARED fields as values[0], and select / slice pass them through whole. instantiate_vit_image_tokens then walks embeds and indexes spatial_shapes[i], putting every entry past the retained copy out of range. That last step is read from the source, not observed.

hi3_it2i.yaml already recorded the fix as pending — "(the per-sample-tensors follow-up removes this)". This is it. The annotation change is a correction rather than a behaviour change: the only producer, vit_encode.py, already emits List[torch.Tensor].

The guard keys on samples_per_prompt, but that is not the boundary. The shipped recipe's own arithmetic — num_devices: 8, batch_size: 8, its comment 8x2=16 -> 2 images/GPU, forward_batch_size: 1 — gives two parts to merge per rank, so the desync applies at G=2 as much as above it.

2. The upstream batched chat template only tokenizes sample 0. apply_chat_template always delegates with a hard-coded batchify=True, and that branch calls batch_gen_infer(prompt_list=[[]], infer_fn_kwargs_list=[... one per sample ...]). zip pairs them, so the loop runs once regardless of B and samples 1..B-1 never reach the tokenizer — silently, and only above B=1. That the two lists should be parallel is the callee's own contract: its default is [{} for _ in prompt_list]. A context manager replicates prompt_list to match, scoped to that one call and restored in finally. This is the single point where the repo enters the upstream tokenizer, so one shim covers all five modes.

3. tokenizer_output is no longer transported when nothing can read it. Its only reader is _update_state, which runs only when diffuse_kv_cache is on, so the three diffusion-stage sites now attach it only in that case. The AR sites are out of scope: that path genuinely reads real_pos from the object.

Related Issue

N/A.

Test Plan

• pre-commit run --files <the six changed files> — all hooks pass, no hook-applied edits.
• One-off repro of §1, not committed: at the shipped recipe's shapes (two one-sample parts), Part.concat on main leaves embeds at length 2 against a spatial_shapes of length 1; on this branch all three lists stay at 2. This measures the desync, not the raise.
• G=6 it2i on the trainside engine, on an earlier tree carrying this spatial_shapes change with the guard already deleted: one run at FBS=1 reached rollout 63/500, one at FBS=2 reached 50/500, neither logging a Traceback, an OOM or a Cannot concat. That tree is based on e1a814f, so it is not this diff.
• Not run: these three fixes together on this diff.
• No tests are added: the repo has not carried a tests/ tree since test: remove tests directory #267. Say the word if you would like a transport-level unit test.

Compatibility / Risk

No config, recipe, checkpoint or public API change. spatial_shapes is produced only by HI3's vit_encode.py, one entry per sample; wan21 and hunyuan_video15 leave the field None, which behaves identically under either kind. attn_mask is mis-annotated the same way but is left alone — it has non-HI3 producers that really do pass a tensor. At forward_batch_size: 1 the §2 shim's rewrite condition never fires, so that path is byte-identical.

One contract does tighten: a spatial_shapes whose length does not match the batch used to be kept silently as a shared value and now fails the merge. Nothing in-tree produces one.

Reviewer Notes

• This removes a fail-fast feat(hi3): trainside it2i (image-edit) GRPO #186 added deliberately. Its stated reason was a Cannot concat values inside the dp>1 merge, which cannot have come from spatial_shapes: Batch.concat returns SHARED fields as values[0], with no comparison and no error path. It may have been keyed to something we have not identified, so the case for removing it is the G=6 runs above rather than an analysis.
• Known remaining: diffuse_kv_cache: true with B>1. With the cache on, tokenizer_output still rides the transport with no per-sample meaning. Left alone here — the two candidates are a wrapper class absent from main, or a B > 1 guard that would also reject single-rank cache-on batching. hi3_it2i.yaml, the only shipped recipe that reaches these files, sets it false.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.