feat(zimage): Z-Image model integration (Phase 1) #9
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Implements base infrastructure for Z-Image-Turbo support as alternative to FLUX. Enables model selection via IMAGE_MODEL environment variable. ## Changes ### Core Architecture - Add abstract ImageGenerator base class for all generators - Implement ZImageGenerator with diffusers ZImagePipeline - Add model factory pattern for generator selection ### Configuration - Add IMAGE_MODEL setting (flux|zimage) - Add Z-Image specific config: ZIMAGE_MODEL, ZIMAGE_ATTENTION, ZIMAGE_COMPILE - Maintain backward compatibility (defaults to flux) ### Features - bfloat16 dtype support - Memory optimizations (attention slicing, VAE tiling) - Attention backend selection (_native_flash, _flash_3, _sdpa) - Model compilation support for H100/H800 - Week-based output organization - Bilingual prompt support (EN/ZH) ### Documentation - Add comprehensive integration plan (Z_IMAGE_INTEGRATION_PLAN.md) - Update .env.example with Z-Image settings - Created GitHub issue #8 with 4-phase implementation plan ## Status - ✅ Phase 1.1-1.4: Core classes and configuration complete - ⏳ Phase 1.5-1.8: Factory pattern, CLI, tests (next) ## Testing Not yet implemented - requires diffusers from source: pip install git+https://github.com/huggingface/diffusers ## Related - Issue: #8 - Branch: feat/z-image-integration - Timeline: Week 1/4 of Phase 1
Phase 1 implementation complete with all core infrastructure: Core Components: - Factory pattern for model selection (factory.py) - CLI --model flag support for runtime selection - Comprehensive unit test suite (8 tests passing) - Pytest marker for slow/integration tests Changes: - src/generators/factory.py: New factory for creating generators - get_image_generator() with mock/zimage/flux support - get_available_models() for dynamic model detection - src/utils/cli.py: Added --model CLI parameter and factory integration - Both generate and loop commands now use factory pattern - Model validation with helpful error messages - pyproject.toml: Added loguru dependency and pytest markers - tests/test_zimage_generator.py: Full test coverage - Mock-based unit tests for all ZImageGenerator methods - Integration test marked as slow (requires CUDA) - Tests pass without requiring ZImagePipeline installation Testing: - All 8 unit tests passing - 1 integration test skipped (marked slow, requires GPU) - Tests use sys.modules mocking for ZImagePipeline import Related to #8 (Phase 1 tasks 1.5-1.8)
- Remove unused Python modules (s3_storage.py, embeddings.py) - Remove redundant scripts/format.py (use `just format`) - Fix 32 ruff issues (unused imports, bare excepts, ambiguous vars) - Fix bug: undefined `self` in image_generator.py decorator - Move 9 markdown docs from root to docs/ folder - All tests pass, ruff clean
- Rewrote zimage_generator.py to use Z-Image native API instead of non-existent ZImagePipeline from diffusers - Uses ref-repos/Z-Image/src for zimage.generate() and model loading - Updated config: ZIMAGE_MODEL_PATH for local model weights path - Updated factory.py: checks for ref-repos/Z-Image/src availability - Updated tests to match new native API - Added test_zimage_plugin_integration.py - Fixed transformers dtype compatibility (torch_dtype vs dtype) - All 34 tests pass, ruff clean Setup requirements: - Clone Z-Image: git clone https://github.com/Tongyi-MAI/Z-Image ref-repos/Z-Image - Download model: huggingface-cli download Tongyi-MAI/Z-Image-Turbo --local-dir ckpts/Z-Image-Turbo - Set IMAGE_MODEL=zimage in .env Closes #8
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
ref-repos/Z-Image/src(not diffusers pipeline)docs/Changes
src/generators/zimage_generator.py- Z-Image generator using native APItests/test_zimage_plugin_integration.py- Integration testsZIMAGE_MODEL_PATHfor local model weightsref-repos/Z-Image/srcavailabilitys3_storage.py,embeddings.py,scripts/format.pydocs/Setup Requirements
Test plan
Closes #8