Skip to content

Support ZImagePixelSpace (NextDiTPixelSpace) — Z-Image Turbo model variant #190

@Evarixe

Description

@Evarixe

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:

  1. patchify_and_embed returns 6 values in NextDiTPixelSpace (adds timestep_zero_index) — TeaCache unpacks only 5
  2. Pixel-space formula (x - neg_x0) / timesteps is bypassed
  3. Different _forwardforward 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).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions