* Add ZImageImg2ImgPipeline
Updated the pipeline structure to include ZImageImg2ImgPipeline
alongside ZImagePipeline.
Implemented the ZImageImg2ImgPipeline class for image-to-image
transformations, including necessary methods for
encoding prompts, preparing latents, and denoising.
Enhanced the auto_pipeline to map the new ZImageImg2ImgPipeline
for image generation tasks.
Added unit tests for ZImageImg2ImgPipeline to ensure
functionality and performance.
Updated dummy objects to include ZImageImg2ImgPipeline for
testing purposes.
* Address review comments for ZImageImg2ImgPipeline
- Add `# Copied from` annotations to encode_prompt and _encode_prompt
- Add ZImagePipeline to auto_pipeline.py for AutoPipeline support
* Add ZImage pipeline documentation
---------
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Γlvaro Somoza <asomoza@users.noreply.github.com>
* feat: Add `flow_prediction` to `prediction_type`, introduce `use_flow_sigmas`, `flow_shift`, `use_dynamic_shifting`, and `time_shift_type` parameters, and refine type hints for various arguments.
* style: reformat argument wrapping in `_convert_to_beta` and `index_for_timestep` method signatures.
* fix: group offloading to support standalone computational layers in block-level offloading
* test: for models with standalone and deeply nested layers in block-level offloading
* feat: support for block-level offloading in group offloading config
* fix: group offload block modules to AutoencoderKL and AutoencoderKLWan
* fix: update group offloading tests to use AutoencoderKL and adjust input dimensions
* refactor: streamline block offloading logic
* Apply style fixes
* update tests
* update
* fix for failing tests
* clean up
* revert to use skip_keys
* clean up
---------
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
* start zimage model tests.
* up
* up
* up
* up
* up
* up
* up
* up
* up
* up
* up
* up
* Revert "up"
This reverts commit bca3e27c96.
* expand upon compilation failure reason.
* Update tests/models/transformers/test_models_transformer_z_image.py
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* reinitialize the padding tokens to ones to prevent NaN problems.
* updates
* up
* skipping ZImage DiT tests
* up
* up
---------
Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com>
* Fix(peft): Re-apply group offloading after deleting adapters
* Test: Add regression test for group offloading + delete_adapters
* Test: Add assertions to verify output changes after deletion
* Test: Add try/finally to clean up group offloading hooks
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Fixes#12673.
Wrong default_stream is used. leading to wrong execution order when record_steram is enabled.
* update
* Update test
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* Refactor image padding logic to pervent zero tensor in transformer_z_image.py
* Apply style fixes
* Add more support to fix repeat bug on tpu devices.
* Fix for dynamo compile error for multi if-branches.
---------
Co-authored-by: Mingjia Li <mingjiali@tju.edu.cn>
Co-authored-by: Mingjia Li <mail@mingjia.li>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
* Add ZImage LoRA support and integrate into ZImagePipeline
* Add LoRA test for Z-Image
* Move the LoRA test
* Fix ZImage LoRA scale support and test configuration
* Add ZImage LoRA test overrides for architecture differences
- Override test_lora_fuse_nan to use ZImage's 'layers' attribute
instead of 'transformer_blocks'
- Skip block-level LoRA scaling test (not supported in ZImage)
- Add required imports: numpy, torch_device, check_if_lora_correctly_set
* Add ZImageLoraLoaderMixin to LoRA documentation
* Use conditional import for peft.LoraConfig in ZImage tests
* Override test_correct_lora_configs_with_different_ranks for ZImage
ZImage uses 'attention.to_k' naming convention instead of 'attn.to_k',
so the base test's module name search loop never finds a match. This
override uses the correct naming pattern for ZImage architecture.
* Add is_flaky decorator to ZImage LoRA tests initialise padding tokens
* Skip ZImage LoRA test class entirely
Skip the entire ZImageLoRATests class due to non-deterministic behavior
from complex64 RoPE operations and torch.empty padding tokens.
LoRA functionality works correctly with real models.
Clean up removed:
- Individual @unittest.skip decorators
- @is_flaky decorator overrides for inherited methods
- Custom test method overrides
- Global torch deterministic settings
- Unused imports (numpy, is_flaky, check_if_lora_correctly_set)
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
Co-authored-by: Γlvaro Somoza <asomoza@users.noreply.github.com>
* Fix examples not loading LoRA adapter weights from checkpoint
* Updated lora saving logic with accelerate save_model_hook and load_model_hook
* Formatted the changes using ruff
* import and upcasting changed
---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
* 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 fbf26b7ed1.
* 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
* Support for num_images_per_prompt>1; Remove redundant unquote variables.
* Fix bugs for num_images_per_prompt with actual batch.
* Add unit tests for Z-Image.
* Refine unitest and skip for cases needed separate test env; Fix compatibility with unitest in model, mostly precision formating.
* Add clean env for test_save_load_float16 separ test; Add Note; Styling.
* Update dtype mentioned by yiyi.
---------
Co-authored-by: liudongyang <liudongyang0114@gmail.com>
