Generalise in-context functionality (#203)
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: filipstrand <filipstrand@users.noreply.github.com>
This commit is contained in:
Filip Strand
GitHub
parent
6a41329a45
commit
847404748f
36
.github/workflows/release.yml
vendored
Normal file
36
.github/workflows/release.yml
vendored
Normal file
@ -0,0 +1,36 @@
|
||||
name: Create Release and Publish to PyPI
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
version:
|
||||
description: 'Release version (e.g., 0.8.1)'
|
||||
required: true
|
||||
type: string
|
||||
|
||||
jobs:
|
||||
release:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: '3.10'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
python -m pip install requests
|
||||
|
||||
- name: Run release script
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
VERSION: ${{ github.event.inputs.version }}
|
||||
TEST_PYPI_API_TOKEN: ${{ secrets.TEST_PYPI_API_TOKEN }}
|
||||
PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
|
||||
run: |
|
||||
python scripts/release.py --version ${{ github.event.inputs.version }}
|
||||
437
CHANGELOG.md
Normal file
437
CHANGELOG.md
Normal file
@ -0,0 +1,437 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to this project will be documented in this file.
|
||||
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
|
||||
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.8.0] - 2025-01-20
|
||||
|
||||
# MFLUX v.0.8.0 Release Notes
|
||||
|
||||
## Experimental AI Features
|
||||
|
||||
### 👗 CatVTON (Virtual Try-On)
|
||||
- **[EXPERIMENTAL]** Added virtual try-on capabilities using in-context learning via `mflux-generate-in-context-catvton`
|
||||
- Support for person image, person mask, and garment image inputs for comprehensive virtual clothing try-on
|
||||
- Automatic prompting for virtual try-on scenarios with optimized default prompts
|
||||
- Side-by-side generation showing garment product shot alongside styled result
|
||||
- AI-powered virtual clothing fitting with realistic lighting and fabric properties
|
||||
|
||||
### ✏️ IC-Edit (In-Context Editing)
|
||||
- **[EXPERIMENTAL]** Added natural language image editing capabilities via `mflux-generate-in-context-edit`
|
||||
- Natural language image editing using simple text instructions like "make the hair black" or "add sunglasses"
|
||||
- Automatic diptych template formatting for optimal editing results
|
||||
- Optimal resolution auto-sizing for 512px width (the resolution IC-Edit was trained on)
|
||||
- Specialized LoRA automatically downloaded and applied for enhanced editing capabilities
|
||||
|
||||
## Enhanced Generation Control
|
||||
|
||||
### 🔎 Image Upscaling
|
||||
- **Built-in upscaling capabilities**: Enhanced image quality and resolution enhancement for generated images
|
||||
- Seamless integration with existing generation workflow
|
||||
- Professional-grade upscaling for production-ready outputs
|
||||
|
||||
## Interpretability research
|
||||
|
||||
### 🧠 Concept Attention
|
||||
- **Enhanced image generation control**: Fine-grained control over image generation focus areas using attention-based concepts
|
||||
- Improved composition and subject handling for more precise artistic direction
|
||||
- Advanced attention mechanisms for better understanding of prompt concepts
|
||||
|
||||
## Workflow & Performance Improvements
|
||||
|
||||
### 🪫 Battery Saver
|
||||
- **Power management**: Automatic power optimization during extended generation sessions
|
||||
- Configurable power-saving modes specifically designed for laptop users
|
||||
- Smart resource management for long-running batch operations
|
||||
|
||||
### 📝 Prompt File Support
|
||||
- **File-based prompt input**: Batch operations via `--prompt-file` for large-scale generation projects
|
||||
- Dynamic prompt updates for large batch generation workflows
|
||||
- Support for external prompt management and automation systems
|
||||
|
||||
### 🔄 Redux Function Balancing
|
||||
- **Enhanced Redux capabilities**: Improved control over image-to-image transformation strength
|
||||
- Better quality variations with adjustable parameters for more predictable results
|
||||
- Refined Redux algorithm for more natural image variations
|
||||
|
||||
## Developer Experience
|
||||
|
||||
### 🔧 LORA_LIBRARY_PATH Improvements
|
||||
- **Unix-style resource discovery**: Enhanced LoRA library path handling for better organization
|
||||
- Improved path handling for LoRA weight discovery across multiple directories
|
||||
- Better cross-platform compatibility for LoRA management
|
||||
|
||||
### 🧪 Testing & Documentation
|
||||
- New command-line arguments for both experimental features with comprehensive help
|
||||
- Comprehensive argument parser tests for new functionality
|
||||
- Updated documentation with experimental feature warnings and usage guidelines
|
||||
- Added note about upcoming FLUX.1 Kontext model from Black Forest Labs
|
||||
|
||||
## Architecture Improvements
|
||||
|
||||
### 📚 Documentation Structure
|
||||
- Refactored "In-Context LoRA" section to "In-Context Generation" with clear subcategories
|
||||
- Enhanced documentation structure for better organization and user navigation
|
||||
- Improved categorization of experimental vs stable features
|
||||
|
||||
### ⚡ Performance Optimizations
|
||||
- Updated MLX dependency to latest version for improved performance and stability
|
||||
- Removed PyTorch dependency for DepthPro model, significantly reducing installation requirements
|
||||
- Streamlined dependencies for faster installation and reduced disk usage
|
||||
|
||||
## Experimental Notice
|
||||
|
||||
⚠️ **Important**: CatVTON and IC-Edit features are experimental and may be removed or significantly changed in future updates. These features represent cutting-edge AI capabilities that are still under active development.
|
||||
|
||||
## Contributors
|
||||
|
||||
Special thanks to the following contributors for their exceptional work since v0.7.1:
|
||||
- **Anthony Wu (@anthonywu)**: Battery Saver implementation, Prompt File Support, LORA_LIBRARY_PATH improvements
|
||||
- **Alessandro (@alessandro)**: Redux Function Balancing enhancements
|
||||
- **Filip Strand (@filipstrand)**: Core development, experimental features integration, infrastructure improvements
|
||||
|
||||
## [0.7.1] - 2025-05-06
|
||||
|
||||
# MFLUX v.0.7.1 Release Notes
|
||||
|
||||
## New Features
|
||||
|
||||
### 🎭 Multi-LoRA Support
|
||||
- **Multiple LoRA Loading**: Added support for loading multiple LoRA adapters simultaneously when using the in-context feature
|
||||
- Enhanced creative flexibility by combining multiple artistic styles in a single generation
|
||||
- Reference: [GitHub Issue #178](https://github.com/filipstrand/mflux/issues/178)
|
||||
|
||||
## [0.7.0] - 2025-04-25
|
||||
# MFLUX v.0.7.0 Release Notes
|
||||
|
||||
## Major New Features
|
||||
|
||||
### 🖌️ FLUX.1 Tools | Fill
|
||||
|
||||
- Added support for the FLUX.1-Fill model for inpainting and outpainting
|
||||
- Introduced `mflux-generate-fill` command-line tool for selective image editing
|
||||
- Implemented interactive mask creation tool to easily mark areas for regeneration
|
||||
- Added outpainting capabilities with customizable canvas expansion
|
||||
- Includes helper tools for creating outpaint image canvases and masks
|
||||
|
||||
### 🔍 FLUX.1 Tools | Depth
|
||||
|
||||
- Added support for the FLUX.1-Depth model for depth-conditioned image generation
|
||||
- Implemented Apple's ML Depth Pro model in MLX for state-of-the-art depth map extraction
|
||||
- Added `mflux-generate-depth` and `mflux-save-depth` command-line tools
|
||||
- Added ability to use either auto-generated depth maps or custom depth maps
|
||||
|
||||
### 🔄 FLUX.1 Tools | Redux
|
||||
|
||||
- Added Redux tool as a new image variation technique
|
||||
- Implemented a different approach compared to standard image-to-image generation
|
||||
- Uses image embedding joined with T5 text encodings for more natural variations
|
||||
- Added Redux-specific weight handlers and initialization
|
||||
|
||||
## New Models
|
||||
|
||||
### 🔎 Apple ML Depth Pro
|
||||
|
||||
- Added native MLX implementation of Apple's ML Depth Pro model for both separate use, and as a part of the Depth tool functionality
|
||||
|
||||
### 🖼️ Google SigLIP Vision Transformer
|
||||
|
||||
- Added SigLIP vision model for the Redux functionality
|
||||
|
||||
## Architecture Improvements
|
||||
|
||||
### 💾 Weight Management Improvements
|
||||
|
||||
- Added support for saving MFLUX version information in model metadata
|
||||
|
||||
### 🧠 Memory Optimization
|
||||
|
||||
- Additional improvements to the `--low-ram` option
|
||||
- Better memory management for image generation models
|
||||
|
||||
## Contributors
|
||||
|
||||
- @anthonywu
|
||||
- @ssakar
|
||||
- @akx
|
||||
|
||||
## [0.6.2] - 2025-03-13
|
||||
|
||||
# MFLUX v.0.6.2 Release Notes
|
||||
|
||||
## Bug Fixes
|
||||
|
||||
### 💾 Model Saving Fix
|
||||
- **Fixed local model saving**: Resolved bug preventing users from saving models locally with `mflux-save`
|
||||
- Restored full functionality for local model storage and management
|
||||
|
||||
## [0.6.1] - 2025-03-11
|
||||
|
||||
# MFLUX v.0.6.1 Release Notes
|
||||
|
||||
## Bug Fixes
|
||||
|
||||
### 🛑 Image Generation Interruption
|
||||
- **Fixed interruption flow**: Properly handles interruptions during image generation, ensuring graceful stops even when no callbacks are registered
|
||||
- **Keyboard interrupt handling**: Ensures image generation can be stopped via Ctrl+C in all diffusion model variants (standard Flux, ControlNet, and In-Context LoRA)
|
||||
- Relocated `StopImageGenerationException` from stepwise handler to main generation functions for more robust interruption system
|
||||
|
||||
## Test Stability Improvements
|
||||
|
||||
### 🧪 Test Reliability
|
||||
- **Fixed sporadic test failures**: Resolved intermittent failures in auto-seeds test case when using random seed count of 1
|
||||
- Improved test consistency and reliability
|
||||
|
||||
## Code Quality Improvements
|
||||
|
||||
### 🔧 Code Standards
|
||||
- **Formatting and linting fixes**: Fixed various formatting issues that were missed in the v0.6.0 release
|
||||
- Enhanced code consistency and maintainability
|
||||
|
||||
## [0.6.0] - 2025-03-05
|
||||
# MFLUX v.0.6.0 Release Notes
|
||||
|
||||
## Major New Features
|
||||
|
||||
### 🌐 Third-Party HuggingFace Model Support
|
||||
- Comprehensive ModelConfig refactor to support compatible HuggingFace dev/schnell models
|
||||
- Added ability to use models like `Freepik/flux.1-lite-8B-alpha` and `shuttleai/shuttle-3-diffusion`
|
||||
- New `--base-model` parameter to specify which base architecture (dev or schnell) a third-party model is derived from
|
||||
- Maintains backward compatibility while opening up the ecosystem to community-created models
|
||||
|
||||
### 🎭 In-Context LoRA
|
||||
- Added support for In-Context LoRA, a powerful technique that allows you to generate images in a specific style based on a reference image without requiring model fine-tuning
|
||||
- Introduced a new command-line tool: `mflux-generate-in-context`
|
||||
- Includes 10 pre-defined styles from the Hugging Face ali-vilab/In-Context-LoRA repository
|
||||
- Detailed documentation on how to use this feature effectively with prompting tips and best practices
|
||||
|
||||
### 🔌 Automatic LoRA Downloads
|
||||
- Added ability to automatically download LoRAs from Hugging Face when specified by repository ID
|
||||
- Simplifies workflow by eliminating the need to manually download LoRA files before use
|
||||
|
||||
### 🧠 Memory Optimizations
|
||||
- Added `--low-ram` option to reduce GPU memory usage by constraining the MLX cache size and releasing text encoders and transformer components after use
|
||||
- Implemented memory saver for ControlNet to reduce RAM requirements
|
||||
- General memory usage optimizations throughout the codebase
|
||||
|
||||
### 🗜️ Enhanced Quantization Options
|
||||
- Added support for 3-bit and 6-bit quantization (requires mlx > v0.21.0)
|
||||
- Expanded quantization options now include 3, 4, 6, and 8-bit precision
|
||||
|
||||
## ⚠️Breaking changes
|
||||
|
||||
Previously saved quantized models will not work for v.0.6.0 and later. See #149 for more details.
|
||||
|
||||
## Interface Improvements
|
||||
|
||||
### 🔧 Modified Parameters
|
||||
|
||||
- The previous `--init-image-path` parameter is now `--image-path`
|
||||
- The previous `--init-image-strength` parameter is now `--image-strength`
|
||||
|
||||
### 🖼️ Image Generation Enhancements
|
||||
- Added `--auto-seeds` option to generate multiple images with random seeds in a single command
|
||||
- Added option to override previously saved test images
|
||||
- Added `--controlnet-save-canny` option to save the Canny edge detection reference image used by ControlNet
|
||||
- Improved handling of edge cases for img2img generation
|
||||
|
||||
### 🔄 Callback System
|
||||
- Implemented a general callback mechanism for more flexible image generation pipelines
|
||||
- Added support for before-loop callbacks to accept latents
|
||||
- Enhanced StepwiseHandler to include initial latent
|
||||
|
||||
## Architecture Improvements
|
||||
|
||||
### 🏗️ Code Refactoring
|
||||
- Removed 'init' prefix for a more general interface
|
||||
- Removed `ConfigControlnet` - the `controlnet_strength` attribute is now on `Config`
|
||||
- Simplified quantization by removing unnecessary class predicates
|
||||
- Refactored model configuration system
|
||||
- Refactored transformer blocks for better maintainability
|
||||
- Unified attention mechanism in single and joint attention blocks
|
||||
- Added support for variable numbers of transformer blocks
|
||||
- Optimized with fast SDPA (Scaled Dot-Product Attention)
|
||||
- Added PromptCache for small optimization when generating with repeated prompts
|
||||
|
||||
### 🧰 Developer Tools
|
||||
- Added Batch Image Renamer tool as an isolated uv run script
|
||||
- Added descriptive comments for attention computations
|
||||
|
||||
## Compatibility Updates
|
||||
- Updated to support the latest mlx version
|
||||
- Fixed compatibility issues with HuggingFace dev/schnell models
|
||||
|
||||
## Bug Fixes
|
||||
- Fixed handling of edge cases for img2img generation
|
||||
- Various small fixes and improvements throughout the codebase
|
||||
|
||||
|
||||
## Contributors
|
||||
|
||||
- @anthonywu
|
||||
- @ssakar
|
||||
- @azrahello
|
||||
- @DanaCase
|
||||
|
||||
## [0.5.1] - 2024-12-23
|
||||
|
||||
# MFLUX v.0.5.1 Release Notes
|
||||
|
||||
## Bug Fixes
|
||||
|
||||
### 🔧 LoRA Loading Fix
|
||||
- **Quantized model LoRA compatibility**: Fixed critical bug where locally saved quantized models failed to set LoRA weights
|
||||
- Users can now successfully combine local quantized models with external LoRA adapters
|
||||
- Improved reliability for advanced workflows combining quantization and LoRA fine-tuning
|
||||
|
||||
## [0.5.0] - 2024-12-22
|
||||
|
||||
# MFLUX v.0.5.0 Release Notes
|
||||
|
||||
## Major New Features
|
||||
|
||||
### 🎛️ DreamBooth Fine-tuning
|
||||
- **DreamBooth support**: Introduced V1 of fine-tuning support in MFLUX
|
||||
- Enables custom model training for personalized image generation
|
||||
- Full fine-tuning pipeline with training configuration options
|
||||
|
||||
## Architecture Improvements
|
||||
|
||||
### 🔧 Weight Management Overhaul
|
||||
- **Rewritten LoRA handling**: Completely rewritten LoRA weight handling system
|
||||
- Improved performance and reliability for LoRA operations
|
||||
- Better support for complex LoRA workflows
|
||||
|
||||
## Developer Experience
|
||||
|
||||
### 🧪 Testing & Quality
|
||||
- **Enhanced test coverage**: Added comprehensive tests for new and existing features
|
||||
- Multi-LoRA testing support
|
||||
- Local model saving test coverage
|
||||
|
||||
### 📊 New Dependencies
|
||||
- **Matplotlib integration**: Added matplotlib for visualizing training loss during fine-tuning
|
||||
- **TOML support**: Added TOML library for better handling of MFLUX version metadata
|
||||
- Enhanced configuration management
|
||||
|
||||
## [0.4.1] - 2024-10-29
|
||||
|
||||
# MFLUX v.0.4.1 Release Notes
|
||||
|
||||
## Bug Fixes
|
||||
|
||||
### 🐛 Image Generation Fixes
|
||||
- **Img2img resolution fix**: Fixed img2img functionality for non-square image resolutions
|
||||
- Improved compatibility with various aspect ratios
|
||||
|
||||
## [0.4.0] - 2024-10-28
|
||||
|
||||
# MFLUX v.0.4.0 Release Notes
|
||||
|
||||
## Major New Features
|
||||
|
||||
### 🖼️ Image-to-Image Generation
|
||||
- **Img2Img Support**: Introduced the ability to generate images based on an initial reference image
|
||||
- Transform existing images using AI-powered generation techniques
|
||||
- Control the strength of transformation to balance between original image preservation and creative generation
|
||||
- Perfect for iterating on designs and creating variations of existing artwork
|
||||
|
||||
### 📊 Metadata-Driven Generation
|
||||
- **Image Generation from Metadata**: Added support to generate images directly from provided metadata files
|
||||
- Streamlined workflow for recreating images with specific parameters
|
||||
- Enhanced reproducibility for professional and research workflows
|
||||
- Automated parameter loading from previously generated images
|
||||
|
||||
### 🔍 Real-time Generation Monitoring
|
||||
- **Progressive Step Output**: Optionally output each step of the image generation process for real-time monitoring
|
||||
- Visual feedback during generation for better understanding of the AI process
|
||||
- Debug and fine-tune generation parameters by observing intermediate steps
|
||||
- Educational tool for understanding diffusion model progression
|
||||
|
||||
## Developer Experience Improvements
|
||||
|
||||
### 🛠️ Enhanced Command-Line Interface
|
||||
- **Improved argument handling**: Enhanced parsing and validation for command-line arguments
|
||||
- Better error messages and user guidance for parameter configuration
|
||||
- More intuitive command structure for complex generation workflows
|
||||
|
||||
### 🧪 Testing & Quality Assurance
|
||||
- **Automated Testing**: Added comprehensive automatic tests for image generation and command-line argument handling
|
||||
- Improved reliability and stability for all generation modes
|
||||
- Continuous integration testing for better code quality
|
||||
|
||||
### 🔧 Development Workflow
|
||||
- **Pre-Commit Hooks**: Integrated pre-commit hooks with `ruff`, `isort`, and typo checks for better code consistency
|
||||
- Enhanced developer experience with automated code quality checks
|
||||
- Streamlined contribution process for open source development
|
||||
|
||||
## [0.3.0] - 2024-09-24
|
||||
|
||||
# MFLUX v.0.3.0 Release Notes
|
||||
|
||||
## Major New Features
|
||||
|
||||
### 🕹️ ControlNet Support
|
||||
- **ControlNet Canny support**: Added Canny edge detection ControlNet functionality for precise image control
|
||||
- Enhanced control over image generation with edge-guided conditioning
|
||||
|
||||
## Model Export Improvements
|
||||
|
||||
### 📦 Advanced Model Export
|
||||
- **Quantized model export with LoRA**: Added ability to export quantized models with LoRA weights baked in
|
||||
- Streamlined deployment for fine-tuned models
|
||||
|
||||
## Developer Experience
|
||||
|
||||
### 🛠️ Development Tools
|
||||
- **Enhanced development workflow**: Improved developer experience with uv, ruff, makefile, pre-commit hooks
|
||||
- Better code quality tools and automated checks
|
||||
- Streamlined contribution process
|
||||
|
||||
## Legal & Licensing
|
||||
|
||||
### ⚖️ Open Source License
|
||||
- **Official MIT license**: Established clear open source licensing for the project
|
||||
- Legal clarity for users and contributors
|
||||
|
||||
## [0.2.1] - 2024-09-14
|
||||
|
||||
# MFLUX v.0.2.1 Release Notes
|
||||
|
||||
## Improvements
|
||||
|
||||
### 🔧 LoRA Enhancements
|
||||
- **Enhanced LoRA support**: Improved compatibility and performance for LoRA weight loading
|
||||
- Better integration with existing workflows
|
||||
- Refined handling of LoRA adapters
|
||||
|
||||
## [0.2.0] - 2024-09-07
|
||||
|
||||
# MFLUX v.0.2.0 Release Notes
|
||||
|
||||
## Major Milestone
|
||||
|
||||
### 🚀 Official PyPI Release
|
||||
- **First official PyPI release**: `pip install mflux` - making MFLUX easily installable for everyone
|
||||
- Big thanks to @deto for letting us have the "mflux" name on PyPI!
|
||||
|
||||
## New Features
|
||||
|
||||
### 🎨 Core Image Generation
|
||||
- **Command-line tools**: Introduced dedicated commands for better user experience
|
||||
- `mflux-generate` for generating images
|
||||
- `mflux-save` for saving quantized models to disk
|
||||
- **🗜️ Quantization support**: Added support for quantized models with 4-bit and 8-bit precision
|
||||
- **LoRA weights**: Added support for loading trained LoRA (Low-Rank Adaptation) weights
|
||||
- **Automatic metadata**: Images now automatically save metadata when generated
|
||||
|
||||
## Developer Experience
|
||||
|
||||
### 📦 Distribution
|
||||
- Official packaging and distribution through PyPI
|
||||
- Simplified installation process for end users
|
||||
- Professional project structure and naming
|
||||
144
README.md
144
README.md
@ -27,7 +27,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
|
||||
* [Multi-LoRA](#multi-lora)
|
||||
* [LoRA Library Path](#lora-library-path)
|
||||
* [Supported LoRA formats (updated)](#supported-lora-formats-updated)
|
||||
- [🎭 In-Context LoRA](#-in-context-lora)
|
||||
- [🎭 In-Context Generation](#-in-context-generation)
|
||||
* [Available Styles](#available-styles)
|
||||
* [How It Works](#how-it-works)
|
||||
* [Tips for Best Results](#tips-for-best-results)
|
||||
@ -230,7 +230,35 @@ The `mflux-generate-in-context` command supports most of the same arguments as `
|
||||
|
||||
- **`--lora-style`** (optional, `str`, default: `None`): The style to use for In-Context LoRA generation. Choose from: `couple`, `storyboard`, `font`, `home`, `illustration`, `portrait`, `ppt`, `sandstorm`, `sparklers`, or `identity`.
|
||||
|
||||
See the [In-Context LoRA](#-in-context-lora) section for more details on how to use this feature effectively.
|
||||
See the [In-Context Generation](#-in-context-generation) section for more details on how to use this feature effectively.
|
||||
|
||||
#### 📜 CatVTON Command-Line Arguments
|
||||
|
||||
The `mflux-generate-in-context-catvton` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
- **`--person-image`** (required, `str`): Path to the person image who will "wear" the garment.
|
||||
|
||||
- **`--person-mask`** (required, `str`): Path to a binary mask image indicating where the garment should be applied on the person.
|
||||
|
||||
- **`--garment-image`** (required, `str`): Path to the garment image to be virtually worn.
|
||||
|
||||
- **`--save-full-image`** (optional, flag): Additionally save the full side-by-side image containing both the garment reference and the result.
|
||||
|
||||
See the [CatVTON (Virtual Try-On)](#-catvton-virtual-try-on) section for more details on this feature.
|
||||
|
||||
#### 📜 IC-Edit Command-Line Arguments
|
||||
|
||||
The `mflux-generate-in-context-edit` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
- **`--reference-image`** (required, `str`): Path to the reference image to be edited.
|
||||
|
||||
- **`--instruction`** (optional, `str`): Simple text instruction describing the desired edit (e.g., "make the hair black"). This will be automatically formatted as a diptych template. Either `--instruction` or `--prompt` is required.
|
||||
|
||||
- **`--save-full-image`** (optional, flag): Additionally save the full side-by-side image showing the original and edited versions.
|
||||
|
||||
**Note**: The IC-Edit tool automatically downloads and applies the required LoRA, and optimizes the resolution to 512px width for best results.
|
||||
|
||||
See the [IC-Edit (In-Context Editing)](#-ic-edit-in-context-editing) section for more details on this feature.
|
||||
|
||||
#### 📜 Redux Tool Command-Line Arguments
|
||||
|
||||
@ -842,9 +870,15 @@ To report additional formats, examples or other any suggestions related to LoRA
|
||||
|
||||
---
|
||||
|
||||
### 🎭 In-Context LoRA
|
||||
### 🎭 In-Context Generation
|
||||
|
||||
In-Context LoRA is a powerful technique that allows you to generate images in a specific style based on a reference image, without requiring model fine-tuning. This approach uses specialized LoRA weights that enable the model to understand and apply the visual context from your reference image to a new generation.
|
||||
In-Context Generation is a powerful collection of techniques that allows you to generate images based on reference images and context, without requiring model fine-tuning. MFLUX supports multiple in-context approaches, each designed for specific use cases ranging from style transfer to virtual try-on and image editing.
|
||||
|
||||
**Note**: Black Forest Labs has announced [FLUX.1 Kontext](https://bfl.ai/models/flux-kontext), a new model suite designed specifically for text-and-image-driven generation and editing with advanced in-context capabilities. This official model may replace existing third-party implementations in the future. FLUX.1 Kontext offers character consistency, local editing, style reference, and interactive speed optimizations. The dev variant is coming soon as an open-weights model.
|
||||
|
||||
#### 🎨 In-Context LoRA
|
||||
|
||||
In-Context LoRA allows you to generate images in a specific style based on a reference image. This approach uses specialized LoRA weights that enable the model to understand and apply the visual context from your reference image to a new generation.
|
||||
|
||||
This feature is based on the [In-Context LoRA for Diffusion Transformers](https://github.com/ali-vilab/In-Context-LoRA) project by Ali-ViLab.
|
||||
|
||||
@ -854,8 +888,7 @@ To use In-Context LoRA, you need:
|
||||
1. A reference image
|
||||
2. A style LoRA (optional - the in-context ability works without LoRAs, but they can significantly enhance the results)
|
||||
|
||||
|
||||
#### Available Styles
|
||||
##### Available Styles
|
||||
|
||||
MFLUX provides several pre-defined styles from the [Hugging Face ali-vilab/In-Context-LoRA repository](https://huggingface.co/ali-vilab/In-Context-LoRA) that you can use with the `--lora-style` argument:
|
||||
|
||||
@ -872,7 +905,7 @@ MFLUX provides several pre-defined styles from the [Hugging Face ali-vilab/In-Co
|
||||
| `sparklers` | Sparklers visual effect |
|
||||
| `identity` | Visual identity and branding design style |
|
||||
|
||||
#### How It Works
|
||||
##### How It Works
|
||||
|
||||
The In-Context LoRA generation creates a side-by-side image where:
|
||||
- The left side shows your reference image with noise applied
|
||||
@ -882,8 +915,7 @@ The final output is automatically cropped to show only the right half (the gener
|
||||
|
||||
|
||||
|
||||
|
||||
#### Prompting for In-Context LoRA
|
||||
##### Prompting for In-Context LoRA
|
||||
|
||||
For best results with In-Context LoRA, your prompt should describe both the reference image and the target image you want to generate. Use markers like `[IMAGE1]`, `[LEFT]`, or `[RIGHT]` to distinguish between the two parts.
|
||||
|
||||
@ -909,7 +941,7 @@ This prompt clearly describes both the reference image (after `[IMAGE1]`) and th
|
||||
|
||||
**Important**: In the current implementation, the reference image is ALWAYS placed on the left side of the composition, and the generated image on the right side. When using marker pairs in your prompt, the first marker (e.g., `[IMAGE1]`, `[LEFT]`, `[REFERENCE]`) always refers to your reference image, while the second marker (e.g., `[IMAGE2]`, `[RIGHT]`, `[OUTPUT]`) refers to what you want to generate.
|
||||
|
||||
#### Tips for Best Results
|
||||
##### Tips for Best Results
|
||||
|
||||
1. **Choose the right reference image**: Select a reference image with a clear composition and structure that matches your intended output.
|
||||
2. **Adjust guidance**: Higher guidance values (7.0-9.0) tend to produce results that more closely follow your prompt.
|
||||
@ -918,6 +950,95 @@ This prompt clearly describes both the reference image (after `[IMAGE1]`) and th
|
||||
5. **Detailed prompting**: Be specific about both the reference image and your desired output in your prompt.
|
||||
6. **Try without LoRA**: While LoRAs enhance the results, you can experiment without them to see the base in-context capabilities.
|
||||
|
||||
#### 👕 CatVTON (Virtual Try-On)
|
||||
|
||||
⚠️ **Experimental Feature**: CatVTON is an experimental feature and may be removed or significantly changed in future updates.
|
||||
|
||||
CatVTON enables virtual try-on capabilities using in-context learning. This approach allows you to generate images of people wearing different garments by providing a person image, a person mask, and a garment image.
|
||||
|
||||
The implementation is based on [@nftblackmagic/catvton-flux](https://github.com/nftblackmagic/catvton-flux) and uses the model weights from [xiaozaa/catvton-flux-beta](https://huggingface.co/xiaozaa/catvton-flux-beta) (approximately 24 GB download).
|
||||
|
||||
##### How to Use CatVTON
|
||||
|
||||
```sh
|
||||
mflux-generate-in-context-catvton \
|
||||
--person-image "person.jpg" \
|
||||
--person-mask "person_mask.png" \
|
||||
--garment-image "garment.jpg" \
|
||||
--steps 20 \
|
||||
--seed 42 \
|
||||
--height 1024 \
|
||||
--width 1024
|
||||
```
|
||||
|
||||
##### Required Inputs
|
||||
|
||||
- **Person Image**: A photo of the person who will "wear" the garment
|
||||
- **Person Mask**: A binary mask indicating the areas where the garment should be applied
|
||||
- **Garment Image**: The clothing item to be virtually worn
|
||||
|
||||
##### CatVTON Features
|
||||
|
||||
- **Automatic Prompting**: If no prompt is provided, CatVTON uses an optimized default prompt designed for virtual try-on scenarios
|
||||
- **Side-by-Side Generation**: Creates a diptych showing the garment product shot alongside the styled result
|
||||
- **Optimized for Clothing**: Specifically tuned for clothing and fashion applications
|
||||
- **High-Quality Results**: Maintains realistic lighting, shadows, and fabric properties
|
||||
|
||||
##### Tips for CatVTON
|
||||
|
||||
1. **High-Quality Inputs**: Use high-resolution, well-lit images for best results
|
||||
2. **Accurate Masks**: Ensure the person mask precisely covers the areas where the garment should appear
|
||||
3. **Consistent Lighting**: Match lighting conditions between person and garment images when possible
|
||||
4. **Garment Type**: Works best with clearly defined clothing items (shirts, dresses, jackets, etc.)
|
||||
|
||||
#### ✏️ IC-Edit (In-Context Editing)
|
||||
|
||||
⚠️ **Experimental Feature**: IC-Edit is an experimental feature and may be removed or significantly changed in future updates.
|
||||
|
||||
IC-Edit provides intuitive image editing capabilities using natural language instructions. This approach automatically applies a specialized LoRA and generates edited versions of your reference image based on simple text instructions.
|
||||
|
||||
The implementation is based on [@River-Zhang/ICEdit](https://github.com/River-Zhang/ICEdit).
|
||||
|
||||
##### How to Use IC-Edit
|
||||
|
||||
```sh
|
||||
mflux-generate-in-context-edit \
|
||||
--reference-image "original.jpg" \
|
||||
--instruction "make the hair black" \
|
||||
--steps 20 \
|
||||
--seed 42
|
||||
```
|
||||
|
||||
##### Key Features
|
||||
|
||||
- **Natural Language Instructions**: Use simple, descriptive instructions like "make the hair black", "add sunglasses", or "change the background to a beach"
|
||||
- **Automatic Prompting**: Instructions are automatically wrapped in a diptych template for optimal results
|
||||
- **Optimal Resolution**: Automatically resizes to 512px width (the resolution IC-Edit was trained on) while maintaining aspect ratio
|
||||
- **Specialized LoRA**: Automatically downloads and applies the IC-Edit LoRA for enhanced editing capabilities
|
||||
|
||||
##### IC-Edit Options
|
||||
|
||||
You can use either `--instruction` for simple edits or `--prompt` for full control:
|
||||
|
||||
```sh
|
||||
# Using instruction (recommended)
|
||||
mflux-generate-in-context-edit \
|
||||
--reference-image "photo.jpg" \
|
||||
--instruction "remove the glasses"
|
||||
|
||||
# Using custom prompt (advanced)
|
||||
mflux-generate-in-context-edit \
|
||||
--reference-image "photo.jpg" \
|
||||
--prompt "A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but with a vintage filter applied"
|
||||
```
|
||||
|
||||
##### Tips for IC-Edit
|
||||
|
||||
1. **Clear Instructions**: Use specific, actionable instructions for best results
|
||||
2. **Single Changes**: Focus on one edit at a time for more predictable results
|
||||
3. **Reference Quality**: Higher quality reference images generally produce better edits
|
||||
4. **Iterative Editing**: Use the output of one edit as input for the next to build complex changes
|
||||
|
||||
---
|
||||
|
||||
### 🛠️ Flux Tools
|
||||
@ -1422,7 +1543,7 @@ This will generate the following image
|
||||
- Some LoRA adapters does not work.
|
||||
- Currently, the supported controlnet is the [canny-only version](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny).
|
||||
- Dreambooth training currently does not support sending in training parameters as flags.
|
||||
- In-Context LoRA currently only supports a left-right image setup (reference image on left, generated image on right).
|
||||
- In-Context Generation features currently only support a left-right image setup (reference image on left, generated image on right).
|
||||
|
||||
### Optional Tool: Batch Image Renamer
|
||||
|
||||
@ -1457,7 +1578,6 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
|
||||
|
||||
### 🔬 Cool research / features to support
|
||||
- [ ] [PuLID](https://github.com/ToTheBeginning/PuLID)
|
||||
- [ ] [catvton-flux](https://github.com/nftblackmagic/catvton-flux)
|
||||
|
||||
### 🌱 Related projects
|
||||
|
||||
|
||||
@ -57,7 +57,9 @@ homepage = "https://github.com/filipstrand/mflux"
|
||||
[project.scripts]
|
||||
mflux-generate = "mflux.generate:main"
|
||||
mflux-generate-controlnet = "mflux.generate_controlnet:main"
|
||||
mflux-generate-in-context = "mflux.generate_in_context:main"
|
||||
mflux-generate-in-context = "mflux.generate_in_context_dev:main"
|
||||
mflux-generate-in-context-edit = "mflux.generate_in_context_edit:main"
|
||||
mflux-generate-in-context-catvton = "mflux.generate_in_context_catvton:main"
|
||||
mflux-generate-fill = "mflux.generate_fill:main"
|
||||
mflux-generate-depth = "mflux.generate_depth:main"
|
||||
mflux-generate-redux = "mflux.generate_redux:main"
|
||||
|
||||
288
scripts/release.py
Normal file
288
scripts/release.py
Normal file
@ -0,0 +1,288 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
class ReleaseManager:
|
||||
def __init__(self, github_token: str):
|
||||
self.github_token = github_token
|
||||
self.github_repo = "filipstrand/mflux"
|
||||
self.base_url = "https://api.github.com"
|
||||
self.headers = {
|
||||
"Authorization": f"token {github_token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
"User-Agent": "MFLUX-Release-Script",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def extract_version_from_ref(github_ref: str) -> str:
|
||||
if not github_ref.startswith("refs/tags/v."):
|
||||
raise ValueError(f"Invalid GitHub ref format: {github_ref}")
|
||||
return github_ref.replace("refs/tags/v.", "")
|
||||
|
||||
@staticmethod
|
||||
def extract_changelog_entry(version: str, changelog_path: Path = Path("CHANGELOG.md")) -> str:
|
||||
import re
|
||||
|
||||
if not changelog_path.exists():
|
||||
raise FileNotFoundError(f"Changelog file not found: {changelog_path}")
|
||||
|
||||
content = changelog_path.read_text(encoding="utf-8")
|
||||
|
||||
# Pattern to match version headers: ## [VERSION]
|
||||
version_pattern = rf"^## \[{re.escape(version)}\].*$"
|
||||
next_version_pattern = r"^## \[.*\].*$"
|
||||
|
||||
lines = content.splitlines()
|
||||
|
||||
# Find the start of the target version section
|
||||
start_idx = None
|
||||
for i, line in enumerate(lines):
|
||||
if re.match(version_pattern, line):
|
||||
start_idx = i + 1 # Start after the version header
|
||||
break
|
||||
|
||||
if start_idx is None:
|
||||
raise ValueError(f"Version {version} not found in changelog")
|
||||
|
||||
# Find the end of the section (next version header)
|
||||
end_idx = len(lines)
|
||||
for i in range(start_idx, len(lines)):
|
||||
if re.match(next_version_pattern, lines[i]):
|
||||
end_idx = i
|
||||
break
|
||||
|
||||
# Extract the content between version headers
|
||||
entry_lines = lines[start_idx:end_idx]
|
||||
|
||||
# Clean up: remove empty lines at start/end and strip trailing whitespace
|
||||
while entry_lines and not entry_lines[0].strip():
|
||||
entry_lines.pop(0)
|
||||
while entry_lines and not entry_lines[-1].strip():
|
||||
entry_lines.pop()
|
||||
|
||||
entry_content = "\n".join(line.rstrip() for line in entry_lines)
|
||||
|
||||
if not entry_content.strip():
|
||||
raise ValueError(f"No content found for version {version}")
|
||||
|
||||
return entry_content
|
||||
|
||||
def check_release_exists(self, tag_name: str) -> bool:
|
||||
url = f"{self.base_url}/repos/{self.github_repo}/releases/tags/{tag_name}"
|
||||
response = requests.get(url, headers=self.headers)
|
||||
return response.status_code == 200
|
||||
|
||||
def create_github_release(self, tag_name: str, version: str, changelog_entry: str) -> dict:
|
||||
url = f"{self.base_url}/repos/{self.github_repo}/releases"
|
||||
|
||||
data = {
|
||||
"tag_name": tag_name,
|
||||
"name": f"Release {version}",
|
||||
"body": changelog_entry,
|
||||
"draft": False,
|
||||
"prerelease": False,
|
||||
}
|
||||
|
||||
response = requests.post(url, json=data, headers=self.headers)
|
||||
|
||||
if response.status_code == 201:
|
||||
print(f"✅ Successfully created GitHub release for {tag_name}")
|
||||
return response.json()
|
||||
else:
|
||||
raise Exception(f"Failed to create GitHub release: {response.status_code} - {response.text}")
|
||||
|
||||
@staticmethod
|
||||
def run_command(cmd: list, description: str, check: bool = True) -> subprocess.CompletedProcess:
|
||||
print(f"🔄 {description}...")
|
||||
print(f" Running: {' '.join(cmd)}")
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, check=False)
|
||||
|
||||
if result.returncode != 0 and check:
|
||||
print(f"❌ {description} failed!")
|
||||
print(f" stdout: {result.stdout}")
|
||||
print(f" stderr: {result.stderr}")
|
||||
raise Exception(f"Command failed: {' '.join(cmd)}")
|
||||
|
||||
if result.stdout:
|
||||
print(f" stdout: {result.stdout.strip()}")
|
||||
|
||||
print(f"✅ {description} completed successfully")
|
||||
return result
|
||||
|
||||
def build_package(self):
|
||||
# Install build dependencies
|
||||
self.run_command(
|
||||
[sys.executable, "-m", "pip", "install", "--upgrade", "pip", "build", "twine"],
|
||||
"Installing build dependencies",
|
||||
)
|
||||
|
||||
# Build the package
|
||||
self.run_command([sys.executable, "-m", "build"], "Building package")
|
||||
|
||||
def publish_to_test_pypi(self, test_pypi_token: Optional[str]):
|
||||
if not test_pypi_token:
|
||||
print("⚠️ Test PyPI token not provided, skipping Test PyPI upload")
|
||||
return
|
||||
|
||||
env = os.environ.copy()
|
||||
env["TWINE_USERNAME"] = "__token__"
|
||||
env["TWINE_PASSWORD"] = test_pypi_token
|
||||
|
||||
self.run_command(
|
||||
[sys.executable, "-m", "twine", "upload", "--repository", "testpypi", "dist/*", "--verbose"],
|
||||
"Publishing to Test PyPI",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def publish_to_pypi(pypi_token: str):
|
||||
env = os.environ.copy()
|
||||
env["TWINE_USERNAME"] = "__token__"
|
||||
env["TWINE_PASSWORD"] = pypi_token
|
||||
|
||||
# Set environment variables for subprocess
|
||||
cmd_env = os.environ.copy()
|
||||
cmd_env.update(env)
|
||||
|
||||
print("🔄 Publishing to PyPI...")
|
||||
print(f" Running: {sys.executable} -m twine upload dist/* --verbose")
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-m", "twine", "upload", "dist/*", "--verbose"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=cmd_env,
|
||||
check=False,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print("❌ Publishing to PyPI failed!")
|
||||
print(f" stdout: {result.stdout}")
|
||||
print(f" stderr: {result.stderr}")
|
||||
raise Exception("PyPI upload failed")
|
||||
|
||||
if result.stdout:
|
||||
print(f" stdout: {result.stdout.strip()}")
|
||||
|
||||
print("✅ Publishing to PyPI completed successfully")
|
||||
|
||||
def release(self, version: str, test_pypi_token: Optional[str], pypi_token: str):
|
||||
print("🚀 Starting MFLUX release process...")
|
||||
|
||||
# Extract version
|
||||
tag_name = f"v.{version}"
|
||||
print(f"📦 Releasing version: {version} (tag: {tag_name})")
|
||||
|
||||
# Extract changelog
|
||||
print("📝 Extracting changelog entry...")
|
||||
changelog_entry = self.extract_changelog_entry(version)
|
||||
print(f" Found changelog entry ({len(changelog_entry)} characters)")
|
||||
|
||||
# Check if release exists
|
||||
print("🔍 Checking if release already exists...")
|
||||
if self.check_release_exists(tag_name):
|
||||
print(f"⚠️ Release {tag_name} already exists, skipping GitHub release creation")
|
||||
skip_github_release = True
|
||||
else:
|
||||
skip_github_release = False
|
||||
|
||||
# Create GitHub release
|
||||
if not skip_github_release:
|
||||
self.create_github_release(tag_name, version, changelog_entry)
|
||||
|
||||
# Build package
|
||||
self.build_package()
|
||||
|
||||
# Publish to Test PyPI
|
||||
if not skip_github_release: # Only publish if this is a new release
|
||||
self.publish_to_test_pypi(test_pypi_token)
|
||||
|
||||
# Publish to PyPI
|
||||
self.publish_to_pypi(pypi_token)
|
||||
else:
|
||||
print("⚠️ Skipping PyPI publishing since release already exists")
|
||||
|
||||
print(f"🎉 Release process completed successfully for version {version}!")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="MFLUX Release Management")
|
||||
parser.add_argument("--version", help="Release version (e.g., 0.8.0)")
|
||||
parser.add_argument("--github-token", help="GitHub token")
|
||||
parser.add_argument("--test-pypi-token", help="Test PyPI API token (optional)")
|
||||
parser.add_argument("--pypi-token", help="PyPI API token")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Print what would be done without executing")
|
||||
parser.add_argument("--show-changelog", help="Show changelog entry for a specific version (e.g., 0.8.0) and exit")
|
||||
parser.add_argument("--markdown", action="store_true", help="Output changelog in markdown format (use with --show-changelog)") # fmt: off
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Handle --show-changelog option
|
||||
if args.show_changelog:
|
||||
try:
|
||||
changelog_entry = ReleaseManager.extract_changelog_entry(args.show_changelog)
|
||||
|
||||
if args.markdown:
|
||||
# Output as Markdown
|
||||
print(f"# Release {args.show_changelog}\n")
|
||||
print(changelog_entry)
|
||||
print(f"\n---\n*Changelog entry: {len(changelog_entry)} characters*")
|
||||
else:
|
||||
# Output with separators (original format)
|
||||
print(f"📝 Changelog entry for version {args.show_changelog}:")
|
||||
print("=" * 60)
|
||||
print(changelog_entry)
|
||||
print("=" * 60)
|
||||
print(f"✅ Found changelog entry ({len(changelog_entry)} characters)")
|
||||
except (ValueError, FileNotFoundError) as e:
|
||||
print(f"❌ Failed to extract changelog: {e}")
|
||||
sys.exit(1)
|
||||
return
|
||||
|
||||
# Get values from args or environment variables
|
||||
version = args.version or os.getenv("VERSION")
|
||||
github_token = args.github_token or os.getenv("GITHUB_TOKEN")
|
||||
test_pypi_token = args.test_pypi_token or os.getenv("TEST_PYPI_API_TOKEN")
|
||||
pypi_token = args.pypi_token or os.getenv("PYPI_API_TOKEN")
|
||||
|
||||
# Validate required parameters
|
||||
if not version:
|
||||
print("❌ Version is required (--version or VERSION env var)")
|
||||
sys.exit(1)
|
||||
|
||||
if not github_token:
|
||||
print("❌ GitHub token is required (--github-token or GITHUB_TOKEN env var)")
|
||||
sys.exit(1)
|
||||
|
||||
if not pypi_token:
|
||||
print("❌ PyPI token is required (--pypi-token or PYPI_API_TOKEN env var)")
|
||||
sys.exit(1)
|
||||
|
||||
if args.dry_run:
|
||||
print("🔍 DRY RUN MODE - would execute release process with:")
|
||||
print(f" Version: {version}")
|
||||
print(f" Has Test PyPI token: {bool(test_pypi_token)}")
|
||||
print(f" Has PyPI token: {bool(pypi_token)}")
|
||||
return
|
||||
|
||||
try:
|
||||
release_manager = ReleaseManager(github_token)
|
||||
release_manager.release(version, test_pypi_token, pypi_token)
|
||||
except (ValueError, FileNotFoundError, requests.RequestException) as e:
|
||||
print(f"❌ Release failed: {e}")
|
||||
sys.exit(1)
|
||||
except Exception as e: # noqa: BLE001
|
||||
print(f"❌ Release failed with unexpected error: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,3 +0,0 @@
|
||||
"""
|
||||
Third-party extensions and integrations for MFlux.
|
||||
"""
|
||||
0
src/mflux/community/in_context/__init__.py
Normal file
0
src/mflux/community/in_context/__init__.py
Normal file
@ -19,7 +19,7 @@ from mflux.post_processing.generated_image import GeneratedImage
|
||||
from mflux.post_processing.image_util import ImageUtil
|
||||
|
||||
|
||||
class Flux1InContextLoRA(nn.Module):
|
||||
class Flux1InContextDev(nn.Module):
|
||||
vae: VAE
|
||||
transformer: Transformer
|
||||
t5_text_encoder: T5Encoder
|
||||
@ -66,7 +66,7 @@ class Flux1InContextLoRA(nn.Module):
|
||||
)
|
||||
|
||||
# 2. Create the initial latents and keep the initial static noise for later blending
|
||||
static_noise = Flux1InContextLoRA._create_in_context_latents(seed=seed, config=config)
|
||||
static_noise = Flux1InContextDev._create_in_context_latents(seed=seed, config=config)
|
||||
latents = mx.array(static_noise)
|
||||
|
||||
# 3. Encode the prompt
|
||||
@ -102,8 +102,8 @@ class Flux1InContextLoRA(nn.Module):
|
||||
dt = config.sigmas[t + 1] - config.sigmas[t]
|
||||
latents += noise * dt
|
||||
|
||||
# 6.t Override left hand side of latents by linearly interpolating between latents and static noise
|
||||
latents = Flux1InContextLoRA._update_latents(
|
||||
# 6.t Override the left-hand side of latents by linearly interpolating between latents and static noise
|
||||
latents = Flux1InContextDev._update_latents(
|
||||
t=t,
|
||||
config=config,
|
||||
latents=latents,
|
||||
165
src/mflux/community/in_context/flux_in_context_fill.py
Normal file
165
src/mflux/community/in_context/flux_in_context_fill.py
Normal file
@ -0,0 +1,165 @@
|
||||
import mlx.core as mx
|
||||
from mlx import nn
|
||||
from tqdm import tqdm
|
||||
|
||||
from mflux.callbacks.callbacks import Callbacks
|
||||
from mflux.community.in_context.utils.in_context_mask_util import InContextMaskUtil
|
||||
from mflux.config.config import Config
|
||||
from mflux.config.model_config import ModelConfig
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
from mflux.error.exceptions import StopImageGenerationException
|
||||
from mflux.flux.flux_initializer import FluxInitializer
|
||||
from mflux.latent_creator.latent_creator import LatentCreator
|
||||
from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
|
||||
from mflux.models.text_encoder.prompt_encoder import PromptEncoder
|
||||
from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
|
||||
from mflux.models.transformer.transformer import Transformer
|
||||
from mflux.models.vae.vae import VAE
|
||||
from mflux.post_processing.array_util import ArrayUtil
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
from mflux.post_processing.image_util import ImageUtil
|
||||
|
||||
|
||||
class FluxInContextFill(nn.Module):
|
||||
vae: VAE
|
||||
transformer: Transformer
|
||||
t5_text_encoder: T5Encoder
|
||||
clip_text_encoder: CLIPEncoder
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_config: ModelConfig,
|
||||
quantize: int | None = None,
|
||||
local_path: str | None = None,
|
||||
lora_paths: list[str] | None = None,
|
||||
lora_scales: list[float] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
FluxInitializer.init(
|
||||
flux_model=self,
|
||||
model_config=model_config,
|
||||
quantize=quantize,
|
||||
local_path=local_path,
|
||||
lora_paths=lora_paths,
|
||||
lora_scales=lora_scales,
|
||||
)
|
||||
|
||||
def generate_image(
|
||||
self,
|
||||
seed: int,
|
||||
prompt: str,
|
||||
config: Config,
|
||||
left_image_path: str,
|
||||
right_image_path: str | None = None,
|
||||
) -> GeneratedImage:
|
||||
# 0. Create a new runtime config based on the model type and input parameters
|
||||
config = RuntimeConfig(config, self.model_config)
|
||||
|
||||
# For in-context learning with side-by-side approach, double the width
|
||||
original_width = config.width
|
||||
config.width = original_width * 2
|
||||
|
||||
time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
|
||||
|
||||
# 1. Create the initial latents
|
||||
latents = LatentCreator.create(
|
||||
seed=seed,
|
||||
height=config.height,
|
||||
width=config.width,
|
||||
)
|
||||
|
||||
# 2. Encode the prompt
|
||||
prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
|
||||
prompt=prompt,
|
||||
prompt_cache=self.prompt_cache,
|
||||
t5_tokenizer=self.t5_tokenizer,
|
||||
clip_tokenizer=self.clip_tokenizer,
|
||||
t5_text_encoder=self.t5_text_encoder,
|
||||
clip_text_encoder=self.clip_text_encoder,
|
||||
)
|
||||
|
||||
# 3. Create the static masked latents
|
||||
static_masked_latents = InContextMaskUtil.create_masked_latents(
|
||||
vae=self.vae,
|
||||
height=config.height,
|
||||
width=config.width,
|
||||
original_width=original_width,
|
||||
left_image_path=left_image_path,
|
||||
right_image_path=right_image_path,
|
||||
mask_path=config.masked_image_path,
|
||||
)
|
||||
|
||||
# (Optional) Call subscribers for beginning of loop
|
||||
Callbacks.before_loop(
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
for t in time_steps:
|
||||
try:
|
||||
# 4.t Concatenate the updated latents with the static masked latents
|
||||
hidden_states = mx.concatenate([latents, static_masked_latents], axis=-1)
|
||||
|
||||
# 5.t Predict the noise
|
||||
noise = self.transformer(
|
||||
t=t,
|
||||
config=config,
|
||||
hidden_states=hidden_states,
|
||||
prompt_embeds=prompt_embeds,
|
||||
pooled_prompt_embeds=pooled_prompt_embeds,
|
||||
)
|
||||
|
||||
# 6.t Take one denoise step
|
||||
dt = config.sigmas[t + 1] - config.sigmas[t]
|
||||
latents += noise * dt
|
||||
|
||||
# (Optional) Call subscribers in-loop
|
||||
Callbacks.in_loop(
|
||||
t=t,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
time_steps=time_steps,
|
||||
)
|
||||
|
||||
# (Optional) Evaluate to enable progress tracking
|
||||
mx.eval(latents)
|
||||
|
||||
except KeyboardInterrupt: # noqa: PERF203
|
||||
Callbacks.interruption(
|
||||
t=t,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
time_steps=time_steps,
|
||||
)
|
||||
raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(time_steps)}")
|
||||
|
||||
# (Optional) Call subscribers after loop
|
||||
Callbacks.after_loop(
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
latents=latents,
|
||||
config=config,
|
||||
)
|
||||
|
||||
# 7. Decode the latent array and return the full image
|
||||
latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
|
||||
decoded = self.vae.decode(latents)
|
||||
return ImageUtil.to_image(
|
||||
decoded_latents=decoded,
|
||||
config=config,
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
quantization=self.bits,
|
||||
lora_paths=self.lora_paths,
|
||||
lora_scales=self.lora_scales,
|
||||
image_path=right_image_path,
|
||||
image_strength=config.image_strength,
|
||||
masked_image_path=config.masked_image_path,
|
||||
generation_time=time_steps.format_dict["elapsed"],
|
||||
)
|
||||
0
src/mflux/community/in_context/utils/__init__.py
Normal file
0
src/mflux/community/in_context/utils/__init__.py
Normal file
@ -1,5 +1,7 @@
|
||||
"""Configuration for In-Context LoRAs from Hugging Face."""
|
||||
|
||||
from mflux.weights.weight_handler_lora_huggingface import WeightHandlerLoRAHuggingFace
|
||||
|
||||
# Default Hugging Face repository for In-Context LoRAs
|
||||
LORA_REPO_ID = "ali-vilab/In-Context-LoRA"
|
||||
|
||||
@ -17,9 +19,37 @@ LORA_NAME_MAP = {
|
||||
"identity": "visual-identity-design.safetensors",
|
||||
}
|
||||
|
||||
# IC-Edit specific configuration
|
||||
IC_EDIT_LORA_REPO_ID = "RiverZ/normal-lora"
|
||||
IC_EDIT_LORA_FILENAME = "pytorch_lora_weights.safetensors"
|
||||
IC_EDIT_LORA_SCALE = 1.0
|
||||
|
||||
|
||||
def get_lora_filename(simple_name: str) -> str:
|
||||
if simple_name not in LORA_NAME_MAP:
|
||||
valid_names = ", ".join(sorted(LORA_NAME_MAP.keys()))
|
||||
raise ValueError(f"Unknown LoRA name: {simple_name}. Valid names are: {valid_names}")
|
||||
return LORA_NAME_MAP[simple_name]
|
||||
|
||||
|
||||
def prepare_ic_edit_loras(additional_lora_paths: list[str] | None = None) -> list[str]:
|
||||
print(f"🔍 Downloading IC-Edit LoRA from {IC_EDIT_LORA_REPO_ID}")
|
||||
|
||||
# Download the required IC-Edit LoRA
|
||||
ic_edit_lora_path = WeightHandlerLoRAHuggingFace.download_lora(
|
||||
repo_id=IC_EDIT_LORA_REPO_ID,
|
||||
lora_name=IC_EDIT_LORA_FILENAME,
|
||||
)
|
||||
|
||||
# IC-Edit LoRA is always required and goes first
|
||||
lora_paths = [ic_edit_lora_path]
|
||||
|
||||
# Add any additional user-specified LoRAs
|
||||
if additional_lora_paths:
|
||||
lora_paths.extend(additional_lora_paths)
|
||||
|
||||
print(f"✅ IC-Edit LoRA ready: {ic_edit_lora_path}")
|
||||
if len(lora_paths) > 1:
|
||||
print(f"📋 Additional LoRAs: {lora_paths[1:]}")
|
||||
|
||||
return lora_paths
|
||||
88
src/mflux/community/in_context/utils/in_context_mask_util.py
Normal file
88
src/mflux/community/in_context/utils/in_context_mask_util.py
Normal file
@ -0,0 +1,88 @@
|
||||
import mlx.core as mx
|
||||
|
||||
from mflux.flux_tools.fill.mask_util import MaskUtil
|
||||
from mflux.models.vae.vae import VAE
|
||||
from mflux.post_processing.array_util import ArrayUtil
|
||||
from mflux.post_processing.image_util import ImageUtil
|
||||
|
||||
|
||||
class InContextMaskUtil:
|
||||
@staticmethod
|
||||
def create_masked_latents(
|
||||
vae: VAE,
|
||||
height: int,
|
||||
width: int,
|
||||
original_width: int,
|
||||
left_image_path: str,
|
||||
right_image_path: str | None,
|
||||
mask_path: str,
|
||||
):
|
||||
# Determine mode based on whether right_image_path is provided
|
||||
is_pure_mode = right_image_path is None
|
||||
|
||||
# Step 1: Load left image (reference) - always required
|
||||
left_image = ImageUtil.scale_to_dimensions(
|
||||
image=ImageUtil.load_image(left_image_path).convert("RGB"),
|
||||
target_width=original_width,
|
||||
target_height=height,
|
||||
)
|
||||
left_array = ImageUtil.to_array(left_image)
|
||||
|
||||
# Step 2: Handle right image based on mode
|
||||
if is_pure_mode:
|
||||
# Pure mode: create empty/noise for right side
|
||||
right_array = mx.zeros_like(left_array) # or could use noise
|
||||
else:
|
||||
# Selective mode: load target image for right side
|
||||
right_image = ImageUtil.scale_to_dimensions(
|
||||
image=ImageUtil.load_image(right_image_path).convert("RGB"),
|
||||
target_width=original_width,
|
||||
target_height=height,
|
||||
)
|
||||
right_array = ImageUtil.to_array(right_image)
|
||||
|
||||
# Step 3: Determine the right-side mask based on mode
|
||||
if is_pure_mode:
|
||||
# Pure in-context: generate entire right side
|
||||
mask_array = mx.ones((1, 1, height, original_width))
|
||||
else:
|
||||
# Selective: use provided mask exactly
|
||||
mask_image = ImageUtil.scale_to_dimensions(
|
||||
image=ImageUtil.load_image(mask_path).convert("RGB"),
|
||||
target_width=original_width,
|
||||
target_height=height,
|
||||
)
|
||||
mask_array = ImageUtil.to_array(mask_image, is_mask=True)
|
||||
|
||||
# Step 4: Create concatenated inputs for in-context learning
|
||||
# Layout: [LEFT_IMAGE] | [RIGHT_IMAGE]
|
||||
concatenated_image = mx.concatenate([left_array, right_array], axis=3)
|
||||
|
||||
# Layout: [NO_MASK] | [ACTUAL_MASK]
|
||||
reference_mask = mx.zeros_like(mask_array) # Empty mask for reference (preserve fully)
|
||||
concatenated_mask = mx.concatenate([reference_mask, mask_array], axis=3)
|
||||
|
||||
# Step 5: Apply standard FLUX Fill processing to concatenated inputs
|
||||
masked_concatenated_image = concatenated_image * (1 - concatenated_mask)
|
||||
|
||||
# Encode and process exactly like MaskUtil.create_masked_latents
|
||||
encoded_image = vae.encode(masked_concatenated_image)
|
||||
encoded_image = ArrayUtil.pack_latents(
|
||||
latents=encoded_image,
|
||||
height=height,
|
||||
width=width,
|
||||
)
|
||||
|
||||
processed_mask = MaskUtil.reshape_mask(
|
||||
the_mask=concatenated_mask,
|
||||
height=height,
|
||||
width=width,
|
||||
)
|
||||
processed_mask = ArrayUtil.pack_latents(
|
||||
latents=processed_mask,
|
||||
height=height,
|
||||
width=width,
|
||||
num_channels_latents=64,
|
||||
)
|
||||
|
||||
return mx.concatenate([encoded_image, processed_mask], axis=-1)
|
||||
@ -1,6 +0,0 @@
|
||||
"""
|
||||
In-context LoRA implementation for MFlux.
|
||||
|
||||
This module provides functionality for conditioning image generation using a reference image
|
||||
directly in the latent space, without requiring model fine-tuning.
|
||||
"""
|
||||
@ -2,6 +2,7 @@ from mflux import Config, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.community.concept_attention.flux_concept import Flux1Concept
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -18,6 +19,10 @@ def main():
|
||||
parser.add_concept_attention_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Set default guidance value if not provided by user
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.GUIDANCE_SCALE
|
||||
|
||||
# 1. Load the concept attention model
|
||||
flux = Flux1Concept(
|
||||
model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
|
||||
|
||||
@ -2,6 +2,7 @@ from mflux import Config, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.community.concept_attention.flux_concept_from_image import FluxConceptFromImage
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -18,6 +19,10 @@ def main():
|
||||
parser.add_concept_from_image_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Set default guidance value if not provided by user
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.GUIDANCE_SCALE
|
||||
|
||||
# 1. Load the concept attention model
|
||||
flux = FluxConceptFromImage(
|
||||
model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
|
||||
|
||||
@ -11,6 +11,7 @@ class ModelConfig:
|
||||
model_name: str,
|
||||
base_model: str | None,
|
||||
controlnet_model: str | None,
|
||||
custom_transformer_model: str | None,
|
||||
num_train_steps: int,
|
||||
max_sequence_length: int,
|
||||
supports_guidance: bool,
|
||||
@ -21,6 +22,7 @@ class ModelConfig:
|
||||
self.model_name = model_name
|
||||
self.base_model = base_model
|
||||
self.controlnet_model = controlnet_model
|
||||
self.custom_transformer_model = custom_transformer_model
|
||||
self.num_train_steps = num_train_steps
|
||||
self.max_sequence_length = max_sequence_length
|
||||
self.supports_guidance = supports_guidance
|
||||
@ -37,6 +39,11 @@ class ModelConfig:
|
||||
def dev_fill() -> "ModelConfig":
|
||||
return AVAILABLE_MODELS["dev-fill"]
|
||||
|
||||
@staticmethod
|
||||
@lru_cache
|
||||
def dev_fill_catvton() -> "ModelConfig":
|
||||
return AVAILABLE_MODELS["dev-fill-catvton"]
|
||||
|
||||
@staticmethod
|
||||
@lru_cache
|
||||
def dev_depth() -> "ModelConfig":
|
||||
@ -68,7 +75,7 @@ class ModelConfig:
|
||||
return AVAILABLE_MODELS["schnell"]
|
||||
|
||||
def x_embedder_input_dim(self) -> int:
|
||||
if self.alias == "dev-fill":
|
||||
if self.alias and "dev-fill" in self.alias:
|
||||
return 384
|
||||
if self.alias == "dev-depth":
|
||||
return 128
|
||||
@ -116,6 +123,7 @@ class ModelConfig:
|
||||
model_name=model_name,
|
||||
base_model=default_base.model_name,
|
||||
controlnet_model=default_base.controlnet_model,
|
||||
custom_transformer_model=default_base.custom_transformer_model,
|
||||
num_train_steps=default_base.num_train_steps,
|
||||
max_sequence_length=default_base.max_sequence_length,
|
||||
supports_guidance=default_base.supports_guidance,
|
||||
@ -130,6 +138,7 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-schnell",
|
||||
base_model=None,
|
||||
controlnet_model=None,
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=256,
|
||||
supports_guidance=False,
|
||||
@ -141,6 +150,7 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-dev",
|
||||
base_model=None,
|
||||
controlnet_model=None,
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=True,
|
||||
@ -152,6 +162,7 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-Fill-dev",
|
||||
base_model=None,
|
||||
controlnet_model=None,
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=True,
|
||||
@ -163,6 +174,7 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-Depth-dev",
|
||||
base_model=None,
|
||||
controlnet_model=None,
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=True,
|
||||
@ -174,17 +186,31 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-Redux-dev",
|
||||
base_model=None,
|
||||
controlnet_model=None,
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=True,
|
||||
requires_sigma_shift=True,
|
||||
priority=3,
|
||||
),
|
||||
"dev-fill-catvton": ModelConfig(
|
||||
alias="dev-fill-catvton",
|
||||
model_name="black-forest-labs/FLUX.1-Fill-dev",
|
||||
base_model=None,
|
||||
controlnet_model=None,
|
||||
custom_transformer_model="xiaozaa/catvton-flux-beta",
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=True,
|
||||
requires_sigma_shift=False, # Not sure why, but produced better results this way...
|
||||
priority=6,
|
||||
),
|
||||
"dev-controlnet-canny": ModelConfig(
|
||||
alias="dev-controlnet-canny",
|
||||
model_name="black-forest-labs/FLUX.1-dev",
|
||||
base_model=None,
|
||||
controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny",
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=True,
|
||||
@ -196,6 +222,7 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-schnell",
|
||||
base_model=None,
|
||||
controlnet_model="InstantX/FLUX.1-dev-Controlnet-Canny",
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=256,
|
||||
supports_guidance=False,
|
||||
@ -207,6 +234,7 @@ AVAILABLE_MODELS = {
|
||||
model_name="black-forest-labs/FLUX.1-dev",
|
||||
base_model=None,
|
||||
controlnet_model="jasperai/Flux.1-dev-Controlnet-Upscaler",
|
||||
custom_transformer_model=None,
|
||||
num_train_steps=1000,
|
||||
max_sequence_length=512,
|
||||
supports_guidance=False,
|
||||
|
||||
@ -108,7 +108,7 @@ class LoRALayers:
|
||||
base_path=base_path,
|
||||
)
|
||||
|
||||
if parts[1] == "single_transformer_blocks":
|
||||
elif parts[1] == "single_transformer_blocks":
|
||||
LoRALayers._handle_single_transformer_blocks(
|
||||
weights=weights,
|
||||
scale=scale,
|
||||
@ -117,8 +117,44 @@ class LoRALayers:
|
||||
base_path=base_path,
|
||||
)
|
||||
|
||||
# Handle top-level transformer components like x_embedder, context_embedder, proj_out
|
||||
elif len(parts) == 2 and parts[0] == "transformer":
|
||||
LoRALayers._handle_top_level_component(
|
||||
weights=weights,
|
||||
scale=scale,
|
||||
transformer=transformer,
|
||||
lora_layers=lora_layers,
|
||||
base_path=base_path,
|
||||
component_name=parts[1],
|
||||
)
|
||||
|
||||
return lora_layers
|
||||
|
||||
@staticmethod
|
||||
def _resolve_legacy_paths(module, module_name: str, attr_name: str, parts: list, block_idx: int) -> tuple:
|
||||
# Handle legacy naming: ff.net.2 -> ff.linear2
|
||||
if module_name == "ff" and attr_name == "net" and len(parts) >= 6 and parts[5] == "2":
|
||||
return getattr(module, "linear2"), f"transformer.transformer_blocks.{block_idx}.ff.linear2"
|
||||
|
||||
# Handle attn.to_out.0 -> attn.to_out[0] (list access)
|
||||
if module_name == "attn" and attr_name == "to_out" and len(parts) >= 6 and parts[5] == "0":
|
||||
return module.to_out[0], f"transformer.transformer_blocks.{block_idx}.attn.to_out"
|
||||
|
||||
# Default case
|
||||
original_layer = getattr(module, attr_name)
|
||||
if len(parts) == 6:
|
||||
original_layer = original_layer[0]
|
||||
return original_layer, ".".join(parts)
|
||||
|
||||
@staticmethod
|
||||
def _create_lora_layer(weights: dict, base_path: str, original_layer, scale: float):
|
||||
lora_A = weights[f"{base_path}.lora_A"]
|
||||
rank = lora_A.shape[1]
|
||||
lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
|
||||
lora_layer.lora_A = lora_A
|
||||
lora_layer.lora_B = weights[f"{base_path}.lora_B"]
|
||||
return lora_layer
|
||||
|
||||
@staticmethod
|
||||
def _handle_transformer_blocks(
|
||||
weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str
|
||||
@ -129,22 +165,19 @@ class LoRALayers:
|
||||
attr_name = parts[4]
|
||||
block = transformer.transformer_blocks[block_idx]
|
||||
module = getattr(block, module_name)
|
||||
original_layer = getattr(module, attr_name)
|
||||
|
||||
if len(parts) == 6:
|
||||
original_layer = original_layer[0] # Special case here
|
||||
# Resolve legacy naming and get the actual layer and storage path
|
||||
original_layer, storage_path = LoRALayers._resolve_legacy_paths(
|
||||
module=module,
|
||||
module_name=module_name,
|
||||
attr_name=attr_name,
|
||||
parts=parts,
|
||||
block_idx=block_idx,
|
||||
)
|
||||
|
||||
# Create LoRA layer
|
||||
lora_A = weights[f"{base_path}.lora_A"]
|
||||
rank = lora_A.shape[1]
|
||||
lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
|
||||
|
||||
# Set the weights
|
||||
lora_layer.lora_A = weights[f"{base_path}.lora_A"]
|
||||
lora_layer.lora_B = weights[f"{base_path}.lora_B"]
|
||||
|
||||
# Store the layer
|
||||
lora_layers[base_path] = lora_layer
|
||||
# Create and store LoRA layer
|
||||
lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
|
||||
lora_layers[storage_path] = lora_layer
|
||||
|
||||
@staticmethod
|
||||
def _handle_single_transformer_blocks(
|
||||
@ -153,24 +186,25 @@ class LoRALayers:
|
||||
parts = base_path.split(".")
|
||||
block_idx = int(parts[2])
|
||||
module_name = parts[3]
|
||||
|
||||
if len(parts) == 4:
|
||||
original_layer = getattr(transformer.single_transformer_blocks[block_idx], module_name)
|
||||
elif len(parts) == 5:
|
||||
else:
|
||||
attr_name = parts[4]
|
||||
block = transformer.single_transformer_blocks[block_idx]
|
||||
module = getattr(block, module_name)
|
||||
original_layer = getattr(module, attr_name)
|
||||
|
||||
# Create LoRA layer
|
||||
lora_A = weights[f"{base_path}.lora_A"]
|
||||
rank = lora_A.shape[1]
|
||||
lora_layer = LoRALinear.from_linear(linear=original_layer, r=rank, scale=scale)
|
||||
# Create and store LoRA layer
|
||||
lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
|
||||
lora_layers[base_path] = lora_layer
|
||||
|
||||
# Set the weights
|
||||
lora_layer.lora_A = weights[f"{base_path}.lora_A"]
|
||||
lora_layer.lora_B = weights[f"{base_path}.lora_B"]
|
||||
|
||||
# Store the layer
|
||||
@staticmethod
|
||||
def _handle_top_level_component(
|
||||
weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str, component_name: str
|
||||
):
|
||||
original_layer = getattr(transformer, component_name)
|
||||
lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale)
|
||||
lora_layers[base_path] = lora_layer
|
||||
|
||||
@staticmethod
|
||||
@ -210,7 +244,11 @@ class LoRALayers:
|
||||
@staticmethod
|
||||
def _set_attribute(block, key: str, val: dict, name: str):
|
||||
if block[key].get(name, False) and val.get(name, False):
|
||||
block[key][name] = val[name]
|
||||
if name == "to_out" and isinstance(block[key][name], list):
|
||||
if len(block[key][name]) > 0:
|
||||
block[key][name][0] = val[name]
|
||||
else:
|
||||
block[key][name] = val[name]
|
||||
|
||||
@staticmethod
|
||||
def _get_nested_attr(obj, attr_path):
|
||||
|
||||
@ -54,12 +54,13 @@ class FluxInitializer:
|
||||
weights = WeightHandler.load_regular_weights(
|
||||
repo_id=model_config.model_name,
|
||||
local_path=local_path,
|
||||
transformer_repo_id=model_config.custom_transformer_model,
|
||||
)
|
||||
|
||||
# 3. Initialize all models
|
||||
flux_model.vae = VAE()
|
||||
|
||||
# Use custom transformer if provided, otherwise create default
|
||||
flux_model.t5_text_encoder = T5Encoder()
|
||||
flux_model.clip_text_encoder = CLIPEncoder()
|
||||
if custom_transformer is not None:
|
||||
flux_model.transformer = custom_transformer
|
||||
else:
|
||||
@ -69,9 +70,6 @@ class FluxInitializer:
|
||||
num_single_transformer_blocks=weights.num_single_transformer_blocks(),
|
||||
)
|
||||
|
||||
flux_model.t5_text_encoder = T5Encoder()
|
||||
flux_model.clip_text_encoder = CLIPEncoder()
|
||||
|
||||
# 4. Apply weights and quantize the models
|
||||
flux_model.bits = WeightUtil.set_weights_and_quantize(
|
||||
quantize_arg=quantize,
|
||||
|
||||
@ -42,7 +42,7 @@ class MaskUtil:
|
||||
masked_image = ArrayUtil.pack_latents(latents=masked_image, height=height, width=width)
|
||||
|
||||
# 4. Resize mask and pack latents
|
||||
mask = MaskUtil._reshape_mask(the_mask=the_mask, height=height, width=width)
|
||||
mask = MaskUtil.reshape_mask(the_mask=the_mask, height=height, width=width)
|
||||
mask = ArrayUtil.pack_latents(latents=mask, height=height, width=width, num_channels_latents=64)
|
||||
|
||||
# 5. Concat the masked_image and the mask
|
||||
@ -50,7 +50,7 @@ class MaskUtil:
|
||||
return masked_image_latents
|
||||
|
||||
@staticmethod
|
||||
def _reshape_mask(the_mask: mx.array, height: int, width: int) -> mx.array:
|
||||
def reshape_mask(the_mask: mx.array, height: int, width: int) -> mx.array:
|
||||
mask = the_mask[:, 0, :, :]
|
||||
mask = mx.reshape(mask, (1, height // 8, 8, width // 8, 8))
|
||||
mask = mx.transpose(mask, (0, 2, 4, 1, 3))
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
from mflux import Config, Flux1, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -16,6 +17,10 @@ def main():
|
||||
parser.add_output_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Set default guidance value if not provided by user
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.GUIDANCE_SCALE
|
||||
|
||||
# 1. Load the model
|
||||
flux = Flux1(
|
||||
model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
|
||||
|
||||
@ -1,6 +1,7 @@
|
||||
from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -16,6 +17,10 @@ def main():
|
||||
parser.add_output_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Set default guidance value if not provided by user
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.GUIDANCE_SCALE
|
||||
|
||||
# 1. Load the model
|
||||
flux = Flux1Controlnet(
|
||||
model_config=_get_controlnet_model_config(args.model),
|
||||
|
||||
@ -2,6 +2,7 @@ from mflux import Config, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.flux_tools.depth.flux_depth import Flux1Depth
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -19,7 +20,7 @@ def main():
|
||||
|
||||
# 0. Default to a medium guidance value for depth related tasks.
|
||||
if args.guidance is None:
|
||||
args.guidance = 10
|
||||
args.guidance = ui_defaults.DEFAULT_DEPTH_GUIDANCE
|
||||
|
||||
# 1. Load the model
|
||||
flux = Flux1Depth(
|
||||
|
||||
@ -2,6 +2,7 @@ from mflux import Config, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.flux_tools.fill.flux_fill import Flux1Fill
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -19,7 +20,7 @@ def main():
|
||||
|
||||
# 0. Default to a higher guidance value for fill related tasks.
|
||||
if args.guidance is None:
|
||||
args.guidance = 30
|
||||
args.guidance = ui_defaults.DEFAULT_DEV_FILL_GUIDANCE
|
||||
|
||||
# 1. Load the model
|
||||
flux = Flux1Fill(
|
||||
|
||||
80
src/mflux/generate_in_context_catvton.py
Normal file
80
src/mflux/generate_in_context_catvton.py
Normal file
@ -0,0 +1,80 @@
|
||||
from pathlib import Path
|
||||
|
||||
from mflux import Config, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.community.in_context.flux_in_context_fill import FluxInContextFill
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
|
||||
def main():
|
||||
# 0. Parse command line arguments
|
||||
parser = CommandLineParser(description="Generate virtual try-on images using in-context learning.")
|
||||
parser.add_general_arguments()
|
||||
parser.add_model_arguments(require_model_arg=False)
|
||||
parser.add_lora_arguments()
|
||||
parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
|
||||
parser.add_catvton_arguments()
|
||||
parser.add_in_context_arguments()
|
||||
parser.add_output_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Default to a higher guidance value for fill
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.DEFAULT_DEV_FILL_GUIDANCE
|
||||
|
||||
# Set default CATVTON prompt if none provided
|
||||
if not args.prompt and not args.prompt_file:
|
||||
args.prompt = "The pair of images highlights a clothing and its styling on a model, high resolution, 4K, 8K; [IMAGE1] Detailed product shot of a clothing; [IMAGE2] The same cloth is worn by a model in a lifestyle setting."
|
||||
|
||||
# Set sensible VAE tiling split for in-context generation (side-by-side images)
|
||||
if args.vae_tiling:
|
||||
args.vae_tiling_split = "vertical"
|
||||
|
||||
# 1. Load the model
|
||||
flux = FluxInContextFill(
|
||||
model_config=ModelConfig.dev_fill_catvton(),
|
||||
quantize=args.quantize,
|
||||
local_path=args.path,
|
||||
lora_paths=args.lora_paths,
|
||||
lora_scales=args.lora_scales,
|
||||
)
|
||||
|
||||
# 2. Register callbacks
|
||||
memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
|
||||
|
||||
try:
|
||||
for seed in args.seed:
|
||||
# 3. Generate an image for each seed value
|
||||
image = flux.generate_image(
|
||||
seed=seed,
|
||||
prompt=get_effective_prompt(args),
|
||||
left_image_path=args.garment_image,
|
||||
right_image_path=args.person_image,
|
||||
config=Config(
|
||||
num_inference_steps=args.steps,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
guidance=args.guidance,
|
||||
image_path=args.person_image,
|
||||
masked_image_path=args.person_mask,
|
||||
),
|
||||
)
|
||||
|
||||
# 4. Save the image(s)
|
||||
output_path = Path(args.output.format(seed=seed))
|
||||
image.get_right_half().save(path=output_path, export_json_metadata=args.metadata)
|
||||
if args.save_full_image:
|
||||
image.save(path=output_path.with_stem(output_path.stem + "_full"))
|
||||
|
||||
except (StopImageGenerationException, PromptFileReadError) as exc:
|
||||
print(exc)
|
||||
finally:
|
||||
if memory_saver:
|
||||
print(memory_saver.memory_stats())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -2,10 +2,11 @@ from pathlib import Path
|
||||
|
||||
from mflux import Config, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.community.in_context_lora.flux_in_context_lora import Flux1InContextLoRA
|
||||
from mflux.community.in_context_lora.in_context_loras import LORA_REPO_ID, get_lora_filename
|
||||
from mflux.community.in_context.flux_in_context_dev import Flux1InContextDev
|
||||
from mflux.community.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename
|
||||
from mflux.config.model_config import ModelConfig
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -18,17 +19,20 @@ def main():
|
||||
parser.add_lora_arguments()
|
||||
parser.add_image_generator_arguments(supports_metadata_config=True)
|
||||
parser.add_image_to_image_arguments(required=True)
|
||||
parser.add_in_context_arguments()
|
||||
parser.add_output_arguments()
|
||||
parser.add_argument(
|
||||
"--save-full-image",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Additionally, save the full image containing the reference image. Useful for verifying the in-context usage of the reference image.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Set default guidance value if not provided by user
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.GUIDANCE_SCALE
|
||||
|
||||
# Set sensible VAE tiling split for in-context generation (side-by-side images)
|
||||
if args.vae_tiling:
|
||||
args.vae_tiling_split = "vertical"
|
||||
|
||||
# 1. Load the model
|
||||
flux = Flux1InContextLoRA(
|
||||
flux = Flux1InContextDev(
|
||||
model_config=ModelConfig.dev(),
|
||||
quantize=args.quantize,
|
||||
lora_names=[get_lora_filename(args.lora_style)] if args.lora_style else None,
|
||||
@ -59,6 +63,7 @@ def main():
|
||||
image.get_right_half().save(path=output_path, export_json_metadata=args.metadata)
|
||||
if args.save_full_image:
|
||||
image.save(path=output_path.with_stem(output_path.stem + "_full"))
|
||||
|
||||
except (StopImageGenerationException, PromptFileReadError) as exc:
|
||||
print(exc)
|
||||
finally:
|
||||
103
src/mflux/generate_in_context_edit.py
Normal file
103
src/mflux/generate_in_context_edit.py
Normal file
@ -0,0 +1,103 @@
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from mflux import Config, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.community.in_context.flux_in_context_fill import FluxInContextFill
|
||||
from mflux.community.in_context.utils.in_context_loras import prepare_ic_edit_loras
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
|
||||
def main():
|
||||
# 0. Parse command line arguments
|
||||
parser = CommandLineParser(description="Generate images using in-context editing.")
|
||||
parser.add_general_arguments()
|
||||
parser.add_model_arguments(require_model_arg=False)
|
||||
parser.add_lora_arguments()
|
||||
parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
|
||||
parser.add_in_context_edit_arguments()
|
||||
parser.add_in_context_arguments()
|
||||
parser.add_output_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Default to a higher guidance value for fill related tasks.
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.DEFAULT_DEV_FILL_GUIDANCE
|
||||
|
||||
# Set sensible VAE tiling split for in-context generation (side-by-side images)
|
||||
if args.vae_tiling:
|
||||
args.vae_tiling_split = "vertical"
|
||||
|
||||
# Auto-resize to optimal width for IC-Edit
|
||||
width, height = _resize_for_ic_edit_optimal_width(args)
|
||||
|
||||
# 1. Load the model with IC-Edit LoRA
|
||||
flux = FluxInContextFill(
|
||||
model_config=ModelConfig.dev_fill(),
|
||||
quantize=args.quantize,
|
||||
local_path=args.path,
|
||||
lora_paths=prepare_ic_edit_loras(args.lora_paths),
|
||||
lora_scales=args.lora_scales,
|
||||
)
|
||||
|
||||
# 2. Register callbacks
|
||||
memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
|
||||
|
||||
try:
|
||||
for seed in args.seed:
|
||||
# 3. Generate an image for each seed value
|
||||
image = flux.generate_image(
|
||||
seed=seed,
|
||||
prompt=_get_effective_ic_edit_prompt(args),
|
||||
left_image_path=args.reference_image,
|
||||
right_image_path=None,
|
||||
config=Config(
|
||||
num_inference_steps=args.steps,
|
||||
height=height,
|
||||
width=width,
|
||||
guidance=args.guidance,
|
||||
),
|
||||
)
|
||||
|
||||
# 4. Save the image(s)
|
||||
output_path = Path(args.output.format(seed=seed))
|
||||
image.get_right_half().save(path=output_path, export_json_metadata=args.metadata)
|
||||
if args.save_full_image:
|
||||
image.save(path=output_path.with_stem(output_path.stem + "_full"))
|
||||
|
||||
except (StopImageGenerationException, PromptFileReadError) as exc:
|
||||
print(exc)
|
||||
finally:
|
||||
if memory_saver:
|
||||
print(memory_saver.memory_stats())
|
||||
|
||||
|
||||
def _get_effective_ic_edit_prompt(args):
|
||||
if hasattr(args, "instruction") and args.instruction:
|
||||
return f"A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but {args.instruction}"
|
||||
else:
|
||||
return get_effective_prompt(args)
|
||||
|
||||
|
||||
def _resize_for_ic_edit_optimal_width(args):
|
||||
with Image.open(args.reference_image) as img:
|
||||
actual_width, actual_height = img.size
|
||||
aspect_ratio = actual_height / actual_width
|
||||
original_args_width = args.width
|
||||
original_args_height = args.height
|
||||
optimal_width = 512
|
||||
optimal_height = int(512 * aspect_ratio)
|
||||
optimal_height = (optimal_height // 8) * 8
|
||||
print(f"[INFO] IC-Edit LoRA trained on 512px width. Auto-resizing from actual image {actual_width}x{actual_height} to {optimal_width}x{optimal_height}") # fmt:off
|
||||
print(f"[INFO] Aspect ratio maintained: {aspect_ratio:.3f}")
|
||||
if original_args_width != actual_width or original_args_height != actual_height:
|
||||
print(f"[INFO] Note: Command line args specified {original_args_width}x{original_args_height}, but using actual image dimensions for scaling") # fmt:off
|
||||
return optimal_width, optimal_height
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -4,6 +4,7 @@ from mflux import Config, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.flux_tools.redux.flux_redux import Flux1Redux
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
|
||||
@ -19,6 +20,10 @@ def main():
|
||||
parser.add_output_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
# 0. Set default guidance value if not provided by user
|
||||
if args.guidance is None:
|
||||
args.guidance = ui_defaults.GUIDANCE_SCALE
|
||||
|
||||
# Validate and normalize redux image strengths
|
||||
redux_image_strengths = _validate_redux_image_strengths(
|
||||
redux_image_paths=args.redux_image_paths,
|
||||
|
||||
@ -5,7 +5,7 @@ import time
|
||||
import typing as t
|
||||
from pathlib import Path
|
||||
|
||||
from mflux.community.in_context_lora.in_context_loras import LORA_NAME_MAP, LORA_REPO_ID
|
||||
from mflux.community.in_context.utils.in_context_loras import LORA_NAME_MAP, LORA_REPO_ID
|
||||
from mflux.ui import (
|
||||
box_values,
|
||||
defaults as ui_defaults,
|
||||
@ -72,10 +72,10 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
self.add_argument("--height", type=int, default=ui_defaults.HEIGHT, help=f"Image height (Default is {ui_defaults.HEIGHT})")
|
||||
self.add_argument("--width", type=int, default=ui_defaults.WIDTH, help=f"Image width (Default is {ui_defaults.HEIGHT})")
|
||||
self.add_argument("--steps", type=int, default=None, help="Inference Steps")
|
||||
self.add_argument("--guidance", type=float, default=ui_defaults.GUIDANCE_SCALE, help=f"Guidance Scale (Default is {ui_defaults.GUIDANCE_SCALE})")
|
||||
self.add_argument("--guidance", type=float, default=None, help=f"Guidance Scale (Default varies by tool: {ui_defaults.GUIDANCE_SCALE} for most, {ui_defaults.DEFAULT_DEV_FILL_GUIDANCE} for fill tools, {ui_defaults.DEFAULT_DEPTH_GUIDANCE} for depth)")
|
||||
|
||||
def add_image_generator_arguments(self, supports_metadata_config=False) -> None:
|
||||
prompt_group = self.add_mutually_exclusive_group(required=(not supports_metadata_config))
|
||||
def add_image_generator_arguments(self, supports_metadata_config=False, require_prompt=True) -> None:
|
||||
prompt_group = self.add_mutually_exclusive_group(required=(require_prompt and not supports_metadata_config))
|
||||
prompt_group.add_argument("--prompt", type=str, help="The textual description of the image to generate.")
|
||||
prompt_group.add_argument("--prompt-file", type=Path, help="Path to a file containing the prompt text. The file will be re-read before each generation, allowing you to edit the prompt between iterations when using multiple seeds without restarting the program.")
|
||||
self.add_argument("--seed", type=int, default=None, nargs='+', help="Specify 1+ Entropy Seeds (Default is 1 time-based random-seed)")
|
||||
@ -83,6 +83,7 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
self._add_image_generator_common_arguments()
|
||||
if supports_metadata_config:
|
||||
self.add_metadata_config()
|
||||
self.require_prompt = require_prompt
|
||||
|
||||
def add_image_to_image_arguments(self, required=False) -> None:
|
||||
self.supports_image_to_image = True
|
||||
@ -98,6 +99,19 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
self.add_argument("--image-path", type=Path, required=True, help="Local path to the source image")
|
||||
self.add_argument("--masked-image-path", type=Path, required=True, help="Local path to the mask image")
|
||||
|
||||
def add_catvton_arguments(self) -> None:
|
||||
self.add_argument("--person-image", type=str, required=True, help="Path to person image")
|
||||
self.add_argument("--person-mask", type=str, required=True, help="Path to person mask image")
|
||||
self.add_argument("--garment-image", type=str, required=True, help="Garment Image")
|
||||
|
||||
def add_in_context_edit_arguments(self) -> None:
|
||||
self.supports_in_context_edit = True
|
||||
self.add_argument("--reference-image", type=str, required=True, help="Path to reference image")
|
||||
self.add_argument("--instruction", type=str, help="User instruction to be wrapped in diptych template (e.g., 'make the hair black'). This will be automatically formatted as 'A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but {instruction}'. Either --instruction or --prompt is required.") # fmt:off
|
||||
|
||||
def add_in_context_arguments(self) -> None:
|
||||
self.add_argument("--save-full-image", action="store_true", default=False, help="Additionally, save the full image containing the reference image. Useful for verifying the in-context usage of the reference image.")
|
||||
|
||||
def add_depth_arguments(self) -> None:
|
||||
self.add_argument("--image-path", type=Path, required=False, help="Local path to the source image")
|
||||
self.add_argument("--depth-image-path", type=Path, required=False, help="Local path to the depth image")
|
||||
@ -242,11 +256,21 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
|
||||
if self.supports_image_generation and namespace.prompt is None and namespace.prompt_file is None:
|
||||
# when metadata config is supported but neither prompt nor prompt-file is provided
|
||||
self.error("Either --prompt or --prompt-file argument is required, or 'prompt' required in metadata config file")
|
||||
# Only error if prompt is actually required
|
||||
if getattr(self, 'require_prompt', True):
|
||||
self.error("Either --prompt or --prompt-file argument is required, or 'prompt' required in metadata config file")
|
||||
|
||||
if self.supports_image_generation and namespace.steps is None:
|
||||
namespace.steps = ui_defaults.MODEL_INFERENCE_STEPS.get(namespace.model, 14)
|
||||
|
||||
# In-context edit specific validations
|
||||
if getattr(self, 'supports_in_context_edit', False):
|
||||
if not getattr(namespace, 'prompt', None) and not getattr(namespace, 'instruction', None):
|
||||
self.error("Either --prompt or --instruction argument is required for in-context editing")
|
||||
|
||||
if getattr(namespace, 'prompt', None) and getattr(namespace, 'instruction', None):
|
||||
self.error("Cannot use both --prompt and --instruction. Choose one.")
|
||||
|
||||
if self.supports_image_outpaint and namespace.image_outpaint_padding is not None:
|
||||
# parse and normalize any acceptable 1,2,3,4-tuple box value to 4-tuple
|
||||
namespace.image_outpaint_padding = box_values.parse_box_value(namespace.image_outpaint_padding)
|
||||
|
||||
@ -1,5 +1,7 @@
|
||||
BATTERY_PERCENTAGE_STOP_LIMIT = 5
|
||||
CONTROLNET_STRENGTH = 0.4
|
||||
DEFAULT_DEV_FILL_GUIDANCE = 30
|
||||
DEFAULT_DEPTH_GUIDANCE = 10
|
||||
GUIDANCE_SCALE = 3.5
|
||||
HEIGHT, WIDTH = 1024, 1024
|
||||
IMAGE_STRENGTH = 0.4
|
||||
|
||||
@ -36,14 +36,22 @@ class WeightHandler:
|
||||
def load_regular_weights(
|
||||
repo_id: str | None = None,
|
||||
local_path: str | None = None,
|
||||
transformer_repo_id: str | None = None,
|
||||
) -> "WeightHandler":
|
||||
# Load the weights from disk, huggingface cache, or download from huggingface
|
||||
root_path = Path(local_path) if local_path else WeightHandler._download_or_get_cached_weights(repo_id)
|
||||
|
||||
# Some custom models might have a different specific transformer setup
|
||||
if transformer_repo_id:
|
||||
transformer_path = WeightHandler._download_transformer_weights(transformer_repo_id)
|
||||
else:
|
||||
transformer_path = root_path
|
||||
|
||||
# Load the weights
|
||||
transformer, quantization_level, mflux_version = WeightHandler.load_transformer(root_path=transformer_path)
|
||||
clip_encoder, _, _ = WeightHandler._load_clip_encoder(root_path=root_path)
|
||||
t5_encoder, _, _ = WeightHandler._load_t5_encoder(root_path=root_path)
|
||||
vae, _, _ = WeightHandler._load_vae(root_path=root_path)
|
||||
transformer, quantization_level, mflux_version = WeightHandler.load_transformer(root_path=root_path)
|
||||
|
||||
return WeightHandler(
|
||||
clip_encoder=clip_encoder,
|
||||
@ -96,6 +104,17 @@ class WeightHandler:
|
||||
weights.pop("encoder")
|
||||
return weights, quantization_level, mflux_version
|
||||
|
||||
@staticmethod
|
||||
def _safely_extract_ff_weights(ff_dict: dict) -> dict | None:
|
||||
try:
|
||||
net = ff_dict["net"]
|
||||
return {
|
||||
"linear1": net[0]["proj"],
|
||||
"linear2": net[2],
|
||||
}
|
||||
except (KeyError, IndexError, TypeError):
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def load_transformer(root_path: Path | None = None, lora_path: str | None = None) -> tuple[dict, int, str | None]:
|
||||
weights, quantization_level, mflux_version = WeightHandler.get_weights("transformer", root_path, lora_path)
|
||||
@ -112,16 +131,18 @@ class WeightHandler:
|
||||
# Reshape and process the huggingface weights
|
||||
if "transformer_blocks" in weights:
|
||||
for block in weights["transformer_blocks"]:
|
||||
if block.get("ff") is not None:
|
||||
block["ff"] = {
|
||||
"linear1": block["ff"]["net"][0]["proj"],
|
||||
"linear2": block["ff"]["net"][2],
|
||||
}
|
||||
if block.get("ff_context") is not None:
|
||||
block["ff_context"] = {
|
||||
"linear1": block["ff_context"]["net"][0]["proj"],
|
||||
"linear2": block["ff_context"]["net"][2],
|
||||
}
|
||||
# Safely process ff weights
|
||||
if "ff" in block:
|
||||
extracted_ff = WeightHandler._safely_extract_ff_weights(block["ff"])
|
||||
if extracted_ff is not None:
|
||||
block["ff"] = extracted_ff
|
||||
|
||||
# Safely process ff_context weights
|
||||
if "ff_context" in block:
|
||||
extracted_ff_context = WeightHandler._safely_extract_ff_weights(block["ff_context"])
|
||||
if extracted_ff_context is not None:
|
||||
block["ff_context"] = extracted_ff_context
|
||||
|
||||
return weights, quantization_level, mflux_version
|
||||
|
||||
@staticmethod
|
||||
@ -141,6 +162,17 @@ class WeightHandler:
|
||||
weights["encoder"]["conv_norm_out"] = {"norm": weights["encoder"]["conv_norm_out"]}
|
||||
return weights, quantization_level, mflux_version
|
||||
|
||||
@staticmethod
|
||||
def _get_model_file_pattern(model_name: str, root_path: Path):
|
||||
if model_name == "transformer":
|
||||
nested_files = list(root_path.glob("transformer/*.safetensors"))
|
||||
if nested_files:
|
||||
return root_path.glob("transformer/*.safetensors")
|
||||
else:
|
||||
return root_path.glob("*.safetensors")
|
||||
else:
|
||||
return root_path.glob(model_name + "/*.safetensors")
|
||||
|
||||
@staticmethod
|
||||
def get_weights(
|
||||
model_name: str,
|
||||
@ -152,7 +184,8 @@ class WeightHandler:
|
||||
mflux_version = None
|
||||
|
||||
if root_path is not None:
|
||||
for file in sorted(root_path.glob(model_name + "/*.safetensors")):
|
||||
file_glob = WeightHandler._get_model_file_pattern(model_name, root_path)
|
||||
for file in sorted(file_glob):
|
||||
data = mx.load(str(file), return_metadata=True)
|
||||
weight = list(data[0].items())
|
||||
if len(data) > 1:
|
||||
@ -190,3 +223,15 @@ class WeightHandler:
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _download_transformer_weights(repo_id: str) -> Path:
|
||||
return Path(
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
allow_patterns=[
|
||||
"transformer/*.safetensors",
|
||||
"*.safetensors",
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
@ -138,6 +138,12 @@ class WeightHandlerLoRA:
|
||||
def set_lora_layers(transformer_module: nn.Module, lora_layers: LoRALayers) -> None:
|
||||
transformer = lora_layers.layers.transformer
|
||||
|
||||
# Handle top-level transformer components (x_embedder, context_embedder, proj_out, etc.)
|
||||
for attr_name in ["x_embedder", "context_embedder", "proj_out"]:
|
||||
component = transformer.get(attr_name, None)
|
||||
if component is not None:
|
||||
setattr(transformer_module, attr_name, component)
|
||||
|
||||
# Handle transformer_blocks
|
||||
transformer_blocks = transformer.get("transformer_blocks", [])
|
||||
for i, weights in enumerate(transformer_blocks):
|
||||
|
||||
@ -17,7 +17,7 @@ class WeightHandlerLoRAHuggingFace:
|
||||
lora_paths = []
|
||||
if lora_names:
|
||||
for lora_name in lora_names:
|
||||
lora_path = WeightHandlerLoRAHuggingFace._download_lora(
|
||||
lora_path = WeightHandlerLoRAHuggingFace.download_lora(
|
||||
repo_id=repo_id,
|
||||
lora_name=lora_name,
|
||||
cache_dir=cache_dir,
|
||||
@ -27,7 +27,7 @@ class WeightHandlerLoRAHuggingFace:
|
||||
return lora_paths
|
||||
|
||||
@staticmethod
|
||||
def _download_lora(
|
||||
def download_lora(
|
||||
repo_id: str,
|
||||
lora_name: str,
|
||||
cache_dir: str | None = None,
|
||||
|
||||
@ -164,6 +164,56 @@ def mflux_concept_minimal_argv() -> list[str]:
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mflux_catvton_parser() -> CommandLineParser:
|
||||
parser = CommandLineParser(description="Generate virtual try-on images using in-context learning.")
|
||||
parser.add_general_arguments()
|
||||
parser.add_model_arguments(require_model_arg=False)
|
||||
parser.add_lora_arguments()
|
||||
parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
|
||||
parser.add_catvton_arguments()
|
||||
parser.add_in_context_arguments()
|
||||
parser.add_output_arguments()
|
||||
return parser
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mflux_catvton_minimal_argv() -> list[str]:
|
||||
return [
|
||||
"mflux-generate-in-context-catvton",
|
||||
"--person-image",
|
||||
"person.png",
|
||||
"--person-mask",
|
||||
"person_mask.png",
|
||||
"--garment-image",
|
||||
"garment.png",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mflux_in_context_edit_parser() -> CommandLineParser:
|
||||
parser = CommandLineParser(description="Generate images using in-context editing.")
|
||||
parser.add_general_arguments()
|
||||
parser.add_model_arguments(require_model_arg=False)
|
||||
parser.add_lora_arguments()
|
||||
parser.add_image_generator_arguments(supports_metadata_config=False, require_prompt=False)
|
||||
parser.add_in_context_edit_arguments()
|
||||
parser.add_in_context_arguments()
|
||||
parser.add_output_arguments()
|
||||
return parser
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mflux_in_context_edit_minimal_argv() -> list[str]:
|
||||
return [
|
||||
"mflux-generate-in-context-edit",
|
||||
"--reference-image",
|
||||
"reference.png",
|
||||
"--instruction",
|
||||
"make the hair black",
|
||||
]
|
||||
|
||||
|
||||
def test_model_path_requires_model_arg(mflux_generate_parser):
|
||||
# when loading a model via --path, the model name still need to be specified
|
||||
with patch("sys.argv", "mflux-generate", "--path", "/some/saved/model"):
|
||||
@ -613,18 +663,15 @@ def test_fill_default_guidance():
|
||||
):
|
||||
args = parser.parse_args()
|
||||
|
||||
# Verify initial guidance value is the UI default
|
||||
assert args.guidance == ui_defaults.GUIDANCE_SCALE
|
||||
# Verify initial guidance value is None (no default set by parser)
|
||||
assert args.guidance is None
|
||||
|
||||
# Simulate what happens in generate_fill.py
|
||||
if args.guidance is None:
|
||||
args.guidance = 30
|
||||
else:
|
||||
# In our test we'll just override it to simulate the behavior
|
||||
args.guidance = 30
|
||||
args.guidance = ui_defaults.DEFAULT_DEV_FILL_GUIDANCE
|
||||
|
||||
# Now check that guidance is correctly set to 30
|
||||
assert args.guidance == 30
|
||||
# Now check that guidance is correctly set to the default dev fill guidance
|
||||
assert args.guidance == ui_defaults.DEFAULT_DEV_FILL_GUIDANCE
|
||||
|
||||
|
||||
def test_save_depth_args(mflux_save_depth_parser, mflux_save_depth_minimal_argv):
|
||||
@ -727,3 +774,102 @@ def test_concept_attention_args(mflux_concept_parser, mflux_concept_minimal_argv
|
||||
assert args.concept == "car"
|
||||
assert args.heatmap_layer_indices == [10, 11, 12]
|
||||
assert args.heatmap_timesteps == [0, 1, 2]
|
||||
|
||||
|
||||
def test_catvton_args(mflux_catvton_parser, mflux_catvton_minimal_argv):
|
||||
# Test required arguments
|
||||
with patch("sys.argv", mflux_catvton_minimal_argv):
|
||||
args = mflux_catvton_parser.parse_args()
|
||||
assert args.person_image == "person.png"
|
||||
assert args.person_mask == "person_mask.png"
|
||||
assert args.garment_image == "garment.png"
|
||||
|
||||
# Test prompt is None by default (set by the app, not parser)
|
||||
assert args.prompt is None
|
||||
|
||||
# Test missing required arguments
|
||||
with patch("sys.argv", ["mflux-generate-in-context-catvton"]):
|
||||
pytest.raises(SystemExit, mflux_catvton_parser.parse_args)
|
||||
|
||||
with patch("sys.argv", ["mflux-generate-in-context-catvton", "--person-image", "person.png"]):
|
||||
pytest.raises(SystemExit, mflux_catvton_parser.parse_args)
|
||||
|
||||
with patch(
|
||||
"sys.argv", ["mflux-generate-in-context-catvton", "--person-image", "person.png", "--person-mask", "mask.png"]
|
||||
):
|
||||
pytest.raises(SystemExit, mflux_catvton_parser.parse_args)
|
||||
|
||||
# Test custom prompt can be set
|
||||
with patch("sys.argv", mflux_catvton_minimal_argv + ["--prompt", "custom prompt"]):
|
||||
args = mflux_catvton_parser.parse_args()
|
||||
assert args.prompt == "custom prompt"
|
||||
|
||||
# Test VAE tiling split argument
|
||||
with patch("sys.argv", mflux_catvton_minimal_argv + ["--vae-tiling"]):
|
||||
args = mflux_catvton_parser.parse_args()
|
||||
assert args.vae_tiling is True
|
||||
assert args.vae_tiling_split == "horizontal" # Default value from parser
|
||||
|
||||
|
||||
def test_in_context_edit_args(mflux_in_context_edit_parser, mflux_in_context_edit_minimal_argv):
|
||||
# Test required arguments with instruction
|
||||
with patch("sys.argv", mflux_in_context_edit_minimal_argv):
|
||||
args = mflux_in_context_edit_parser.parse_args()
|
||||
assert args.reference_image == "reference.png"
|
||||
assert args.instruction == "make the hair black"
|
||||
assert hasattr(mflux_in_context_edit_parser, "supports_in_context_edit")
|
||||
assert mflux_in_context_edit_parser.supports_in_context_edit is True
|
||||
|
||||
# Test with prompt instead of instruction
|
||||
with patch(
|
||||
"sys.argv",
|
||||
[
|
||||
"mflux-generate-in-context-edit",
|
||||
"--reference-image",
|
||||
"reference.png",
|
||||
"--prompt",
|
||||
"A diptych with custom prompt",
|
||||
],
|
||||
):
|
||||
args = mflux_in_context_edit_parser.parse_args()
|
||||
assert args.reference_image == "reference.png"
|
||||
assert args.prompt == "A diptych with custom prompt"
|
||||
|
||||
# Test missing required arguments
|
||||
with patch("sys.argv", ["mflux-generate-in-context-edit"]):
|
||||
pytest.raises(SystemExit, mflux_in_context_edit_parser.parse_args)
|
||||
|
||||
with patch("sys.argv", ["mflux-generate-in-context-edit", "--reference-image", "reference.png"]):
|
||||
pytest.raises(SystemExit, mflux_in_context_edit_parser.parse_args)
|
||||
|
||||
# Test both prompt and instruction provided (should error)
|
||||
with patch(
|
||||
"sys.argv",
|
||||
[
|
||||
"mflux-generate-in-context-edit",
|
||||
"--reference-image",
|
||||
"reference.png",
|
||||
"--prompt",
|
||||
"test prompt",
|
||||
"--instruction",
|
||||
"test instruction",
|
||||
],
|
||||
):
|
||||
pytest.raises(SystemExit, mflux_in_context_edit_parser.parse_args)
|
||||
|
||||
# Test VAE tiling split argument
|
||||
with patch("sys.argv", mflux_in_context_edit_minimal_argv + ["--vae-tiling"]):
|
||||
args = mflux_in_context_edit_parser.parse_args()
|
||||
assert args.vae_tiling is True
|
||||
assert args.vae_tiling_split == "horizontal" # Default value from parser
|
||||
|
||||
|
||||
def test_in_context_args(mflux_catvton_parser, mflux_catvton_minimal_argv):
|
||||
# Test save_full_image flag
|
||||
with patch("sys.argv", mflux_catvton_minimal_argv):
|
||||
args = mflux_catvton_parser.parse_args()
|
||||
assert args.save_full_image is False
|
||||
|
||||
with patch("sys.argv", mflux_catvton_minimal_argv + ["--save-full-image"]):
|
||||
args = mflux_catvton_parser.parse_args()
|
||||
assert args.save_full_image is True
|
||||
|
||||
@ -5,6 +5,7 @@ from PIL import Image
|
||||
|
||||
from mflux import Config, ModelConfig
|
||||
from mflux.flux_tools.fill.flux_fill import Flux1Fill
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
|
||||
|
||||
|
||||
@ -49,7 +50,7 @@ class ImageGeneratorFillTestHelper:
|
||||
width=width,
|
||||
image_path=image_path,
|
||||
masked_image_path=masked_image_path,
|
||||
guidance=30,
|
||||
guidance=ui_defaults.DEFAULT_DEV_FILL_GUIDANCE,
|
||||
),
|
||||
)
|
||||
image.save(path=output_image_path, overwrite=True)
|
||||
|
||||
@ -0,0 +1,73 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from mflux import Config, ModelConfig
|
||||
from mflux.community.in_context.flux_in_context_fill import FluxInContextFill
|
||||
from mflux.community.in_context.utils.in_context_loras import prepare_ic_edit_loras
|
||||
|
||||
|
||||
class ImageGeneratorICEditTestHelper:
|
||||
@staticmethod
|
||||
def assert_matches_reference_image(
|
||||
reference_image_path: str,
|
||||
output_image_path: str,
|
||||
prompt: str,
|
||||
steps: int,
|
||||
seed: int,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
reference_image: str | None = None,
|
||||
lora_paths: list[str] | None = None,
|
||||
lora_scales: list[float] | None = None,
|
||||
):
|
||||
# resolve paths
|
||||
reference_image_path = ImageGeneratorICEditTestHelper.resolve_path(reference_image_path)
|
||||
output_image_path = ImageGeneratorICEditTestHelper.resolve_path(output_image_path)
|
||||
reference_image = ImageGeneratorICEditTestHelper.resolve_path(reference_image)
|
||||
lora_paths = [str(ImageGeneratorICEditTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
|
||||
|
||||
try:
|
||||
# given
|
||||
flux = FluxInContextFill(
|
||||
model_config=ModelConfig.dev_fill(),
|
||||
quantize=8,
|
||||
lora_paths=prepare_ic_edit_loras(lora_paths),
|
||||
lora_scales=lora_scales,
|
||||
)
|
||||
|
||||
# when
|
||||
image = flux.generate_image(
|
||||
seed=seed,
|
||||
prompt=prompt,
|
||||
left_image_path=str(reference_image),
|
||||
right_image_path=None,
|
||||
config=Config(
|
||||
num_inference_steps=steps,
|
||||
height=height,
|
||||
width=width,
|
||||
),
|
||||
)
|
||||
|
||||
# Save the result
|
||||
image.get_right_half().save(path=output_image_path, overwrite=True)
|
||||
|
||||
# then
|
||||
np.testing.assert_array_equal(
|
||||
np.array(Image.open(output_image_path)),
|
||||
np.array(Image.open(reference_image_path)),
|
||||
err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}",
|
||||
)
|
||||
|
||||
finally:
|
||||
# cleanup
|
||||
if os.path.exists(output_image_path):
|
||||
os.remove(output_image_path)
|
||||
|
||||
@staticmethod
|
||||
def resolve_path(path) -> Path | None:
|
||||
if path is None:
|
||||
return None
|
||||
return Path(__file__).parent.parent.parent / "resources" / path
|
||||
@ -5,8 +5,8 @@ import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from mflux import Config, ModelConfig
|
||||
from mflux.community.in_context_lora.flux_in_context_lora import Flux1InContextLoRA
|
||||
from mflux.community.in_context_lora.in_context_loras import LORA_REPO_ID, get_lora_filename
|
||||
from mflux.community.in_context.flux_in_context_dev import Flux1InContextDev
|
||||
from mflux.community.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename
|
||||
|
||||
|
||||
class ImageGeneratorInContextTestHelper:
|
||||
@ -35,7 +35,7 @@ class ImageGeneratorInContextTestHelper:
|
||||
|
||||
try:
|
||||
# given
|
||||
flux = Flux1InContextLoRA(
|
||||
flux = Flux1InContextDev(
|
||||
model_config=model_config,
|
||||
quantize=8,
|
||||
lora_names=[get_lora_filename(lora_style)] if lora_style else None,
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
from mflux import ModelConfig
|
||||
from tests.image_generation.helpers.image_generation_ic_edit_test_helper import ImageGeneratorICEditTestHelper
|
||||
from tests.image_generation.helpers.image_generation_in_context_test_helper import ImageGeneratorInContextTestHelper
|
||||
|
||||
|
||||
@ -16,3 +17,15 @@ class TestImageGeneratorInContext:
|
||||
image_path="ic_lora_reference_logo.png",
|
||||
lora_style="identity",
|
||||
)
|
||||
|
||||
def test_ic_edit_with_instruction(self):
|
||||
ImageGeneratorICEditTestHelper.assert_matches_reference_image(
|
||||
reference_image_path="ic_edit_reference.png",
|
||||
output_image_path="output_ic_edit.png",
|
||||
steps=20,
|
||||
seed=2799665,
|
||||
height=224,
|
||||
width=512,
|
||||
prompt="A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but remove the cheese",
|
||||
reference_image="reference_redux_dev.png",
|
||||
)
|
||||
|
||||
BIN
tests/resources/ic_edit_reference.png
Normal file
BIN
tests/resources/ic_edit_reference.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 168 KiB |
Loading…
Reference in New Issue
Block a user