-
Notifications
You must be signed in to change notification settings - Fork 6.7k
Feature/zimage inpaint pipeline #13006
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Feature/zimage inpaint pipeline #13006
Conversation
Updated the pipeline structure to include ZImageInpaintPipeline
alongside ZImagePipeline and ZImageImg2ImgPipeline.
Implemented the ZImageInpaintPipeline class for inpainting
tasks, including necessary methods for encoding prompts,
preparing masked latents, and denoising.
Enhanced the auto_pipeline to map the new ZImageInpaintPipeline
for inpainting generation tasks.
Added unit tests for ZImageInpaintPipeline to ensure
functionality and performance.
Updated dummy objects to include ZImageInpaintPipeline for
testing purposes.
- Add torch.empty fix for x_pad_token and cap_pad_token in test - Add # Copied from annotations for encode_prompt methods - Add documentation with usage example and autodoc directive
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
This PR adds a comprehensive inpainting pipeline for Z-Image, extending the existing Z-Image family of pipelines (text-to-image, img2img, controlnet) with mask-based inpainting capabilities. The implementation follows established patterns from other inpainting pipelines in the diffusers library while adapting to Z-Image's specific requirements (e.g., complex64 RoPE embeddings, flow matching scheduler).
Changes:
- Implemented ZImageInpaintPipeline with full inpainting support including mask blending and strength-based denoising control
- Added comprehensive test suite covering inference, batch processing, strength validation, mask functionality, VAE tiling, and device offloading
- Integrated the pipeline into auto_pipeline infrastructure with proper model mapping for "z-image" model type
- Updated all necessary init.py files and dummy objects for proper exports
- Added documentation with usage examples
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| src/diffusers/pipelines/z_image/pipeline_z_image_inpaint.py | New inpainting pipeline implementation with prepare_mask_latents, prepare_latents, and main call method for mask-based image inpainting |
| tests/pipelines/z_image/test_z_image_inpaint.py | Comprehensive test suite including inference tests, strength parameter validation, mask functionality tests, and compatibility tests |
| src/diffusers/pipelines/z_image/init.py | Added ZImageInpaintPipeline to module exports |
| src/diffusers/pipelines/init.py | Added ZImageInpaintPipeline to main pipelines module exports |
| src/diffusers/init.py | Added ZImageInpaintPipeline to top-level diffusers exports |
| src/diffusers/pipelines/auto_pipeline.py | Mapped ZImageInpaintPipeline to "z-image" in AUTO_INPAINT_PIPELINES_MAPPING |
| src/diffusers/utils/dummy_torch_and_transformers_objects.py | Added dummy ZImageInpaintPipeline class for when dependencies are not available |
| docs/source/en/api/pipelines/z_image.md | Added inpainting section with usage example and API reference |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Add batch size validation and callback handling fixes per review, using diffusers conventions rather than suggested code verbatim.
|
@yiyixuxu Ready for re-review. |
What does this PR do?
This PR adds an inpainting pipeline for Z-Image. The summary of changes are below:
Closes issue #12752
Tested using a simple script:
Testing script
LoRA functionality is also supported (inherited from ZImageLoraLoaderMixin).
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@yiyixuxu @asomoza @sayakpaul