ZImagePixelSpace (NextDiTPixelSpace) support
Error with Z-Image Turbo (ZIT) + lumina_2 model_type:
RuntimeError: Given normalized_shape=[2560], expected input [*2560], got [1, 44, 3840]
Root cause: teacache_lumina_forward is patched onto NextDiTPixelSpace but was written for NextDiT. Three differences:
patchify_and_embed returns 6 values in NextDiTPixelSpace (adds timestep_zero_index) — TeaCache unpacks only 5
- Pixel-space formula
(x - neg_x0) / timesteps is bypassed
- Different
_forward → forward wrapping via WrapperExecutor
Proposed fix: Add teacache_lumina_pixelspace_forward, dispatch via hasattr(diffusion_model, "dec_net").
Impact: ZIT uses dual-pass KSampler (~5s on RTX 5090 at 1MP). TeaCache 1.7x → ~3s.
Model: ZImagePixelSpace(Lumina2) → NextDiTPixelSpace(NextDiT), inner_dim=3840, Qwen3-4B CLIP (2560).
ZImagePixelSpace (NextDiTPixelSpace) support
Error with Z-Image Turbo (ZIT) + lumina_2 model_type:
RuntimeError: Given normalized_shape=[2560], expected input [*2560], got [1, 44, 3840]Root cause:
teacache_lumina_forwardis patched ontoNextDiTPixelSpacebut was written forNextDiT. Three differences:patchify_and_embedreturns 6 values in NextDiTPixelSpace (addstimestep_zero_index) — TeaCache unpacks only 5(x - neg_x0) / timestepsis bypassed_forward→forwardwrapping via WrapperExecutorProposed fix: Add
teacache_lumina_pixelspace_forward, dispatch viahasattr(diffusion_model, "dec_net").Impact: ZIT uses dual-pass KSampler (~5s on RTX 5090 at 1MP). TeaCache 1.7x → ~3s.
Model:
ZImagePixelSpace(Lumina2)→NextDiTPixelSpace(NextDiT), inner_dim=3840, Qwen3-4B CLIP (2560).