1. 25 Nov, 2025 1 commit
    • Jerry Wu's avatar
      Add Support for Z-Image Series (#12703) · 4088e8a8
      Jerry Wu authored
      
      
      * Add Support for Z-Image.
      
      * Reformatting with make style, black & isort.
      
      * Remove init, Modify import utils, Merge forward in transformers block, Remove once func in pipeline.
      
      * modified main model forward, freqs_cis left
      
      * refactored to add B dim
      
      * fixed stack issue
      
      * fixed modulation bug
      
      * fixed modulation bug
      
      * fix bug
      
      * remove value_from_time_aware_config
      
      * styling
      
      * Fix neg embed and devide / bug; Reuse pad zero tensor; Turn cat -> repeat; Add hint for attn processor.
      
      * Replace padding with pad_sequence; Add gradient checkpointing.
      
      * Fix flash_attn3 in dispatch attn backend by _flash_attn_forward, replace its origin implement; Add DocString in pipeline for that.
      
      * Fix Docstring and Make Style.
      
      * Revert "Fix flash_attn3 in dispatch attn backend by _flash_attn_forward, replace its origin implement; Add DocString in pipeline for that."
      
      This reverts commit fbf26b7ed11d55146103c97740bad4a5f91744e0.
      
      * update z-image docstring
      
      * Revert attention dispatcher
      
      * update z-image docstring
      
      * styling
      
      * Recover attention_dispatch.py with its origin impl, later would special commit for fa3 compatibility.
      
      * Fix prev bug, and support for prompt_embeds pass in args after prompt pre-encode as List of torch Tensor.
      
      * Remove einop dependency.
      
      * remove redundant imports & make fix-copies
      
      * fix import
      
      ---------
      Co-authored-by: default avatarliudongyang <liudongyang0114@gmail.com>
      4088e8a8