diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..9d76e5c --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,37 @@ +name: Fast Tests + +on: + pull_request: + branches: [main] + push: + branches: [main] + +jobs: + fast-tests: + runs-on: macos-14 # Apple Silicon runner (M1) + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + version: "latest" + + - name: Install dependencies + run: | + uv venv + uv pip install -e . + uv pip install pytest + + - name: Run fast tests + run: | + source .venv/bin/activate + python -m pytest -m fast -v + diff --git a/CHANGELOG.md b/CHANGELOG.md index 82b4394..700d3c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,109 @@ 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). +## [0.13.0] - 2025-12-03 + +# MFLUX v.0.13.0 Release Notes + +### ๐ŸŽจ New Model Support + +- **Z-Image Turbo Support**: Added support for [Z-Image Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo), a fast distilled Z-Image variant optimized for speed +- **New command**: `mflux-generate-z-image-turbo` for rapid image generation (with LoRA support, img2img, and quantization) + +### โœจ New Features + +- **FIBO VLM Quantization Support**: The FIBO VLM commands (`mflux-fibo-inspire`, `mflux-fibo-refine`) now support quantization via the `-q` flag (3, 4, 5, 6, or 8-bit) + +- **Unified `--model` argument**: The `--model` flag now accepts local paths, HuggingFace repos, or predefined model names + - Local paths: `--model /Users/me/models/fibo-4bit` or `--model ~/my-model` + - HuggingFace repos: `--model briaai/Fibo-mlx-4bit` + - Predefined names: `--model dev`, `--model schnell`, `--model fibo` + - This mirrors how LoRA paths work for a consistent UX + +- **Scale Factor Dimensions for Img2Img**: Generalized the scale factor feature (e.g., `2x`, `0.5x`, `auto`) from upscaling to all img2img commands + - Specify output dimensions relative to input image: `--width 2x --height 2x` + - Use `auto` to match input image dimensions: `--width auto --height auto` + - Mix scale factors with absolute values: `--width 2x --height 512` + - Supported in: `mflux-generate`, `mflux-generate-z-image-turbo`, `mflux-generate-fibo`, `mflux-generate-kontext`, `mflux-generate-qwen` +- **DimensionResolver utility**: New `DimensionResolver.resolve()` for consistent dimension handling across commands + +### ๐Ÿ”ง Architecture Improvements + +- **Unified Resolution System**: New `resolution/` module for consistent parameter resolution across all models + - `PathResolution`: Resolves model paths from local paths, HuggingFace repos, or predefined names + - `LoRAResolution`: Handles LoRA path resolution from all supported formats + - `ConfigResolution`: Centralizes configuration resolution logic + - `QuantizationResolution`: Determines quantization from saved models or CLI args +- **Unified Weight Loading System**: Complete rewrite of weight handling with declarative mappings + - New `WeightLoader` with single `load(model_path)` interface + - `WeightDefinition` classes define model structure per model family + - `WeightMapping` declarative mappings replace imperative weight handlers + - Removed all per-model `weight_handler_*.py` files in favor of unified system +- **Unified Tokenizer System**: New common tokenizer module + - `TokenizerLoader.load_all()` with unified `model_path` interface + - Removed model-specific tokenizer handlers (`clip_tokenizer.py`, `t5_tokenizer.py`, etc.) +- **Unified LoRA API**: Simplified LoRA loading to a single `lora_paths` parameter + - All LoRA formats now resolved through `LoRALibrary.resolve_paths()`: + - Local paths: `/path/to/lora.safetensors` + - Registry names: `my-lora` (from `LORA_LIBRARY_PATH`) + - HuggingFace repos: `author/model` + - **New**: HuggingFace collections: `repo_id:filename.safetensors` + - Simplified model initialization: just pass `lora_paths` and everything resolves automatically +- **Unified Latent Creator Interface**: Standardized `unpack_latents(latents, height, width)` signature across all model families + - `FluxLatentCreator`, `ZImageLatentCreator`, `FiboLatentCreator`, and `QwenLatentCreator` now share the same interface + - Moved `FIBO._unpack_latents` to `FiboLatentCreator.unpack_latents` for consistency +- **StepwiseHandler Refactor**: Fixed `StepwiseHandler` to work with all model types by accepting a `latent_creator` parameter + - Previously hardcoded to `FluxLatentCreator`, now model-agnostic + - Each command passes its appropriate latent creator to `CallbackManager.register_callbacks()` +- **CLI Reorganization**: Moved CLI entry points to model-specific directories (e.g., `mflux/models/flux/cli/`) + +### ๐Ÿ”„ Breaking Changes + +- **Simplified `generate_image()` API** (programmatic users only): + - Removed `Config` class - parameters are now passed directly to `generate_image()` + - Removed `RuntimeConfig` class - internal complexity eliminated + - Added `Flux1` export to main `mflux` module for cleaner imports +- **LoRA API simplified** (programmatic users only): + - Removed `lora_names` and `lora_repo_id` parameters from all model classes (`Flux1`, `QwenImage`, `QwenImageEdit`, etc.) + - Removed `--lora-name` and `--lora-repo-id` CLI arguments + - Removed `LoRAHuggingFaceDownloader` class + +### ๐Ÿ”„ Breaking Changes (CLI) + +- **`--path` flag removed**: The deprecated `--path` flag for loading models has been removed. Use `--model` instead for local paths, HuggingFace repos, or predefined model names. + +### ๐Ÿ› Bug Fixes + +- **`--model` flag not working**: Fixed bug where the `--model` argument wasn't being used for loading models from HuggingFace or local paths. All CLI commands now correctly use `--model` for model path resolution. +- **Model Saving Index File**: Fixed issue where locally saved models (via `mflux-save`) would fail to load when uploaded to HuggingFace, due to missing `model.safetensors.index.json`. The model saver now generates this index file alongside the safetensor shards, ensuring compatibility with both mflux and standard HuggingFace loading paths. (see [#285](https://github.com/filipstrand/mflux/issues/285)) + +### ๐Ÿงช Test Infrastructure + +- **Test markers**: Added `fast` and `slow` pytest markers to categorize tests + - Fast tests: Unit tests that don't generate images (parsers, schedulers, resolution, utilities) + - Slow tests: Integration tests that generate actual images and compare to references +- **New Makefile targets**: + - `make test-fast` - Run fast tests only (quick feedback during development) + - `make test-slow` - Run slow tests only (image generation tests) + - `make test` - Run all tests (unchanged) +- Run specific test categories: `pytest -m fast` or `pytest -m slow` +- **GitHub Actions CI**: Fast tests now run automatically on PRs and pushes to main + +### ๐Ÿ”ง Internal Changes + +- Simplified `WeightLoader.load()` to take a single `model_path` parameter instead of separate `repo_id` and `local_path` +- Simplified `TokenizerLoader.load_all()` with the same unified `model_path` interface +- Renamed `local_path` parameter to `model_path` in all model constructors for clarity +- Removed `quantization_util.py` - quantization now handled through `QuantizationResolution` +- Removed `lora_huggingface_downloader.py` - downloading integrated into `LoRAResolution` +- Added comprehensive test coverage for resolution modules + +### ๐Ÿ‘ฉโ€๐Ÿ’ป Contributors + +- **Filip Strand (@filipstrand)**: Z-Image Turbo support, architecture improvements, core development + +--- + ## [0.12.1] - 2025-11-27 ### ๐Ÿ› Bug Fixes diff --git a/Makefile b/Makefile index be58b5c..6c63409 100644 --- a/Makefile +++ b/Makefile @@ -97,6 +97,21 @@ test: ensure-pytest $(PYTHON) -m pytest # โœ… Tests completed +# Run fast tests only (no image generation) +.PHONY: test-fast +test-fast: ensure-pytest + # ๐Ÿ—๏ธ Running fast tests (no image generation)... + $(PYTHON) -m pytest -m fast + # โœ… Fast tests completed + +# Run slow tests only (image generation tests) +.PHONY: test-slow +test-slow: ensure-pytest + # ๐Ÿ—๏ธ Running slow tests (image generation)... + uv pip install mlx==0.29.2 # Install pinned MLX version specifically for testing + $(PYTHON) -m pytest -m slow + # โœ… Slow tests completed + # Run uv build and check dist sizes for optimized user installs .PHONY: build @@ -129,7 +144,9 @@ help: @echo " make lint - Run ruff python linter" @echo " make format - Run ruff code formatter" @echo " make check - Run linters auto fixes *and* style formatter via pre-commit hook" - @echo " make test - Run tests" + @echo " make test - Run all tests" + @echo " make test-fast - Run fast tests only (no image generation)" + @echo " make test-slow - Run slow tests only (image generation)" @echo " make build - Build distribution packages and check sizes" @echo " make clean - Remove the virtual environment" @echo " make help - Show this help message" diff --git a/README.md b/README.md index df454d6..0ea35fb 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ ### About -Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux), [Qwen Image](https://github.com/QwenLM/Qwen-Image) and [FIBO](https://huggingface.co/briaai/FIBO) models locally on your Mac! +Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux), [Qwen Image](https://github.com/QwenLM/Qwen-Image), [FIBO](https://huggingface.co/briaai/FIBO), and [Z-Image](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) models locally on your Mac! ### Table of contents @@ -17,13 +17,14 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux), [Qwen Image](http - [โฑ๏ธ Image generation speed (updated)](#%EF%B8%8F-image-generation-speed-updated) - [โ†”๏ธ Equivalent to Diffusers implementation](#%EF%B8%8F-equivalent-to-diffusers-implementation) - [๐Ÿ—œ๏ธ Quantization](#%EF%B8%8F-quantization) -- [๐Ÿ’ฝ Running a non-quantized model directly from disk](#-running-a-non-quantized-model-directly-from-disk) +- [๐Ÿ’ฝ Running a model directly from disk](#-running-a-model-directly-from-disk) - [๐ŸŒ Third-Party HuggingFace Model Support](#-third-party-huggingface-model-support) - [๐ŸŽจ Image-to-Image](#-image-to-image) - [๐Ÿฆ™ Qwen Models](#-qwen-models) * [๐Ÿ–ผ๏ธ Qwen Image](#%EF%B8%8F-qwen-image) * [โœ๏ธ Qwen Image Edit](#%EF%B8%8F-qwen-image-edit) - [๐ŸŒ€ FIBO](#-fibo) +- [โšก Z-Image](#-z-image) - [๐Ÿ”Œ LoRA](#-lora) - [๐ŸŽญ In-Context Generation](#-in-context-generation) * [๐Ÿ“ธ Kontext](#-kontext) @@ -51,7 +52,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux), [Qwen Image](http ### Philosophy -MFLUX is a line-by-line port of the FLUX, Qwen and Bria models implementations in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) and [Huggingface Transformers](https://github.com/huggingface/transformers) libraries to [Apple MLX](https://github.com/ml-explore/mlx). +MFLUX is a line-by-line port of the FLUX, Qwen, Bria and Z-Image models implementations in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) and [Huggingface Transformers](https://github.com/huggingface/transformers) libraries to [Apple MLX](https://github.com/ml-explore/mlx). MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models (thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#%EF%B8%8F-image-generation-speed-updated), [and even faster quantized](#%EF%B8%8F-quantization). @@ -171,8 +172,7 @@ This is useful for integrating MFLUX into shell scripts or dynamically generatin Alternatively, you can use MFLUX directly in Python: ```python -from mflux.flux.flux import Flux1 -from mflux.config.config import Config +from mflux import Flux1 # Load the model flux = Flux1.from_name( @@ -184,11 +184,9 @@ flux = Flux1.from_name( image = flux.generate_image( seed=2, prompt="Luxury food photograph", - config=Config( - num_inference_steps=2, # "schnell" works well with 2-4 steps, "dev" and "krea-dev" work well with 20-25 steps - height=1024, - width=1024, - ) + num_inference_steps=2, # "schnell" works well with 2-4 steps, "dev" and "krea-dev" work well with 20-25 steps + height=1024, + width=1024, ) image.save(path="image.png") @@ -200,7 +198,7 @@ For more advanced Python usage and additional configuration options, you can exp *By default, mflux caches files in `~/Library/Caches/mflux/`. The Hugging Face model files themselves are cached separately in the Hugging Face cache directory (e.g., `~/.cache/huggingface/`).* -*To change the mflux cache location, set the `MFLUX_CACHE_DIR` environment variable. To change the Hugging Face cache location, you can modify the `HF_HOME` environment variable. For more details on Hugging Face cache settings, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*. +*To change the mflux cache location, set the `MFLUX_CACHE_DIR` environment variable. To change the Hugging Face cache location, you can modify the `HF_HOME` environment variable (e.g. to `HF_HOME=/Volumes/T7/.cache/huggingface`). For more details on Hugging Face cache settings, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*. ๐Ÿ”’ [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) ๐Ÿ”’ @@ -236,9 +234,12 @@ mflux-generate \ - **`--prompt`** (required, `str`): Text description of the image to generate. Use `-` to read the prompt from stdin (e.g., `echo "A beautiful sunset" | mflux-generate --prompt -`). -- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official Flux models (`"schnell"`, `"dev"`, or `"krea-dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`). For Qwen models, use `mflux-generate-qwen` instead. +- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Accepts: + - Predefined names: `"schnell"`, `"dev"`, `"krea-dev"`, `"fibo"`, `"z-image-turbo"` + - HuggingFace repos: `"Freepik/flux.1-lite-8B-alpha"`, `"briaai/Fibo-mlx-4bit"` + - Local paths: `"/Users/me/models/my-model"`, `"~/my-model"` -- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"`, `"dev"`, or `"krea-dev"`). Required when using third-party models from HuggingFace. +- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"`, `"dev"`, or `"krea-dev"`). Required when using third-party models from HuggingFace or local paths. - **`--output`** (optional, `str`, default: `"image.png"`): Output image filename. If `--seed` or `--auto-seeds` establishes multiple seed values, the output filename will automatically be modified to include the seed value (e.g., `image_seed_42.png`). @@ -254,11 +255,15 @@ mflux-generate \ - **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale (only used for `"dev"` and `"krea-dev"` models). -- **`--path`** (optional, `str`, default: `None`): Path to a local model on disk. +- **`--path`** (optional, `str`, default: `None`): **[DEPRECATED: use `--model` instead]** Path to a local model on disk. - **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `3`, `4`, `5`, `6`, or `8` bits). -- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#-LoRA) weights. +- **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#-LoRA) weights. Supports multiple formats: + - Local files: `/path/to/lora.safetensors` + - HuggingFace repos: `author/model` (auto-downloads) + - HuggingFace collections: `repo_id:filename.safetensors` (downloads specific file) + - Registry names: `my-lora` (via `LORA_LIBRARY_PATH`) - **`--lora-scales`** (optional, `[float]`, default: `None`): The scale for each respective [LoRA](#-LoRA) (will default to `1.0` if not specified and only one LoRA weight is loaded.) @@ -274,10 +279,6 @@ mflux-generate \ - **`--battery-percentage-stop-limit`** or **`-B`** (optional, `int`, default: `5`): On Mac laptops powered by battery, automatically stops image generation when battery percentage reaches this threshold. Prevents your Mac from shutting down and becoming unresponsive during long generation sessions. -- **`--lora-name`** (optional, `str`, default: `None`): The name of the LoRA to download from Hugging Face. - -- **`--lora-repo-id`** (optional, `str`, default: `"ali-vilab/In-Context-LoRA"`): The Hugging Face repository ID for LoRAs. - - **`--stepwise-image-output-dir`** (optional, `str`, default: `None`): [EXPERIMENTAL] Output directory to write step-wise images and their final composite image to. This feature may change in future versions. When specified, MFLUX will save an image for each denoising step, allowing you to visualize the generation process from noise to final image. - **`--vae-tiling`** (optional, flag): Enable VAE tiling to reduce memory usage during the decoding phase. This splits the image into smaller chunks for processing, which can prevent out-of-memory errors when generating high-resolution images. Note that this optimization may occasionally produce a subtle seam in the middle of the image, but it's often worth the tradeoff for being able to generate images that would otherwise cause your system to run out of memory. @@ -701,7 +702,7 @@ However, if we were to import a fixed instance of this latent array saved from t The images below illustrate this equivalence. In all cases the Schnell model was run for 2 time steps. The Diffusers implementation ran in CPU mode. -The precision for MFLUX can be set in the [Config](src/mflux/config/config.py) class. +The precision for MFLUX can be set in the [Config](src/mflux/models/common/config/config.py) class. There is typically a noticeable but very small difference in the final image when switching between 16bit and 32bit precision. --- @@ -837,6 +838,9 @@ In other words, you can reclaim the 34GB diskspace (per model) by deleting the f - [akx/FLUX.1-Kontext-dev-mflux-4bit](https://huggingface.co/akx/FLUX.1-Kontext-dev-mflux-4bit) - [filipstrand/FLUX.1-Krea-dev-mflux-4bit](https://huggingface.co/filipstrand/FLUX.1-Krea-dev-mflux-4bit) - [filipstrand/Qwen-Image-mflux-6bit](https://huggingface.co/filipstrand/Qwen-Image-mflux-6bit) + - [filipstrand/Z-Image-Turbo-mflux-4bit](https://huggingface.co/filipstrand/Z-Image-Turbo-mflux-4bit) + - [briaai/Fibo-mlx-4bit](https://huggingface.co/briaai/Fibo-mlx-4bit) + - [briaai/Fibo-mlx-8bit](https://huggingface.co/briaai/Fibo-mlx-8bit) Using the [community model support](#-third-party-huggingface-model-support), the quantized weights can be also be automatically downloaded when running the generate command: @@ -850,23 +854,33 @@ mflux-generate \ --seed 2674888 ``` +```sh +mflux-generate-fibo \ + --model briaai/Fibo-mlx-4bit \ + --prompt-file ~/Desktop/bird.json \ + --width 1024 \ + --height 1024 \ + --steps 20 \ + --seed 42 \ +``` + --- -### ๐Ÿ’ฝ Running a non-quantized model directly from disk +### ๐Ÿ’ฝ Running a model directly from disk -MFLUX also supports running a non-quantized model directly from a custom location. +MFLUX supports running a model directly from a custom location using the `--model` flag with a local path. In the example below, the model is placed in `/Users/filipstrand/Desktop/schnell`: ```sh mflux-generate \ - --path "/Users/filipstrand/Desktop/schnell" \ - --model schnell \ + --model "/Users/filipstrand/Desktop/schnell" \ + --base-model schnell \ --steps 2 \ --seed 2 \ --prompt "Luxury food photograph" ``` -Note that the `--model` flag must be set when loading a model from disk. +When loading from a local path, use `--base-model` to specify the architecture (e.g., `schnell`, `dev`). Also note that unlike when using the typical `alias` way of initializing the model (which internally handles that the required resources are downloaded), when loading a model directly from disk, we require the downloaded models to look like the following: @@ -909,17 +923,30 @@ processed a bit differently, which is why we require this structure above.* ### ๐ŸŒ Third-Party HuggingFace Model Support -MFLUX now supports compatible third-party models from HuggingFace that follow the FLUX architecture. This opens up the ecosystem to community-created models that may offer different capabilities, sizes, or specializations. +MFLUX supports compatible third-party models from HuggingFace that follow the FLUX architecture. The `--model` parameter accepts: -To use a third-party model, specify the HuggingFace repository ID with the `--model` parameter and indicate which base architecture (dev or schnell) it's derived from using the `--base-model` parameter: +- **Predefined names**: `dev`, `schnell`, `fibo`, `z-image-turbo`, etc. +- **HuggingFace repos**: `Freepik/flux.1-lite-8B`, `briaai/Fibo-mlx-4bit` +- **Local paths**: `/Users/me/models/my-model`, `~/my-model` + +This unified interface mirrors how LoRA paths work, making it easy to switch between local and remote models. ```sh +# Using a HuggingFace repo mflux-generate \ --model Freepik/flux.1-lite-8B \ --base-model schnell \ --steps 4 \ --seed 42 \ --prompt "A beautiful landscape with mountains and a lake" + +# Using a local path +mflux-generate \ + --model /Users/me/models/flux-lite \ + --base-model schnell \ + --steps 4 \ + --seed 42 \ + --prompt "A beautiful landscape with mountains and a lake" ``` Some examples of compatible third-party models include: @@ -1426,6 +1453,36 @@ Note: The optional `--prompt` parameter allows you to guide the VLM's interpreta --- +### โšก Z-Image + +MFLUX supports [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) from Tongyi Lab (Alibaba), released in November 2025. Z-Image is an efficient 6B-parameter image generation model with a single-stream DiT architecture. Z-Image-Turbo delivers high-quality images in just 9 steps, making it one of the fastest open-source models available. All the standard modes such as, img2img, LoRA and quantizations are supported for this model. See the [technical paper](https://arxiv.org/abs/2511.22699) for more details. + +![Z-Image-Turbo Example](src/mflux/assets/z_image_turbo_example.jpg) + +#### Example + +The following uses the pre-quantized 4-bit model from [filipstrand/Z-Image-Turbo-mflux-4bit](https://huggingface.co/filipstrand/Z-Image-Turbo-mflux-4bit) to generate a vibrant 1960s style image with a LoRA adapter [Technically Color](https://huggingface.co/renderartist/Technically-Color-Z-Image-Turbo) for enhanced film color: + +```sh +mflux-generate-z-image-turbo \ + --model filipstrand/Z-Image-Turbo-mflux-4bit \ + --prompt "t3chnic4lly vibrant 1960s close-up of a woman sitting under a tree in a blue skirt and white blouse, she has blonde wavy short hair and a smile with green eyes lake scene by a garden with flowers in the foreground 1960s style film She's holding her hand out there is a small smooth frog in her palm, she's making eye contact with the toad." \ + --width 1280 \ + --height 720 \ + --seed 456 \ + --steps 9 \ + --lora-paths renderartist/Technically-Color-Z-Image-Turbo \ + --lora-scales 0.5 +``` + +โš ๏ธ *Note: Z-Image-Turbo requires downloading the `Tongyi-MAI/Z-Image-Turbo` model weights (~31GB), or use quantization for smaller sizes.* + +*Dreambooth fine-tuning for Z-Image is not yet supported in MFLUX but is planned. In the meantime, you can train Z-Image-Turbo LoRAs using [AI Toolkit](https://github.com/ostris/ai-toolkit) - see [How to Train a Z-Image-Turbo LoRA with AI Toolkit](https://www.youtube.com/watch?v=Kmve1_jiDpQ) by Ostris AI.* + +*For a Swift MLX implementation of Z-Image, see [zimage.swift](https://github.com/mzbac/zimage.swift) by [@mzbac](https://github.com/mzbac).* + +--- + ### ๐Ÿ”Œ LoRA MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming). @@ -1469,6 +1526,28 @@ mflux-generate \ Just to see the difference, this image displays the four cases: One of having both adapters fully active, partially active and no LoRA at all. The example above also show the usage of `--lora-scales` flag. +#### HuggingFace LoRA Downloads + +MFLUX can automatically download LoRAs directly from HuggingFace. Simply pass the repository ID to `--lora-paths`: + +```sh +# Download from a HuggingFace repo (auto-finds the .safetensors file) +mflux-generate \ + --prompt "a portrait" \ + --lora-paths "author/lora-model" +``` + +For repositories with multiple LoRA files (collections), use the `repo_id:filename` format to specify which file to download: + +```sh +# Download a specific file from a collection +mflux-generate \ + --prompt "film storyboard style, a cat" \ + --lora-paths "ali-vilab/In-Context-LoRA:film-storyboard.safetensors" +``` + +Downloaded LoRAs are cached locally and reused on subsequent runs. + #### LoRA Library Path MFLUX supports a convenient LoRA library feature that allows you to reference LoRA files by their basename instead of full paths. This is particularly useful when you have a collection of LoRA files organized in one or more directories. diff --git a/pyproject.toml b/pyproject.toml index 5bb71d6..a7f06b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,21 +3,18 @@ requires = ["uv_build>=0.7.19,<0.8.0"] build-backend = "uv_build" [tool.uv.build-backend] -default-excludes = true # __pycache__, *.pyc, and *.pyo +default-excludes = true module-name = "mflux" namespace = true source-exclude = [ - # documentation assets were 27MB on 2025-07-05 "**/assets/**", - # dreambooth examples ~=5MB on 2025-07-05 "**/models/flux/variants/dreambooth/_example/images/**", - # loss pdf/tex does not need to be distributed "**/optimization/_loss_derivation/**", ] [project] name = "mflux" -version = "0.12.1" +version = "0.13.0.dev0" description = "A MLX port of FLUX based on the Huggingface Diffusers implementation." readme = "README.md" keywords = ["diffusers", "flux", "mlx"] @@ -42,9 +39,8 @@ dependencies = [ "requests>=2.32.4", "safetensors>=0.4.4,<1.0", "sentencepiece>=0.2.0,<1.0; python_version<'3.13'", - # sentencepiece 0.2.1 is first release with 3.13 and 3.14 wheels: https://pypi.org/project/sentencepiece/0.2.1/ "sentencepiece>=0.2.1,<1.0; python_version>='3.13'", - "tokenizers>=0.20.3; python_version>='3.13'", # transformers -> tokenizers + "tokenizers>=0.20.3; python_version>='3.13'", "toml>=0.10.2,<1.0", "torch>=2.7.1", "torch>=2.3.1,<3.0; python_version<'3.13'", @@ -68,36 +64,37 @@ dev = [ "matplotlib>3.10,<4.0", "pytest>=8.3.0,<9.0", "pytest-timer>=1.0,<2.0", - "mlx==0.29.2", # Used ONLY during test runs to ensure deterministic test results + "mlx==0.29.2", ] [project.urls] 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_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" -mflux-generate-kontext = "mflux.generate_kontext:main" -mflux-generate-qwen = "mflux.generate_qwen:main" -mflux-generate-qwen-edit = "mflux.generate_qwen_edit:main" -mflux-generate-fibo = "mflux.generate_fibo:main" -mflux-refine-fibo = "mflux.refine_fibo:main" -mflux-inspire-fibo = "mflux.inspire_fibo:main" -mflux-concept = "mflux.concept:main" -mflux-concept-from-image = "mflux.concept_from_image:main" -mflux-save = "mflux.save:main" -mflux-save-depth = "mflux.save_depth:main" -mflux-train = "mflux.train:main" -mflux-upscale = "mflux.upscale:main" -mflux-lora-library = "mflux.lora_library:main" -mflux-info = "mflux.info:main" -mflux-completions = "mflux.ui.cli.completions.install:main" +mflux-generate = "mflux.models.flux.cli.flux_generate:main" +mflux-generate-controlnet = "mflux.models.flux.cli.flux_generate_controlnet:main" +mflux-generate-in-context = "mflux.models.flux.cli.flux_generate_in_context_dev:main" +mflux-generate-in-context-edit = "mflux.models.flux.cli.flux_generate_in_context_edit:main" +mflux-generate-in-context-catvton = "mflux.models.flux.cli.flux_generate_in_context_catvton:main" +mflux-generate-fill = "mflux.models.flux.cli.flux_generate_fill:main" +mflux-generate-depth = "mflux.models.flux.cli.flux_generate_depth:main" +mflux-generate-redux = "mflux.models.flux.cli.flux_generate_redux:main" +mflux-generate-kontext = "mflux.models.flux.cli.flux_generate_kontext:main" +mflux-generate-qwen = "mflux.models.qwen.cli.qwen_image_generate:main" +mflux-generate-qwen-edit = "mflux.models.qwen.cli.qwen_image_edit_generate:main" +mflux-generate-fibo = "mflux.models.fibo.cli.fibo_generate:main" +mflux-generate-z-image-turbo = "mflux.models.z_image.cli.z_image_turbo_generate:main" +mflux-refine-fibo = "mflux.models.fibo_vlm.cli.fibo_refine:main" +mflux-inspire-fibo = "mflux.models.fibo_vlm.cli.fibo_inspire:main" +mflux-concept = "mflux.models.flux.cli.flux_concept:main" +mflux-concept-from-image = "mflux.models.flux.cli.flux_concept_from_image:main" +mflux-save = "mflux.models.common.cli.save:main" +mflux-save-depth = "mflux.models.depth_pro.cli.save_depth:main" +mflux-train = "mflux.models.common.cli.train:main" +mflux-upscale = "mflux.models.flux.cli.flux_upscale:main" +mflux-lora-library = "mflux.models.common.cli.lora_library:main" +mflux-info = "mflux.models.common.cli.info:main" +mflux-completions = "mflux.cli.completions.install:main" [tool.ruff] @@ -107,44 +104,18 @@ target-version = "py310" respect-gitignore = true [tool.ruff.lint] -# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`) codes by default. -# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or -# McCabe complexity (`C901`) by default. select = ["BLE", "E4", "E7", "E9", "F", "I", "ICN", "LOG", "PERF", "W"] ignore = [] - -# Allow fix for all enabled rules (when `--fix`) is provided. fixable = ["ALL"] unfixable = [] - -# Allow unused variables when underscore-prefixed. dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$" [tool.ruff.format] -# Like Black, use double quotes for strings. quote-style = "double" - -# Like Black, indent with spaces, rather than tabs. indent-style = "space" - -# Like Black, respect magic trailing commas. skip-magic-trailing-comma = false - -# Like Black, automatically detect the appropriate line ending. line-ending = "auto" - -# Enable auto-formatting of code examples in docstrings. Markdown, -# reStructuredText code/literal blocks and doctests are all supported. -# -# This is currently disabled by default, but it is planned for this -# to be opt-out in the future. docstring-code-format = false - -# Set the line length limit used when formatting code snippets in -# docstrings. -# -# This only has an effect when the `docstring-code-format` setting is -# enabled. docstring-code-line-length = "dynamic" [tool.pytest.ini_options] @@ -155,7 +126,6 @@ markers = [ "high_memory_requirement: marks tests that require high memory (deselect with '-m \"not high_memory_requirement\"')", ] -# https://docs.astral.sh/ruff/settings/#lintisort [tool.ruff.lint.isort] case-sensitive = false combine-as-imports = true @@ -173,7 +143,6 @@ section-order = [ [tool.mypy] disable_error_code = [ - # TODO: for each error code - clean them up in a single PR then remove entry "annotation-unchecked", "arg-type", "assignment", @@ -188,5 +157,4 @@ disable_error_code = [ error_summary = true ignore_missing_imports = true implicit_optional = true -# we should feel free to use new hinting features even if we back support to earlier versions python_version = "3.12" diff --git a/src/mflux/__init__.py b/src/mflux/__init__.py index c72ee76..9795341 100644 --- a/src/mflux/__init__.py +++ b/src/mflux/__init__.py @@ -4,3 +4,8 @@ import os # This must be set before any tokenizers are imported/used if "TOKENIZERS_PARALLELISM" not in os.environ: os.environ["TOKENIZERS_PARALLELISM"] = "false" + +# Export main classes for convenient import +from mflux.models.flux.variants.txt2img.flux import Flux1 + +__all__ = ["Flux1"] diff --git a/src/mflux/assets/z_image_turbo_example.jpg b/src/mflux/assets/z_image_turbo_example.jpg new file mode 100644 index 0000000..9e91b15 Binary files /dev/null and b/src/mflux/assets/z_image_turbo_example.jpg differ diff --git a/src/mflux/callbacks/callback.py b/src/mflux/callbacks/callback.py index 4fe5289..d599eda 100644 --- a/src/mflux/callbacks/callback.py +++ b/src/mflux/callbacks/callback.py @@ -4,7 +4,7 @@ import mlx.core as mx import PIL.Image import tqdm -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config class BeforeLoopCallback(Protocol): @@ -13,7 +13,7 @@ class BeforeLoopCallback(Protocol): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, canny_image: PIL.Image.Image | None = None, depth_image: PIL.Image.Image | None = None, ) -> None: ... @@ -26,7 +26,7 @@ class InLoopCallback(Protocol): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ) -> None: ... @@ -37,7 +37,7 @@ class AfterLoopCallback(Protocol): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, ) -> None: ... @@ -48,6 +48,6 @@ class InterruptCallback(Protocol): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ) -> None: ... diff --git a/src/mflux/callbacks/callback_manager.py b/src/mflux/callbacks/callback_manager.py index c439260..fabb5ec 100644 --- a/src/mflux/callbacks/callback_manager.py +++ b/src/mflux/callbacks/callback_manager.py @@ -1,6 +1,5 @@ from argparse import Namespace -from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.battery_saver import BatterySaver from mflux.callbacks.instances.canny_saver import CannyImageSaver from mflux.callbacks.instances.depth_saver import DepthImageSaver @@ -13,32 +12,33 @@ class CallbackManager: def register_callbacks( args: Namespace, model, + latent_creator, enable_canny_saver: bool = False, enable_depth_saver: bool = False, ) -> MemorySaver | None: # Battery saver (always enabled) - CallbackManager._register_battery_saver(args) + CallbackManager._register_battery_saver(args, model) # VAE Tiling (if requested) CallbackManager._register_vae_tiling(args, model) # Specialized savers (based on flags) if enable_canny_saver: - CallbackManager._register_canny_saver(args) + CallbackManager._register_canny_saver(args, model) if enable_depth_saver: - CallbackManager._register_depth_saver(args) + CallbackManager._register_depth_saver(args, model) # Stepwise handler (if requested) - CallbackManager._register_stepwise_handler(args, model) + CallbackManager._register_stepwise_handler(args, model, latent_creator) # Memory saver (if requested) return CallbackManager._register_memory_saver(args, model) @staticmethod - def _register_battery_saver(args: Namespace) -> None: + def _register_battery_saver(args: Namespace, model) -> None: battery_saver = BatterySaver(battery_percentage_stop_limit=args.battery_percentage_stop_limit) - CallbackRegistry.register_before_loop(battery_saver) + model.callbacks.register(battery_saver) @staticmethod def _register_vae_tiling(args: Namespace, model) -> None: @@ -47,31 +47,31 @@ class CallbackManager: model.vae.decoder.split_direction = args.vae_tiling_split @staticmethod - def _register_canny_saver(args: Namespace) -> None: + def _register_canny_saver(args: Namespace, model) -> None: if hasattr(args, "controlnet_save_canny") and args.controlnet_save_canny: canny_image_saver = CannyImageSaver(path=args.output) - CallbackRegistry.register_before_loop(canny_image_saver) + model.callbacks.register(canny_image_saver) @staticmethod - def _register_depth_saver(args: Namespace) -> None: + def _register_depth_saver(args: Namespace, model) -> None: if hasattr(args, "save_depth_map") and args.save_depth_map: depth_image_saver = DepthImageSaver(path=args.output) - CallbackRegistry.register_before_loop(depth_image_saver) + model.callbacks.register(depth_image_saver) @staticmethod - def _register_stepwise_handler(args: Namespace, model) -> None: + def _register_stepwise_handler(args: Namespace, model, latent_creator) -> None: if args.stepwise_image_output_dir: - handler = StepwiseHandler(model=model, output_dir=args.stepwise_image_output_dir) - CallbackRegistry.register_before_loop(handler) - CallbackRegistry.register_in_loop(handler) - CallbackRegistry.register_interrupt(handler) + handler = StepwiseHandler( + model=model, + latent_creator=latent_creator, + output_dir=args.stepwise_image_output_dir, + ) + model.callbacks.register(handler) @staticmethod def _register_memory_saver(args: Namespace, model) -> MemorySaver | None: memory_saver = None if args.low_ram: memory_saver = MemorySaver(model=model, keep_transformer=len(args.seed) > 1) - CallbackRegistry.register_before_loop(memory_saver) - CallbackRegistry.register_in_loop(memory_saver) - CallbackRegistry.register_after_loop(memory_saver) + model.callbacks.register(memory_saver) return memory_saver diff --git a/src/mflux/callbacks/callback_registry.py b/src/mflux/callbacks/callback_registry.py index 7429075..4e610df 100644 --- a/src/mflux/callbacks/callback_registry.py +++ b/src/mflux/callbacks/callback_registry.py @@ -1,40 +1,44 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + from mflux.callbacks.callback import AfterLoopCallback, BeforeLoopCallback, InLoopCallback, InterruptCallback +if TYPE_CHECKING: + from mflux.callbacks.generation_context import GenerationContext + from mflux.models.common.config.config import Config + class CallbackRegistry: - in_loop = [] - before_loop = [] - interrupt = [] - after_loop = [] + def __init__(self): + self.in_loop: list[InLoopCallback] = [] + self.before_loop: list[BeforeLoopCallback] = [] + self.interrupt: list[InterruptCallback] = [] + self.after_loop: list[AfterLoopCallback] = [] - @staticmethod - def register_in_loop(callback: InLoopCallback) -> None: - CallbackRegistry.in_loop.append(callback) + def register(self, callback) -> None: + if hasattr(callback, "call_before_loop"): + self.before_loop.append(callback) + if hasattr(callback, "call_in_loop"): + self.in_loop.append(callback) + if hasattr(callback, "call_after_loop"): + self.after_loop.append(callback) + if hasattr(callback, "call_interrupt"): + self.interrupt.append(callback) - @staticmethod - def register_before_loop(callback: BeforeLoopCallback) -> None: - CallbackRegistry.before_loop.append(callback) + def start(self, seed: int, prompt: str, config: Config) -> GenerationContext: + from mflux.callbacks.generation_context import GenerationContext - @staticmethod - def register_after_loop(callback: AfterLoopCallback) -> None: - CallbackRegistry.after_loop.append(callback) + return GenerationContext(self, seed, prompt, config) - @staticmethod - def register_interrupt(callback: InterruptCallback) -> None: - CallbackRegistry.interrupt.append(callback) + def before_loop_callbacks(self) -> list[BeforeLoopCallback]: + return self.before_loop - @staticmethod - def before_loop_callbacks() -> list[BeforeLoopCallback]: - return CallbackRegistry.before_loop + def in_loop_callbacks(self) -> list[InLoopCallback]: + return self.in_loop - @staticmethod - def in_loop_callbacks() -> list[InLoopCallback]: - return CallbackRegistry.in_loop + def after_loop_callbacks(self) -> list[AfterLoopCallback]: + return self.after_loop - @staticmethod - def after_loop_callbacks() -> list[AfterLoopCallback]: - return CallbackRegistry.after_loop - - @staticmethod - def interrupt_callbacks() -> list[InterruptCallback]: - return CallbackRegistry.interrupt + def interrupt_callbacks(self) -> list[InterruptCallback]: + return self.interrupt diff --git a/src/mflux/callbacks/callbacks.py b/src/mflux/callbacks/callbacks.py index 08d07d9..4c6e55a 100644 --- a/src/mflux/callbacks/callbacks.py +++ b/src/mflux/callbacks/callbacks.py @@ -3,20 +3,21 @@ import PIL.Image import tqdm from mflux.callbacks.callback_registry import CallbackRegistry -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config class Callbacks: @staticmethod def before_loop( + registry: CallbackRegistry, seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, canny_image: PIL.Image.Image | None = None, depth_image: PIL.Image.Image | None = None, ): - for subscriber in CallbackRegistry.before_loop_callbacks(): + for subscriber in registry.before_loop_callbacks(): subscriber.call_before_loop( seed=seed, prompt=prompt, @@ -28,14 +29,15 @@ class Callbacks: @staticmethod def in_loop( + registry: CallbackRegistry, t: int, seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ): - for subscriber in CallbackRegistry.in_loop_callbacks(): + for subscriber in registry.in_loop_callbacks(): subscriber.call_in_loop( t=t, seed=seed, @@ -47,24 +49,26 @@ class Callbacks: @staticmethod def after_loop( + registry: CallbackRegistry, seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, ): - for subscriber in CallbackRegistry.after_loop_callbacks(): + for subscriber in registry.after_loop_callbacks(): subscriber.call_after_loop(seed=seed, prompt=prompt, latents=latents, config=config) @staticmethod def interruption( + registry: CallbackRegistry, t: int, seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ): - for subscriber in CallbackRegistry.interrupt_callbacks(): + for subscriber in registry.interrupt_callbacks(): subscriber.call_interrupt( t=t, seed=seed, diff --git a/src/mflux/callbacks/generation_context.py b/src/mflux/callbacks/generation_context.py new file mode 100644 index 0000000..f6b3e87 --- /dev/null +++ b/src/mflux/callbacks/generation_context.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import mlx.core as mx +import PIL.Image +import tqdm + +if TYPE_CHECKING: + from mflux.callbacks.callback_registry import CallbackRegistry + from mflux.models.common.config.config import Config + + +class GenerationContext: + def __init__( + self, + registry: CallbackRegistry, + seed: int, + prompt: str, + config: Config, + ): + self._registry = registry + self._seed = seed + self._prompt = prompt + self._config = config + + def before_loop( + self, + latents: mx.array, + *, + canny_image: PIL.Image.Image | None = None, + depth_image: PIL.Image.Image | None = None, + ) -> None: + for subscriber in self._registry.before_loop_callbacks(): + subscriber.call_before_loop( + seed=self._seed, + prompt=self._prompt, + latents=latents, + config=self._config, + canny_image=canny_image, + depth_image=depth_image, + ) + + def in_loop(self, t: int, latents: mx.array, time_steps: tqdm = None) -> None: + time_steps = time_steps or self._config.time_steps + for subscriber in self._registry.in_loop_callbacks(): + subscriber.call_in_loop( + t=t, + seed=self._seed, + prompt=self._prompt, + latents=latents, + config=self._config, + time_steps=time_steps, + ) + + def after_loop(self, latents: mx.array) -> None: + for subscriber in self._registry.after_loop_callbacks(): + subscriber.call_after_loop( + seed=self._seed, + prompt=self._prompt, + latents=latents, + config=self._config, + ) + + def interruption(self, t: int, latents: mx.array, time_steps: tqdm = None) -> None: + time_steps = time_steps or self._config.time_steps + for subscriber in self._registry.interrupt_callbacks(): + subscriber.call_interrupt( + t=t, + seed=self._seed, + prompt=self._prompt, + latents=latents, + config=self._config, + time_steps=time_steps, + ) diff --git a/src/mflux/callbacks/instances/battery_saver.py b/src/mflux/callbacks/instances/battery_saver.py index c06d26c..13139fd 100644 --- a/src/mflux/callbacks/instances/battery_saver.py +++ b/src/mflux/callbacks/instances/battery_saver.py @@ -6,57 +6,67 @@ import subprocess from mflux.callbacks.callback import BeforeLoopCallback from mflux.utils.exceptions import StopImageGenerationException -PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'" -PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%" - logger = logging.getLogger(__name__) -def _get_machine_model() -> str: - """Get the Mac machine model using system_profiler.""" - try: - result = subprocess.run( - ["system_profiler", "-json", "SPHardwareDataType"], capture_output=True, text=True, check=True - ) - data = json.loads(result.stdout) - return data["SPHardwareDataType"][0]["machine_model"] - except (subprocess.CalledProcessError, json.JSONDecodeError, IndexError, KeyError) as e: - logger.warning(f"Cannot determine machine model via 'system_profiler -json SPHardwareDataType': {e}") - return "Unknown" - - -MACHINE_MODEL = _get_machine_model() -# assumption: all Apple Silicon models powered by battery are "MacBook"s -MACHINE_IS_BATTERY_POWERED = "MacBook" in MACHINE_MODEL - - -def get_battery_percentage() -> int | None: - """Get the current battery percentage of a battery-powered Mac. - Returns None if Mac is not a battery-powered machine.""" - if not MACHINE_IS_BATTERY_POWERED: - return None - percentage = None - try: - # running the subprocess would be expensive in a tight loop - # but in mflux use case, we would call this only once every - # few minutes due to N-minutes-long generation times - result = subprocess.run(["pmset", "-g", "batt"], capture_output=True, text=True, check=True) - if PMSET_AC_POWER_STATUS not in result.stdout: - if match := re.search(PMSET_BATT_STATUS_PATTERN, result.stdout): - percentage = int(match.group(1)) - except (subprocess.CalledProcessError, TypeError) as e: - logger.warning( - f"Cannot read battery percentage via 'pmset -g batt': {e}. Battery saver functionality is disabled and the program will continue running." - ) - - return percentage - - class BatterySaver(BeforeLoopCallback): - def __init__(self, battery_percentage_stop_limit=10): + PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'" + PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%" + + _machine_model: str | None = None + _is_battery_powered: bool | None = None + + def __init__(self, battery_percentage_stop_limit: int = 10): self.limit = battery_percentage_stop_limit def call_before_loop(self, **kwargs) -> None: # type: ignore - current_pct: int | None = get_battery_percentage() + current_pct = self._get_battery_percentage() if current_pct is not None and current_pct <= self.limit: raise StopImageGenerationException(f"Battery below {self.limit}% threshold: {current_pct}%") + + def _get_battery_percentage(self) -> int | None: + if not self._is_machine_battery_powered(): + return None + + percentage = None + try: + result = subprocess.run( + ["pmset", "-g", "batt"], + capture_output=True, + text=True, + check=True, + ) + if self.PMSET_AC_POWER_STATUS not in result.stdout: + if match := re.search(self.PMSET_BATT_STATUS_PATTERN, result.stdout): + percentage = int(match.group(1)) + except (subprocess.CalledProcessError, TypeError) as e: + logger.warning( + f"Cannot read battery percentage via 'pmset -g batt': {e}. " + f"Battery saver functionality is disabled and the program will continue running." + ) + + return percentage + + @classmethod + def _is_machine_battery_powered(cls) -> bool: + if cls._is_battery_powered is None: + machine_model = cls._get_machine_model() + cls._is_battery_powered = "MacBook" in machine_model + return cls._is_battery_powered + + @classmethod + def _get_machine_model(cls) -> str: + if cls._machine_model is None: + try: + result = subprocess.run( + ["system_profiler", "-json", "SPHardwareDataType"], + capture_output=True, + text=True, + check=True, + ) + data = json.loads(result.stdout) + cls._machine_model = data["SPHardwareDataType"][0]["machine_model"] + except (subprocess.CalledProcessError, json.JSONDecodeError, IndexError, KeyError) as e: + logger.warning(f"Cannot determine machine model via 'system_profiler -json SPHardwareDataType': {e}") + cls._machine_model = "Unknown" + return cls._machine_model diff --git a/src/mflux/callbacks/instances/canny_saver.py b/src/mflux/callbacks/instances/canny_saver.py index 2f2c232..4991269 100644 --- a/src/mflux/callbacks/instances/canny_saver.py +++ b/src/mflux/callbacks/instances/canny_saver.py @@ -5,7 +5,7 @@ import mlx.core as mx import PIL.Image from mflux.callbacks.callback import BeforeLoopCallback -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config from mflux.utils.image_util import ImageUtil @@ -18,7 +18,7 @@ class CannyImageSaver(BeforeLoopCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, canny_image: PIL.Image.Image | None = None, depth_image: PIL.Image.Image | None = None, ) -> None: diff --git a/src/mflux/callbacks/instances/depth_saver.py b/src/mflux/callbacks/instances/depth_saver.py index 6d1f938..5a29c26 100644 --- a/src/mflux/callbacks/instances/depth_saver.py +++ b/src/mflux/callbacks/instances/depth_saver.py @@ -5,7 +5,7 @@ import mlx.core as mx import PIL.Image from mflux.callbacks.callback import BeforeLoopCallback -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config from mflux.utils.image_util import ImageUtil @@ -18,7 +18,7 @@ class DepthImageSaver(BeforeLoopCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, canny_image: PIL.Image.Image | None = None, depth_image: PIL.Image.Image | None = None, ) -> None: @@ -28,5 +28,5 @@ class DepthImageSaver(BeforeLoopCallback): base, ext = os.path.splitext(self.path) ImageUtil.save_image( image=depth_image, - path=f"{base}_depth_map{ext}" - ) # fmt: off + path=f"{base}_depth_map{ext}", + ) diff --git a/src/mflux/callbacks/instances/memory_saver.py b/src/mflux/callbacks/instances/memory_saver.py index bb93416..50bc245 100644 --- a/src/mflux/callbacks/instances/memory_saver.py +++ b/src/mflux/callbacks/instances/memory_saver.py @@ -5,15 +5,10 @@ import PIL.Image from tqdm import tqdm from mflux.callbacks.callback import AfterLoopCallback, BeforeLoopCallback, InLoopCallback -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback): - """ - Optimizes memory usage by clearing caches and removing unused model - components at strategic points in the execution cycle. - """ - def __init__(self, model, keep_transformer: bool = True, cache_limit_bytes: int = 1000**3): self.model = model self.keep_transformer = keep_transformer @@ -27,7 +22,7 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, canny_image: PIL.Image.Image | None = None, depth_image: PIL.Image.Image | None = None, ) -> None: @@ -40,7 +35,7 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ) -> None: self.peak_memory = mx.get_peak_memory() @@ -50,7 +45,7 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, ) -> None: self.peak_memory = mx.get_peak_memory() if not self.keep_transformer: @@ -66,8 +61,9 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback): self.model.text_encoder = None if hasattr(self.model, "qwen_vl_encoder") and self.model.qwen_vl_encoder is not None: self.model.qwen_vl_encoder = None - if hasattr(self.model, "qwen_vl_tokenizer") and self.model.qwen_vl_tokenizer is not None: - self.model.qwen_vl_tokenizer = None + # Clear VLM tokenizers from the tokenizers dict if present + if hasattr(self.model, "tokenizers") and "qwen_vl" in self.model.tokenizers: + self.model.tokenizers["qwen_vl"] = None gc.collect() mx.clear_cache() diff --git a/src/mflux/callbacks/instances/stepwise_handler.py b/src/mflux/callbacks/instances/stepwise_handler.py index fd14cc9..f7c5f61 100644 --- a/src/mflux/callbacks/instances/stepwise_handler.py +++ b/src/mflux/callbacks/instances/stepwise_handler.py @@ -5,8 +5,7 @@ import PIL.Image import tqdm from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback -from mflux.config.runtime_config import RuntimeConfig -from mflux.utils.array_util import ArrayUtil +from mflux.models.common.config.config import Config from mflux.utils.image_util import ImageUtil @@ -15,9 +14,11 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): self, model, output_dir: str, + latent_creator, ): self.model = model self.output_dir = Path(output_dir) + self.latent_creator = latent_creator self.step_wise_images = [] if self.output_dir: @@ -28,7 +29,7 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, canny_image: PIL.Image.Image | None = None, depth_image: PIL.Image.Image | None = None, ) -> None: @@ -47,7 +48,7 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ) -> None: self._save_image( @@ -65,7 +66,7 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ) -> None: self._save_composite(seed=seed) @@ -76,10 +77,10 @@ class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback): seed: int, prompt: str, latents: mx.array, - config: RuntimeConfig, + config: Config, time_steps: tqdm, ) -> None: - unpack_latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + unpack_latents = self.latent_creator.unpack_latents(latents=latents, height=config.height, width=config.width) stepwise_decoded = self.model.vae.decode(unpack_latents) generation_time = time_steps.format_dict["elapsed"] if time_steps is not None else 0 stepwise_img = ImageUtil.to_image( diff --git a/src/mflux/config/__init__.py b/src/mflux/cli/__init__.py similarity index 100% rename from src/mflux/config/__init__.py rename to src/mflux/cli/__init__.py diff --git a/src/mflux/models/common/lora/download/__init__.py b/src/mflux/cli/completions/__init__.py similarity index 100% rename from src/mflux/models/common/lora/download/__init__.py rename to src/mflux/cli/completions/__init__.py diff --git a/src/mflux/ui/cli/completions/generator.py b/src/mflux/cli/completions/generator.py similarity index 76% rename from src/mflux/ui/cli/completions/generator.py rename to src/mflux/cli/completions/generator.py index 263901b..02d9098 100644 --- a/src/mflux/ui/cli/completions/generator.py +++ b/src/mflux/cli/completions/generator.py @@ -1,16 +1,12 @@ -"""Generate ZSH completion scripts for mflux commands.""" - import argparse from pathlib import Path +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_NAME_MAP -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser class CompletionGenerator: - """Generate ZSH completion scripts by introspecting argparse parsers.""" - def __init__(self): self.commands = [ "mflux-generate", @@ -22,6 +18,12 @@ class CompletionGenerator: "mflux-generate-fill", "mflux-generate-depth", "mflux-generate-redux", + "mflux-generate-qwen", + "mflux-generate-qwen-edit", + "mflux-generate-fibo", + "mflux-generate-z-image-turbo", + "mflux-refine-fibo", + "mflux-inspire-fibo", "mflux-concept", "mflux-concept-from-image", "mflux-save", @@ -29,10 +31,10 @@ class CompletionGenerator: "mflux-train", "mflux-upscale", "mflux-lora-library", + "mflux-info", ] def create_parser_for_command(self, command: str) -> CommandLineParser: - """Create the appropriate parser for a given command.""" parser = CommandLineParser(prog=command, add_help=False) if command == "mflux-generate": @@ -111,6 +113,58 @@ class CompletionGenerator: parser.add_redux_arguments() parser.add_output_arguments() + elif command == "mflux-generate-qwen": + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments() + parser.add_output_arguments() + + elif command == "mflux-generate-qwen-edit": + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_argument("--image-paths", type=Path, nargs="+", required=True, help="Local paths to init images") + parser.add_output_arguments() + + elif command == "mflux-generate-fibo": + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments() + parser.add_output_arguments() + + elif command == "mflux-generate-z-image-turbo": + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments() + parser.add_output_arguments() + + elif command == "mflux-refine-fibo": + parser.add_argument("--prompt-file", type=Path, required=True, help="Path to JSON prompt file to refine") + parser.add_argument("--instructions", type=str, required=True, help="Text instructions for refinement") + parser.add_argument("--output", type=Path, help="Output path for refined JSON prompt") + parser.add_argument("--path", type=str, help="Local path for VLM model") + parser.add_argument("--top-p", type=float, help="Top-p sampling for VLM") + parser.add_argument("--temperature", type=float, help="Temperature for VLM") + parser.add_argument("--max-tokens", type=int, help="Max tokens for VLM generation") + parser.add_argument("--seed", type=int, help="Seed for VLM generation") + + elif command == "mflux-inspire-fibo": + parser.add_argument("--image-path", type=Path, required=True, help="Path to image file to inspire from") + parser.add_argument("--prompt", type=str, help="Optional text prompt to blend with the image") + parser.add_argument("--output", type=Path, help="Output path for generated JSON prompt") + parser.add_argument("--path", type=str, help="Local path for VLM model") + parser.add_argument("--top-p", type=float, help="Top-p sampling for VLM") + parser.add_argument("--temperature", type=float, help="Temperature for VLM") + parser.add_argument("--max-tokens", type=int, help="Max tokens for VLM generation") + parser.add_argument("--seed", type=int, help="Seed for VLM generation") + elif command == "mflux-concept": parser.add_general_arguments() parser.add_model_arguments() @@ -149,10 +203,12 @@ class CompletionGenerator: list_parser = subparsers.add_parser("list") list_parser.add_argument("--paths", action="store_true") + elif command == "mflux-info": + parser.add_info_arguments() + return parser def escape_description(self, desc: str) -> str: - """Escape special characters in descriptions for ZSH.""" if not desc: return "" # Escape brackets, quotes, and other special characters @@ -165,7 +221,6 @@ class CompletionGenerator: return desc def format_argument_spec(self, action: argparse.Action) -> list[str]: - """Format an argparse action into ZSH completion syntax.""" specs = [] # Handle options with both short and long forms @@ -195,7 +250,6 @@ class CompletionGenerator: return specs def get_value_spec(self, action: argparse.Action) -> str: - """Get the value specification for an argument.""" # Check for special cases first if action.dest == "model": return ":model:_mflux_models" @@ -241,7 +295,6 @@ class CompletionGenerator: return "" def generate_command_function(self, command: str, parser: CommandLineParser) -> str: - """Generate the ZSH completion function for a specific command.""" func_name = command.replace("-", "_") lines = [f"_{func_name}() {{"] lines.append(" local -a args") @@ -262,7 +315,6 @@ class CompletionGenerator: return "\n".join(lines) def generate_header(self) -> str: - """Generate the completion script header.""" commands = " ".join(self.commands) return f"""#compdef {commands} # ZSH completion for mflux commands @@ -271,7 +323,6 @@ class CompletionGenerator: """ def generate_helper_functions(self) -> str: - """Generate helper functions for common completions.""" helpers = [] # Model completion helper @@ -303,7 +354,6 @@ class CompletionGenerator: return "\n".join(helpers) def generate_main_function(self) -> str: - """Generate the main completion dispatcher.""" lines = ["# Main completion dispatcher", "_mflux() {", " local cmd=$words[1]", " case $cmd in"] for command in self.commands: @@ -317,7 +367,6 @@ class CompletionGenerator: return "\n".join(lines) def generate(self) -> str: - """Generate the complete ZSH completion script.""" script = [self.generate_header()] script.append(self.generate_helper_functions()) diff --git a/src/mflux/ui/cli/completions/install.py b/src/mflux/cli/completions/install.py similarity index 94% rename from src/mflux/ui/cli/completions/install.py rename to src/mflux/cli/completions/install.py index 3d6c47d..28d9fc0 100644 --- a/src/mflux/ui/cli/completions/install.py +++ b/src/mflux/cli/completions/install.py @@ -1,17 +1,14 @@ #!/usr/bin/env python3 -"""Install ZSH completions for mflux commands.""" - import argparse import os import subprocess import sys from pathlib import Path -from mflux.ui.cli.completions.generator import CompletionGenerator +from mflux.cli.completions.generator import CompletionGenerator def get_zsh_fpath(): - """Get the ZSH fpath directories.""" try: result = subprocess.run( ["zsh", "-c", "echo $fpath"], @@ -25,7 +22,6 @@ def get_zsh_fpath(): def find_completion_dir(): - """Find appropriate directory for completion files.""" # Common completion directories in order of preference candidates = [ Path.home() / ".zsh" / "completions", @@ -53,7 +49,6 @@ def find_completion_dir(): def check_installation(): - """Check if completions are properly installed and accessible.""" print("Checking mflux completions installation...\n") # Check if completion file exists @@ -107,7 +102,6 @@ def check_installation(): def main(): - """Main entry point for completion installation.""" parser = argparse.ArgumentParser( description="Install or generate ZSH completions for mflux commands", formatter_class=argparse.RawDescriptionHelpFormatter, diff --git a/src/mflux/models/common/quantization/__init__.py b/src/mflux/cli/defaults/__init__.py similarity index 100% rename from src/mflux/models/common/quantization/__init__.py rename to src/mflux/cli/defaults/__init__.py diff --git a/src/mflux/cli/defaults/defaults.py b/src/mflux/cli/defaults/defaults.py new file mode 100644 index 0000000..4aaadef --- /dev/null +++ b/src/mflux/cli/defaults/defaults.py @@ -0,0 +1,31 @@ +import os +from pathlib import Path + +import platformdirs + +BATTERY_PERCENTAGE_STOP_LIMIT = 5 +CONTROLNET_STRENGTH = 0.4 +DEFAULT_DEV_FILL_GUIDANCE = 30 +DEFAULT_DEPTH_GUIDANCE = 10 +DIMENSION_STEP_PIXELS = 16 +GUIDANCE_SCALE = 3.5 +GUIDANCE_SCALE_KONTEXT = 2.5 +HEIGHT, WIDTH = 1024, 1024 +IMAGE_STRENGTH = 0.4 +MODEL_CHOICES = ["dev", "schnell", "krea-dev", "dev-krea", "qwen", "fibo", "z-image-turbo"] +MODEL_INFERENCE_STEPS = { + "dev": 25, + "schnell": 4, + "krea-dev": 25, + "qwen": 20, + "fibo": 20, + "z-image-turbo": 9, +} +QUANTIZE_CHOICES = [3, 5, 4, 6, 8] + +if os.environ.get("MFLUX_CACHE_DIR"): + MFLUX_CACHE_DIR = Path(os.environ["MFLUX_CACHE_DIR"]).resolve() +else: + MFLUX_CACHE_DIR = Path(platformdirs.user_cache_dir(appname="mflux")) + +MFLUX_LORA_CACHE_DIR = MFLUX_CACHE_DIR / "loras" diff --git a/src/mflux/models/flux/tokenizer/__init__.py b/src/mflux/cli/parser/__init__.py similarity index 100% rename from src/mflux/models/flux/tokenizer/__init__.py rename to src/mflux/cli/parser/__init__.py diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/cli/parser/parsers.py similarity index 93% rename from src/mflux/ui/cli/parsers.py rename to src/mflux/cli/parser/parsers.py index 19ace37..47299b4 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/cli/parser/parsers.py @@ -5,28 +5,19 @@ import time import typing as t from pathlib import Path -from mflux.models.common.lora.download.lora_library import get_lora_path -from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_NAME_MAP, LORA_REPO_ID -from mflux.ui import ( - box_values, - defaults as ui_defaults, - scale_factor, -) +from mflux.cli.defaults import defaults as ui_defaults +from mflux.models.common.resolution.lora_resolution import LoraResolution +from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_NAME_MAP +from mflux.utils import box_values, scale_factor class ModelSpecAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): - if values in ui_defaults.MODEL_CHOICES: - setattr(namespace, self.dest, values) - return - - if values.count("/") != 1: - raise argparse.ArgumentError( - self, - (f'Value must be either {" ".join(ui_defaults.MODEL_CHOICES)} or in format "org/model". Got: {values}'), - ) - - # If we got here, values contains exactly one slash + # Accept: + # 1. Predefined model names (dev, schnell, fibo, etc.) + # 2. HuggingFace repos (org/model format) + # 3. Local paths (/path/to/model, ./model, ~/model) + # The WeightLoader will determine if it's a local path or HuggingFace repo setattr(namespace, self.dest, values) @@ -42,7 +33,7 @@ def int_or_special_value(value) -> int | scale_factor.ScaleFactor: # If not an integer, try to parse as scale factor try: - return scale_factor.parse_scale_factor(value) + return scale_factor.ScaleFactor.parse(value) except ValueError: raise argparse.ArgumentTypeError( f"'{value}' is not a valid integer or 'auto' or a scale factor like '2x' or '3.5x'" @@ -71,10 +62,8 @@ class CommandLineParser(argparse.ArgumentParser): def add_model_arguments(self, path_type: t.Literal["load", "save"] = "load", require_model_arg: bool = True) -> None: self.require_model_arg = require_model_arg - self.add_argument("--model", "-m", type=str, required=require_model_arg, action=ModelSpecAction, help=f"The model to use ({' or '.join(ui_defaults.MODEL_CHOICES)} or a compatible huggingface repo_id org/model).") - if path_type == "load": - self.add_argument("--path", type=str, default=None, help="Local path for loading a model from disk") - else: + self.add_argument("--model", "-m", type=str, required=require_model_arg, action=ModelSpecAction, help=f"The model to use ({' or '.join(ui_defaults.MODEL_CHOICES)}, a HuggingFace repo org/model, or a local path).") + if path_type == "save": self.add_argument("--path", type=str, required=True, help="Local path for saving a model to disk.") self.add_argument("--base-model", type=str, required=False, choices=ui_defaults.MODEL_CHOICES, help="When using a third-party huggingface model, explicitly specify whether the base model is dev or schnell") self.add_argument("--quantize", "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)") @@ -83,10 +72,8 @@ class CommandLineParser(argparse.ArgumentParser): self.supports_lora = True lora_group = self.add_argument_group("LoRA configuration") lora_group.add_argument("--lora-style", type=str, choices=sorted(LORA_NAME_MAP.keys()), help="Style of the LoRA to use (e.g., 'storyboard' for film storyboard style)") - self.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk") + self.add_argument("--lora-paths", type=str, nargs="*", default=None, help="LoRA paths: local files, HuggingFace repos (org/model), or collection format (repo:filename.safetensors)") self.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.") - lora_group.add_argument("--lora-name", type=str, help="Name of the LoRA to download from Hugging Face") - lora_group.add_argument("--lora-repo-id", type=str, default=LORA_REPO_ID, help=f"Hugging Face repository ID for LoRAs (default: {LORA_REPO_ID})") def _add_image_generator_common_arguments(self, supports_dimension_scale_factor=False) -> None: self.supports_image_generation = True @@ -312,7 +299,7 @@ class CommandLineParser(argparse.ArgumentParser): 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) + namespace.image_outpaint_padding = box_values.BoxValues.parse(namespace.image_outpaint_padding) print(f"{namespace.image_outpaint_padding=}") # Resolve lora paths from library if needed @@ -320,10 +307,17 @@ class CommandLineParser(argparse.ArgumentParser): resolved_paths = [] for lora_path in namespace.lora_paths: try: - resolved_path = get_lora_path(lora_path) + resolved_path = LoraResolution.resolve(lora_path) resolved_paths.append(resolved_path) except FileNotFoundError as e: # noqa: PERF203 self.error(str(e)) namespace.lora_paths = resolved_paths + # Compute model_path: None for predefined names, otherwise use the model value + # Predefined names like "schnell", "dev" are handled by ModelConfig, not PathResolution + if hasattr(namespace, "model") and namespace.model is not None: + namespace.model_path = None if namespace.model in ui_defaults.MODEL_CHOICES else namespace.model + else: + namespace.model_path = None + return namespace diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py deleted file mode 100644 index 001923a..0000000 --- a/src/mflux/config/config.py +++ /dev/null @@ -1,40 +0,0 @@ -import logging -from pathlib import Path - -import mlx.core as mx - -log = logging.getLogger(__name__) - - -class Config: - precision: mx.Dtype = mx.bfloat16 - - def __init__( - self, - num_inference_steps: int = 4, - width: int = 1024, - height: int = 1024, - guidance: float = 4.0, - image_path: Path | None = None, - image_strength: float | None = None, - depth_image_path: Path | None = None, - redux_image_paths: list[Path] | None = None, - redux_image_strengths: list[float] | None = None, - masked_image_path: Path | None = None, - controlnet_strength: float | None = None, - scheduler: str = "linear", - ): - if width % 16 != 0 or height % 16 != 0: - log.warning("Width and height should be multiples of 16. Rounding down.") - self.width = 16 * (width // 16) - self.height = 16 * (height // 16) - self.num_inference_steps = num_inference_steps - self.guidance = guidance - self.image_path = image_path - self.image_strength = image_strength - self.depth_image_path = depth_image_path - self.redux_image_paths = redux_image_paths - self.redux_image_strengths = redux_image_strengths - self.masked_image_path = masked_image_path - self.controlnet_strength = controlnet_strength - self.scheduler_str = scheduler diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py deleted file mode 100644 index 7d99a5b..0000000 --- a/src/mflux/config/runtime_config.py +++ /dev/null @@ -1,116 +0,0 @@ -import logging -from pathlib import Path - -import mlx.core as mx - -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.models.common.schedulers import SCHEDULER_REGISTRY, try_import_external_scheduler -from mflux.models.common.schedulers.linear_scheduler import LinearScheduler - -logger = logging.getLogger(__name__) - - -class RuntimeConfig: - def __init__( - self, - config: Config, - model_config: ModelConfig, - ): - self.config = config - self.model_config = model_config - self._scheduler = None - - @property - def height(self) -> int: - return self.config.height - - @property - def width(self) -> int: - return self.config.width - - @width.setter - def width(self, value): - self.config.width = value - - @property - def guidance(self) -> float: - return self.config.guidance - - @property - def num_inference_steps(self) -> int: - return self.config.num_inference_steps - - @property - def precision(self) -> mx.Dtype: - return self.config.precision - - @property - def num_train_steps(self) -> int: - return self.model_config.num_train_steps - - @property - def image_path(self) -> Path | None: - return self.config.image_path - - @property - def image_strength(self) -> float | None: - return self.config.image_strength - - @property - def depth_image_path(self) -> Path | None: - return self.config.depth_image_path - - @property - def redux_image_paths(self) -> list[Path] | None: - return self.config.redux_image_paths - - @property - def redux_image_strengths(self) -> list[float] | None: - return self.config.redux_image_strengths - - @property - def masked_image_path(self) -> Path | None: - return self.config.masked_image_path - - @property - def init_time_step(self) -> int: - is_img2img = ( - self.config.image_path is not None and - self.image_strength is not None and - self.image_strength > 0.0 - ) # fmt: off - - if is_img2img: - # 1. Clamp strength to [0, 1] - strength = max(0.0, min(1.0, self.config.image_strength)) # type: ignore - - # 2. Return start time in [1, floor(num_steps * strength)] - return max(1, int(self.num_inference_steps * strength)) # type: ignore - else: - return 0 - - @property - def controlnet_strength(self) -> float | None: - if self.config.controlnet_strength is not None: - return self.config.controlnet_strength - - return None - - @property - def scheduler(self): - if self._scheduler is not None: - return self._scheduler - - if self.config.scheduler_str == "linear": - self._scheduler = LinearScheduler(self) - elif (registered_scheduler := SCHEDULER_REGISTRY.get(self.config.scheduler_str, None)) is not None: - self._scheduler = registered_scheduler(self) - elif "." in self.config.scheduler_str: - # this raises ValueError if scheduler is not importable - scheduler_cls = try_import_external_scheduler(self.config.scheduler_str) - self._scheduler = scheduler_cls(self) - else: - raise NotImplementedError(f"The scheduler {self.config.scheduler_str!r} is not implemented by mflux.") - - return self._scheduler diff --git a/src/mflux/info.py b/src/mflux/info.py deleted file mode 100644 index c74cc4a..0000000 --- a/src/mflux/info.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Display metadata information from MFLUX generated images.""" - -import sys -from datetime import datetime -from pathlib import Path - -from mflux.ui.cli.parsers import CommandLineParser -from mflux.utils.metadata_reader import MetadataReader - - -def format_metadata(metadata: dict) -> str: - """Format metadata in a clean, readable format.""" - exif = metadata.get("exif", {}) - if not exif: - return "No metadata found" - - lines = [] - lines.append("=" * 60) - lines.append("MFLUX Image Information") - lines.append("=" * 60) - - # Prompt - if prompt := exif.get("prompt"): - lines.append(f"\nPrompt: {prompt}") - - if negative_prompt := exif.get("negative_prompt"): - lines.append(f"Negative Prompt: {negative_prompt}") - - # Model information - lines.append("") - if model := exif.get("model"): - lines.append(f"Model: {model}") - - # Image dimensions - if width := exif.get("width"): - lines.append(f"Width: {width}") - if height := exif.get("height"): - lines.append(f"Height: {height}") - - # Generation parameters - lines.append("") - if seed := exif.get("seed"): - lines.append(f"Seed: {seed}") - if steps := exif.get("steps"): - lines.append(f"Steps: {steps}") - if guidance := exif.get("guidance"): - lines.append(f"Guidance: {guidance}") - - # Technical settings - if quantize := exif.get("quantize"): - lines.append(f"Quantization: {quantize}-bit") - if precision := exif.get("precision"): - lines.append(f"Precision: {precision}") - - # LoRA information - if lora_paths := exif.get("lora_paths"): - lines.append("") - lines.append(f"LoRAs ({len(lora_paths)}):") - lora_scales = exif.get("lora_scales") or [] - for i, lora in enumerate(lora_paths): - scale = lora_scales[i] if i < len(lora_scales) else 1.0 - lora_name = Path(lora).name - lines.append(f" - {lora_name} (scale: {scale})") - - # Image-to-image parameters - if image_path := exif.get("image_path"): - lines.append("") - lines.append(f"Source Image: {Path(image_path).name}") - if image_strength := exif.get("image_strength"): - lines.append(f"Image Strength: {image_strength}") - - # ControlNet parameters - if controlnet_path := exif.get("controlnet_image_path"): - lines.append("") - lines.append(f"ControlNet Image: {Path(controlnet_path).name}") - if controlnet_strength := exif.get("controlnet_strength"): - lines.append(f"ControlNet Strength: {controlnet_strength}") - - # Generation metadata - lines.append("") - if gen_time := exif.get("generation_time_seconds"): - lines.append(f"Generation Time: {gen_time:.2f}s") - - if created_at := exif.get("created_at"): - try: - dt = datetime.fromisoformat(created_at) - lines.append(f"Created: {dt.strftime('%Y-%m-%d %H:%M:%S')}") - except (ValueError, AttributeError): - lines.append(f"Created: {created_at}") - - if version := exif.get("mflux_version"): - lines.append(f"MFLUX Version: {version}") - - lines.append("=" * 60) - return "\n".join(lines) - - -def main(): - # Parse command line arguments - parser = CommandLineParser(description="Display metadata from MFLUX generated images") - parser.add_info_arguments() - args = parser.parse_args() - - # Check if file exists - image_path = Path(args.image_path) - if not image_path.exists(): - print(f"Error: Image file not found: {image_path}") - sys.exit(1) - - # Read metadata - metadata = MetadataReader.read_all_metadata(image_path) - - # Check if metadata was found - if not metadata or (not metadata.get("exif") and not metadata.get("xmp")): - print("No metadata found") - sys.exit(1) - - # Format and display - print(format_metadata(metadata)) - - -if __name__ == "__main__": - main() diff --git a/src/mflux/lora_library.py b/src/mflux/lora_library.py deleted file mode 100644 index ca650f2..0000000 --- a/src/mflux/lora_library.py +++ /dev/null @@ -1,133 +0,0 @@ -#!/usr/bin/env python3 -"""CLI tool for managing the MFLUX LoRA library.""" - -import argparse -import os -import sys -from collections import defaultdict -from pathlib import Path - -from mflux.models.common.lora.download.lora_library import _discover_lora_files - - -def list_loras(paths: list[str] | None = None) -> int: - """List all discovered LoRA files from specified paths or LORA_LIBRARY_PATH. - - Args: - paths: Optional list of paths to use instead of LORA_LIBRARY_PATH - - Returns: - Exit code (0 for success, 1 for error) - """ - if paths: - # Use provided paths - library_paths = [Path(p.strip()) for p in paths] - else: - # Use environment variable - library_path_env = os.environ.get("LORA_LIBRARY_PATH") - - if not library_path_env: - print("LORA_LIBRARY_PATH environment variable is not set.", file=sys.stderr) - print("Set it to one or more colon-separated directories containing .safetensors files.", file=sys.stderr) - print("Alternatively, use --paths to specify directories directly.", file=sys.stderr) - return 1 - - # Parse library paths from environment - library_paths = [Path(p.strip()) for p in library_path_env.split(":") if p.strip()] - - # Check which paths exist - valid_paths = [] - for path in library_paths: - if path.exists() and path.is_dir(): - valid_paths.append(path) - else: - print(f"Warning: Path does not exist or is not a directory: {path}", file=sys.stderr) - - if not valid_paths: - print("No valid directories found in LORA_LIBRARY_PATH.", file=sys.stderr) - return 1 - - # Discover all LoRA files - lora_registry = _discover_lora_files(valid_paths) - - if not lora_registry: - print("No .safetensors files found in the specified directories.") - return 0 - - # Sort by basename for consistent output - sorted_items = sorted(lora_registry.items()) - - # Print all discovered LoRAs - print("Discovered LoRA files:") - print("-" * 80) - for basename, full_path in sorted_items: - print(f"{basename} -> {full_path}") - - # Calculate statistics per top-level directory - stats = defaultdict(int) - for full_path in lora_registry.values(): - # Find which library path this file belongs to - for lib_path in valid_paths: - try: - # Check if full_path is relative to lib_path - full_path.relative_to(lib_path.resolve()) - stats[str(lib_path)] += 1 - break - except ValueError: - # Not relative to this lib_path, continue - continue - - # Print summary - print("-" * 80) - print(f"\nTotal LoRA files found: {len(lora_registry)}") - - if len(valid_paths) > 1: - print("\nBreakdown by library path:") - for lib_path in valid_paths: - count = stats.get(str(lib_path), 0) - print(f" {lib_path}: {count} files") - - return 0 - - -def main(): - """Main entry point for the CLI.""" - parser = argparse.ArgumentParser( - description="MFLUX LoRA Library management tool", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -Environment Variables: - LORA_LIBRARY_PATH Colon-separated list of directories containing .safetensors files - Example: /path/to/loras:/another/path/to/loras - -Examples: - # List all discovered LoRA files using LORA_LIBRARY_PATH - mflux-lora-library list - - # With environment variable set - LORA_LIBRARY_PATH=/home/user/loras:/opt/shared/loras mflux-lora-library list - - # Override with specific paths - mflux-lora-library list --paths /path/to/loras /another/path/to/loras -""", - ) - - subparsers = parser.add_subparsers(dest="command", help="Available commands") - - # Add 'list' command - list_parser = subparsers.add_parser("list", help="List all discovered LoRA files") - list_parser.add_argument( - "--paths", nargs="+", help="Override LORA_LIBRARY_PATH with these directories (space-separated)" - ) - - args = parser.parse_args() - - if args.command == "list": - return list_loras(paths=args.paths) - else: - parser.print_help() - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/src/mflux/models/flux/variants/kontext/utils/__init__.py b/src/mflux/models/common/cli/__init__.py similarity index 100% rename from src/mflux/models/flux/variants/kontext/utils/__init__.py rename to src/mflux/models/common/cli/__init__.py diff --git a/src/mflux/models/common/cli/info.py b/src/mflux/models/common/cli/info.py new file mode 100644 index 0000000..0a2cbcb --- /dev/null +++ b/src/mflux/models/common/cli/info.py @@ -0,0 +1,34 @@ +import sys +from pathlib import Path + +from mflux.cli.parser.parsers import CommandLineParser +from mflux.utils.info_util import InfoUtil +from mflux.utils.metadata_reader import MetadataReader + + +def main(): + # Parse command line arguments + parser = CommandLineParser(description="Display metadata from MFLUX generated images") + parser.add_info_arguments() + args = parser.parse_args() + + # Check if file exists + image_path = Path(args.image_path) + if not image_path.exists(): + print(f"Error: Image file not found: {image_path}") + sys.exit(1) + + # Read metadata + metadata = MetadataReader.read_all_metadata(image_path) + + # Check if metadata was found + if not metadata or (not metadata.get("exif") and not metadata.get("xmp")): + print("No metadata found") + sys.exit(1) + + # Format and display + print(InfoUtil.format_metadata(metadata)) + + +if __name__ == "__main__": + main() diff --git a/src/mflux/models/common/cli/lora_library.py b/src/mflux/models/common/cli/lora_library.py new file mode 100644 index 0000000..89f088c --- /dev/null +++ b/src/mflux/models/common/cli/lora_library.py @@ -0,0 +1,26 @@ +import argparse +import sys + +from mflux.utils.lora_library_util import LoraLibraryUtil + + +def main(): + parser = argparse.ArgumentParser( + description="MFLUX LoRA Library management tool", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=LoraLibraryUtil.epilog(), + ) + subparsers = parser.add_subparsers(dest="command", help="Available commands") + list_parser = subparsers.add_parser("list", help="List all discovered LoRA files") + list_parser.add_argument("--paths", nargs="+", help="Override LORA_LIBRARY_PATH with these directories (space-separated)") # fmt: off + args = parser.parse_args() + + if args.command == "list": + return LoraLibraryUtil.list_loras(paths=args.paths) + else: + parser.print_help() + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/mflux/save.py b/src/mflux/models/common/cli/save.py similarity index 62% rename from src/mflux/save.py rename to src/mflux/models/common/cli/save.py index a157767..7751313 100644 --- a/src/mflux/save.py +++ b/src/mflux/models/common/cli/save.py @@ -1,8 +1,9 @@ -from mflux.config.model_config import ModelConfig +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig from mflux.models.fibo.variants.txt2img.fibo import FIBO from mflux.models.flux.variants.txt2img.flux import Flux1 from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage -from mflux.ui.cli.parsers import CommandLineParser +from mflux.models.z_image.variants.turbo.z_image_turbo import ZImageTurbo def main(): @@ -18,19 +19,18 @@ def main(): model_class = QwenImage elif "fibo" in model_name_lower: model_class = FIBO + elif "z-image" in model_name_lower or "zimage" in model_name_lower: + model_class = ZImageTurbo else: model_class = Flux1 # 2. Load, quantize and save the model - model_kwargs = { - "model_config": ModelConfig.from_name(args.model, base_model=args.base_model), - "quantize": args.quantize, - } - if args.lora_paths and model_class != FIBO: - model_kwargs["lora_paths"] = args.lora_paths - model_kwargs["lora_scales"] = args.lora_scales - - model = model_class(**model_kwargs) + model = model_class( + quantize=args.quantize, + lora_paths=args.lora_paths, + lora_scales=args.lora_scales, + model_config=ModelConfig.from_name(args.model, base_model=args.base_model), + ) model.save_model(args.path) diff --git a/src/mflux/train.py b/src/mflux/models/common/cli/train.py similarity index 80% rename from src/mflux/train.py rename to src/mflux/models/common/cli/train.py index 0dce622..fba7a45 100644 --- a/src/mflux/train.py +++ b/src/mflux/models/common/cli/train.py @@ -1,6 +1,6 @@ +from mflux.cli.parser.parsers import CommandLineParser from mflux.models.flux.variants.dreambooth.dreambooth import DreamBooth from mflux.models.flux.variants.dreambooth.dreambooth_initializer import DreamBoothInitializer -from mflux.ui.cli.parsers import CommandLineParser from mflux.utils.exceptions import StopTrainingException @@ -13,7 +13,7 @@ def main(): args = parser.parse_args() # 1. Initialize the required resources - flux, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + flux, config, training_spec, training_state = DreamBoothInitializer.initialize( config_path=args.train_config, checkpoint_path=args.train_checkpoint, ) @@ -22,12 +22,12 @@ def main(): try: DreamBooth.train( flux=flux, - runtime_config=runtime_config, + config=config, training_spec=training_spec, training_state=training_state, ) except StopTrainingException as stop_exc: - training_state.save(training_spec) + training_state.save(flux, training_spec) print(stop_exc) diff --git a/src/mflux/models/common/config/__init__.py b/src/mflux/models/common/config/__init__.py new file mode 100644 index 0000000..0c26da7 --- /dev/null +++ b/src/mflux/models/common/config/__init__.py @@ -0,0 +1,4 @@ +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig + +__all__ = ["Config", "ModelConfig"] diff --git a/src/mflux/models/common/config/config.py b/src/mflux/models/common/config/config.py new file mode 100644 index 0000000..02d557f --- /dev/null +++ b/src/mflux/models/common/config/config.py @@ -0,0 +1,148 @@ +import logging +from pathlib import Path + +import mlx.core as mx +from tqdm import tqdm + +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.common.schedulers import SCHEDULER_REGISTRY, try_import_external_scheduler +from mflux.models.common.schedulers.linear_scheduler import LinearScheduler + +logger = logging.getLogger(__name__) + + +class Config: + def __init__( + self, + model_config: ModelConfig, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + depth_image_path: Path | str | None = None, + redux_image_paths: list[Path | str] | None = None, + redux_image_strengths: list[float] | None = None, + masked_image_path: Path | str | None = None, + controlnet_strength: float | None = None, + scheduler: str = "linear", + ): + # Ensure dimensions are multiples of 16 + if width % 16 != 0 or height % 16 != 0: + logger.warning("Width and height should be multiples of 16. Rounding down.") + + self.model_config = model_config + self._num_inference_steps = num_inference_steps + self._height = 16 * (height // 16) + self._width = 16 * (width // 16) + self._guidance = guidance + self._image_path = Path(image_path) if isinstance(image_path, str) else image_path + self._image_strength = image_strength + self._depth_image_path = Path(depth_image_path) if isinstance(depth_image_path, str) else depth_image_path + self._redux_image_paths = ( + [Path(p) if isinstance(p, str) else p for p in redux_image_paths] if redux_image_paths else None + ) + self._redux_image_strengths = redux_image_strengths + self._masked_image_path = Path(masked_image_path) if isinstance(masked_image_path, str) else masked_image_path + self._controlnet_strength = controlnet_strength + self._scheduler_str = scheduler + self._scheduler = None + self._time_steps = None + + @property + def height(self) -> int: + return self._height + + @property + def width(self) -> int: + return self._width + + @width.setter + def width(self, value): + self._width = value + + @property + def guidance(self) -> float: + return self._guidance + + @property + def num_inference_steps(self) -> int: + return self._num_inference_steps + + @property + def precision(self) -> mx.Dtype: + return ModelConfig.precision + + @property + def num_train_steps(self) -> int: + return self.model_config.num_train_steps + + @property + def image_path(self) -> Path | None: + return self._image_path + + @property + def image_strength(self) -> float | None: + return self._image_strength + + @property + def depth_image_path(self) -> Path | None: + return self._depth_image_path + + @property + def redux_image_paths(self) -> list[Path] | None: + return self._redux_image_paths + + @property + def redux_image_strengths(self) -> list[float] | None: + return self._redux_image_strengths + + @property + def masked_image_path(self) -> Path | None: + return self._masked_image_path + + @property + def init_time_step(self) -> int: + is_img2img = ( + self._image_path is not None and + self._image_strength is not None and + self._image_strength > 0.0 + ) # fmt: off + + if is_img2img: + # 1. Clamp strength to [0, 1] + strength = max(0.0, min(1.0, self._image_strength)) # type: ignore + + # 2. Return start time in [1, floor(num_steps * strength)] + return max(1, int(self._num_inference_steps * strength)) # type: ignore + else: + return 0 + + @property + def time_steps(self) -> tqdm: + if self._time_steps is None: + self._time_steps = tqdm(range(self.init_time_step, self.num_inference_steps)) + return self._time_steps + + @property + def controlnet_strength(self) -> float | None: + return self._controlnet_strength + + @property + def scheduler(self): + if self._scheduler is not None: + return self._scheduler + + if self._scheduler_str == "linear": + self._scheduler = LinearScheduler(self) + elif (registered_scheduler := SCHEDULER_REGISTRY.get(self._scheduler_str, None)) is not None: + self._scheduler = registered_scheduler(self) + elif "." in self._scheduler_str: + # this raises ValueError if scheduler is not importable + scheduler_cls = try_import_external_scheduler(self._scheduler_str) + self._scheduler = scheduler_cls(self) + else: + raise NotImplementedError(f"The scheduler {self._scheduler_str!r} is not implemented by mflux.") + + return self._scheduler diff --git a/src/mflux/config/model_config.py b/src/mflux/models/common/config/model_config.py similarity index 77% rename from src/mflux/config/model_config.py rename to src/mflux/models/common/config/model_config.py index c51d316..dc11af8 100644 --- a/src/mflux/config/model_config.py +++ b/src/mflux/models/common/config/model_config.py @@ -1,10 +1,14 @@ from functools import lru_cache from typing import Literal -from mflux.utils.exceptions import InvalidBaseModel, ModelConfigError +import mlx.core as mx + +from mflux.models.common.resolution.config_resolution import ConfigResolution class ModelConfig: + precision: mx.Dtype = mx.bfloat16 + def __init__( self, aliases: list[str], @@ -99,6 +103,11 @@ class ModelConfig: def fibo() -> "ModelConfig": return AVAILABLE_MODELS["fibo"] + @staticmethod + @lru_cache + def z_image_turbo() -> "ModelConfig": + return AVAILABLE_MODELS["z-image-turbo"] + def x_embedder_input_dim(self) -> int: if "Fill" in self.model_name: return 384 @@ -115,56 +124,7 @@ class ModelConfig: model_name: str, base_model: Literal["dev", "schnell", "krea-dev"] | None = None, ) -> "ModelConfig": - # 0. Get all base models (where base_model is None) sorted by priority - base_models = sorted( - [model for model in AVAILABLE_MODELS.values() if model.base_model is None], key=lambda x: x.priority - ) - - # 1. Check if model_name matches any base model's aliases or full name - for base in base_models: - if model_name == base.model_name or model_name in base.aliases: - return base - - # 2. Validate explicit base_model - allowed_names = [] - for base in base_models: - allowed_names.extend(base.aliases + [base.model_name]) - if base_model and base_model not in allowed_names: - raise InvalidBaseModel(f"Invalid base_model. Choose one of {allowed_names}") - - # 3. Determine the base model (explicit or inferred) - if base_model: - # Find by explicit base_model name (check all aliases) - default_base = next((b for b in base_models if base_model == b.model_name or base_model in b.aliases), None) - else: - # Infer from model_name substring - prefer longer matches (more specific) - # Use case-insensitive matching for better compatibility - model_name_lower = model_name.lower() - matching_bases = [ - (b, alias) for b in base_models for alias in b.aliases if alias and alias.lower() in model_name_lower - ] - - if matching_bases: - # Sort by alias length descending, then by priority ascending - default_base = sorted(matching_bases, key=lambda x: (-len(x[1]), x[0].priority))[0][0] - else: - default_base = None - if not default_base: - raise ModelConfigError(f"Cannot infer base_model from {model_name}") - - # 4. Construct the config - return ModelConfig( - aliases=default_base.aliases, - 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, - requires_sigma_shift=default_base.requires_sigma_shift, - priority=default_base.priority, - ) + return ConfigResolution.resolve(model_name=model_name, base_model=base_model) AVAILABLE_MODELS = { @@ -336,4 +296,16 @@ AVAILABLE_MODELS = { requires_sigma_shift=False, priority=13, ), + "z-image-turbo": ModelConfig( + aliases=["z-image-turbo", "z-image", "zimage-turbo", "zimage"], + model_name="Tongyi-MAI/Z-Image-Turbo", + base_model=None, + controlnet_model=None, + custom_transformer_model=None, + num_train_steps=1000, + max_sequence_length=512, + supports_guidance=False, # Turbo model uses guidance_scale=0 + requires_sigma_shift=True, + priority=14, + ), } diff --git a/src/mflux/models/common/latent_creator/latent_creator.py b/src/mflux/models/common/latent_creator/latent_creator.py index 1696741..41c72e5 100644 --- a/src/mflux/models/common/latent_creator/latent_creator.py +++ b/src/mflux/models/common/latent_creator/latent_creator.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeAlias import mlx.core as mx from mlx import nn @@ -10,13 +10,16 @@ if TYPE_CHECKING: from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator + from mflux.models.z_image.latent_creator.z_image_latent_creator import ZImageLatentCreator + + LatentCreatorType: TypeAlias = type[FiboLatentCreator | FluxLatentCreator | QwenLatentCreator | ZImageLatentCreator] class Img2Img: def __init__( self, vae: nn.Module, - latent_creator: type["FiboLatentCreator"] | type["FluxLatentCreator"] | type["QwenLatentCreator"], + latent_creator: "LatentCreatorType", sigmas: mx.array, init_time_step: int, image_path: str | Path | None, diff --git a/src/mflux/models/common/lora/download/lora_huggingface_downloader.py b/src/mflux/models/common/lora/download/lora_huggingface_downloader.py deleted file mode 100644 index 502ef69..0000000 --- a/src/mflux/models/common/lora/download/lora_huggingface_downloader.py +++ /dev/null @@ -1,101 +0,0 @@ -import shutil -from pathlib import Path - -from mflux.ui.defaults import MFLUX_LORA_CACHE_DIR -from mflux.utils.download import snapshot_download - - -class LoRAHuggingFaceDownloader: - - @staticmethod - def download_loras( - lora_names: list[str] | None = None, - repo_id: str | None = None, - cache_dir: Path | str | None = None, - model_name: str = "LoRA", - ) -> list[str]: - if not lora_names or not repo_id: - return [] - - lora_paths = [] - for lora_name in lora_names: - lora_path = LoRAHuggingFaceDownloader.download_lora( - repo_id=repo_id, - lora_name=lora_name, - cache_dir=cache_dir, - model_name=model_name, - ) - lora_paths.append(lora_path) - - return lora_paths - - @staticmethod - def download_lora( - repo_id: str, - lora_name: str, - cache_dir: Path | str | None = None, - model_name: str = "LoRA", # For logging purposes - ) -> str: - # Ensure cache_dir is a Path object - if cache_dir is None: - cache_path = MFLUX_LORA_CACHE_DIR - else: - cache_path = Path(cache_dir) - - cache_path.mkdir(parents=True, exist_ok=True) - - # Check if already cached - cached_file_path = cache_path / lora_name - if cached_file_path.exists() and cached_file_path.is_file(): - try: - # Verify the file is actually readable (catches broken symlinks) - with open(cached_file_path, "rb") as f: - f.read(1) # Try to read just 1 byte to verify it works - print(f"Using cached {model_name} LoRA: {cached_file_path}") - return str(cached_file_path) - except (OSError, IOError): - # File exists but is not readable (broken symlink, permissions, etc.) - print(f"Cached {model_name} LoRA file is corrupted or inaccessible, re-downloading: {cached_file_path}") - try: - cached_file_path.unlink() # Remove the broken file/symlink - except OSError: - pass # Ignore if we can't remove it - - # Download the LoRA from Hugging Face - print(f"Downloading {model_name} LoRA '{lora_name}' from {repo_id}...") - download_path = Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=[f"*{lora_name}*"], - cache_dir=str(cache_path), - ) - ) - - # Find the downloaded file - print(f"๐Ÿ” Searching for downloaded files in: {download_path}") - found_files = list(download_path.glob(f"**/*{lora_name}*")) - print(f"๐Ÿ“ Found files matching pattern: {found_files}") - - for file in found_files: - print(f"๐Ÿ“„ Checking file: {file} (suffix: {file.suffix}, size: {file.stat().st_size} bytes)") - if file.is_file() and file.suffix in [".safetensors", ".bin"]: - # Ensure the target path has the correct extension - if not lora_name.endswith(file.suffix): - target_name = f"{lora_name}{file.suffix}" - else: - target_name = lora_name - - target_path = cache_path / target_name - if not target_path.exists(): - # Create a symlink or copy the file - try: - target_path.symlink_to(file) - print(f"๐Ÿ”— Created symlink: {target_path} -> {file}") - except (OSError, AttributeError): - shutil.copy2(file, target_path) - print(f"๐Ÿ“‹ Copied file: {file} -> {target_path}") - - print(f"{model_name} LoRA downloaded to: {target_path}") - return str(target_path) - - raise FileNotFoundError(f"Could not find {model_name} LoRA file '{lora_name}' in the downloaded repository.") diff --git a/src/mflux/models/common/lora/download/lora_library.py b/src/mflux/models/common/lora/download/lora_library.py deleted file mode 100644 index cde3ac8..0000000 --- a/src/mflux/models/common/lora/download/lora_library.py +++ /dev/null @@ -1,87 +0,0 @@ -import os -from pathlib import Path - - -def _discover_lora_files(library_paths: list[Path]) -> dict[str, Path]: - """ - Discover all .safetensors files in the library paths and their subdirectories. - Earlier paths in the list have higher precedence for duplicate basenames. - - Args: - library_paths: List of paths to LORA library directories (in precedence order) - - Returns: - Dictionary mapping basename (without extension) to full path - """ - lora_files = {} - - # Process paths in reverse order so earlier paths overwrite later ones - for library_path in reversed(library_paths): - if not library_path.exists() or not library_path.is_dir(): - continue - - # Find all .safetensors files recursively - for safetensor_path in library_path.rglob("*.safetensors"): - # Use the basename without extension as the key - basename = safetensor_path.stem - - # Skip files with digit-only names (0-9) in transformer directories - if basename.isdigit() and safetensor_path.parent.name == "transformer": - continue - - # Earlier paths in the list take precedence (overwrite) - lora_files[basename] = safetensor_path.resolve() - - return lora_files - - -# Global registry that will be populated on module import -_LORA_REGISTRY: dict[str, Path] = {} - - -def _initialize_registry() -> None: - """Initialize the global LORA registry from LORA_LIBRARY_PATH environment variable.""" - global _LORA_REGISTRY - - library_path_env = os.environ.get("LORA_LIBRARY_PATH") - if library_path_env: - # Split by colon to support multiple paths - library_paths = [Path(p.strip()) for p in library_path_env.split(":") if p.strip()] - _LORA_REGISTRY = _discover_lora_files(library_paths) - - -def get_lora_path(path_or_name: str) -> str: - """ - Get the full path for a LORA file, resolving from library if needed. - - Args: - path_or_name: Either a full path or a basename that exists in the library - - Returns: - The resolved path as a string - - Raises: - FileNotFoundError: If the file cannot be found either as a path or in the registry - """ - # If it's already a path that exists, return it as-is - path = Path(path_or_name) - if path.exists(): - return str(path) - - # Otherwise, check if it's in the registry - if path_or_name in _LORA_REGISTRY: - return str(_LORA_REGISTRY[path_or_name]) - - # If not found, raise FileNotFoundError - raise FileNotFoundError( - f"LoRA file not found: '{path_or_name}'. File does not exist and is not in the LoRA library." - ) - - -def get_registry() -> dict[str, Path]: - """Get a copy of the current LORA registry.""" - return _LORA_REGISTRY.copy() - - -# Initialize the registry when the module is imported -_initialize_registry() diff --git a/src/mflux/models/common/lora/mapping/lora_loader.py b/src/mflux/models/common/lora/mapping/lora_loader.py index ef42bc3..377c130 100644 --- a/src/mflux/models/common/lora/mapping/lora_loader.py +++ b/src/mflux/models/common/lora/mapping/lora_loader.py @@ -1,6 +1,7 @@ import re +from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path -from typing import Dict, Tuple import mlx.core as mx import mlx.nn as nn @@ -8,6 +9,16 @@ import mlx.nn as nn from mflux.models.common.lora.layer.fused_linear_lora_layer import FusedLoRALinear from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear from mflux.models.common.lora.mapping.lora_mapping import LoRATarget +from mflux.models.common.resolution.lora_resolution import LoraResolution + + +@dataclass +class PatternMatch: + source_pattern: str + target_path: str + matrix_name: str # "lora_A", "lora_B", or "alpha" + transpose: bool + transform: Callable[[mx.array], mx.array] | None = None class LoRALoader: @@ -15,30 +26,34 @@ class LoRALoader: def load_and_apply_lora( lora_mapping: list[LoRATarget], transformer: nn.Module, - lora_files: list[str], + lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - ) -> None: - if not lora_files: - return + ) -> tuple[list[str], list[float]]: + resolved_paths = LoraResolution.resolve_paths(lora_paths) + if not resolved_paths: + return resolved_paths, [] - # Validate scales - handle both None and empty list cases - if lora_scales is None or len(lora_scales) == 0: - lora_scales = [1.0] * len(lora_files) - elif len(lora_scales) != len(lora_files): + resolved_scales = LoraResolution.resolve_scales(lora_scales, len(resolved_paths)) + if len(resolved_scales) != len(resolved_paths): raise ValueError( - f"Number of LoRA scales ({len(lora_scales)}) must match number of LoRA files ({len(lora_files)})" + f"Number of LoRA scales ({len(resolved_scales)}) must match number of LoRA files ({len(resolved_paths)})" ) - print(f"๐Ÿ“ฆ Loading {len(lora_files)} LoRA file(s)...") + print(f"๐Ÿ“ฆ Loading {len(resolved_paths)} LoRA file(s)...") - for lora_file, scale in zip(lora_files, lora_scales): + for lora_file, scale in zip(resolved_paths, resolved_scales): LoRALoader._apply_single_lora(transformer, lora_file, scale, lora_mapping) print("โœ… All LoRA weights applied successfully") + return resolved_paths, resolved_scales + @staticmethod def _apply_single_lora( - transformer: nn.Module, lora_file: str, scale: float, lora_mapping: list[LoRATarget] + transformer: nn.Module, + lora_file: str, + scale: float, + lora_mapping: list[LoRATarget], ) -> None: # Load the LoRA weights if not Path(lora_file).exists(): @@ -49,73 +64,136 @@ class LoRALoader: try: weights = dict(mx.load(lora_file, return_metadata=True)[0].items()) - mx.eval(weights) except (FileNotFoundError, ValueError, RuntimeError) as e: print(f"โŒ Failed to load LoRA file: {e}") return - # Apply LoRA using the provided mapping - flat_mapping = LoRALoader._get_flat_mapping(lora_mapping) - applied_count = LoRALoader._apply_lora_with_mapping(transformer, weights, scale, flat_mapping) + # Build pattern mappings from LoRATargets + pattern_mappings = LoRALoader._build_pattern_mappings(lora_mapping) - print(f" โœ… Applied to {applied_count} layers") + # Apply LoRA using the mappings (allows multiple targets per source) + applied_count, matched_keys = LoRALoader._apply_lora_with_mapping(transformer, weights, scale, pattern_mappings) + + # Report results + total_keys = len(weights) + unmatched_keys = set(weights.keys()) - matched_keys + + print(f" โœ… Applied to {applied_count} layers ({len(matched_keys)}/{total_keys} keys matched)") + + if unmatched_keys: + print(f" โš ๏ธ {len(unmatched_keys)} unmatched keys in LoRA file:") + for key in sorted(unmatched_keys)[:5]: + print(f" - {key}") + if len(unmatched_keys) > 5: + print(f" ... and {len(unmatched_keys) - 5} more") + + @staticmethod + def _build_pattern_mappings(targets: list[LoRATarget]) -> list[PatternMatch]: + mappings = [] + + for target in targets: + # Add up weight patterns (lora_B) + mappings.extend( + PatternMatch( + source_pattern=pattern, + target_path=target.model_path, + matrix_name="lora_B", + transpose=True, + transform=target.up_transform, + ) + for pattern in target.possible_up_patterns + ) + + # Add down weight patterns (lora_A) + mappings.extend( + PatternMatch( + source_pattern=pattern, + target_path=target.model_path, + matrix_name="lora_A", + transpose=True, + transform=target.down_transform, + ) + for pattern in target.possible_down_patterns + ) + + # Add alpha patterns (no transpose, no transform) + mappings.extend( + PatternMatch( + source_pattern=pattern, + target_path=target.model_path, + matrix_name="alpha", + transpose=False, + transform=None, + ) + for pattern in target.possible_alpha_patterns + ) + + return mappings @staticmethod def _apply_lora_with_mapping( - transformer: nn.Module, weights: dict, scale: float, lora_mappings: Dict[str, Tuple[str, str, bool]] - ) -> int: + transformer: nn.Module, + weights: dict, + scale: float, + pattern_mappings: list[PatternMatch], + ) -> tuple[int, set]: applied_count = 0 - lora_data_by_target = {} + lora_data_by_target: dict[str, dict] = {} + matched_keys: set[str] = set() - # Group LoRA weights by their target layers + # For each weight key, find ALL matching patterns (not just first) + # This allows multiple targets to use the same source (e.g., QKV split) for weight_key, weight_value in weights.items(): - found_mapping = None - block_idx = None + for mapping in pattern_mappings: + match_result = LoRALoader._match_pattern(weight_key, mapping.source_pattern) + if match_result is None: + continue - # Pattern matching logic - for pattern, mapping_info in lora_mappings.items(): - if "{block}" in pattern: - # Extract block number from the weight key - try both . and _ separators - # This handles both standard LoRA formats (dot-separated) and other formats (underscore-separated) - # Find all numbers in the weight key - numbers_in_key = re.findall(r"\d+", weight_key) - for num_str in numbers_in_key: - try: - test_block_idx = int(num_str) - concrete_pattern = pattern.format(block=test_block_idx) - if weight_key == concrete_pattern: - found_mapping = mapping_info - block_idx = test_block_idx - break - except (ValueError, KeyError): - continue - if found_mapping: - break - else: - if weight_key == pattern: - found_mapping = mapping_info - break + matched_keys.add(weight_key) + block_idx = match_result - if found_mapping is None: - continue + # Resolve target path with block index if needed + target_path = mapping.target_path + if block_idx is not None and "{block}" in target_path: + target_path = target_path.format(block=block_idx) - target_path, matrix_name, transpose = found_mapping + # Apply transform if specified + transformed_value = weight_value + if mapping.transform is not None: + transformed_value = mapping.transform(weight_value) - # Handle block substitution in target path - if block_idx is not None and "{block}" in target_path: - target_path = target_path.format(block=block_idx) + # Apply transpose if needed + if mapping.transpose: + transformed_value = transformed_value.T - if target_path not in lora_data_by_target: - lora_data_by_target[target_path] = {} + # Store for this target + if target_path not in lora_data_by_target: + lora_data_by_target[target_path] = {} - lora_data_by_target[target_path][matrix_name] = (weight_value, transpose) + lora_data_by_target[target_path][mapping.matrix_name] = transformed_value # Apply LoRA to each target for target_path, lora_data in lora_data_by_target.items(): if LoRALoader._apply_lora_matrices_to_target(transformer, target_path, lora_data, scale): applied_count += 1 - return applied_count + return applied_count, matched_keys + + @staticmethod + def _match_pattern(weight_key: str, pattern: str) -> int | None: + if "{block}" in pattern: + # Find all numbers in the weight key + numbers_in_key = re.findall(r"\d+", weight_key) + for num_str in numbers_in_key: + test_block_idx = int(num_str) + concrete_pattern = pattern.replace("{block}", str(test_block_idx)) + if weight_key == concrete_pattern: + return test_block_idx + return None + else: + if weight_key == pattern: + return 0 # Return 0 to indicate match (no block) + return None @staticmethod def _apply_lora_matrices_to_target(transformer: nn.Module, target_path: str, lora_data: dict, scale: float) -> bool: @@ -138,19 +216,14 @@ class LoRALoader: print(f"โŒ Missing required LoRA matrices for {target_path}") return False - lora_A, transpose_A = lora_data["lora_A"] - lora_B, transpose_B = lora_data["lora_B"] - - # Handle transposition - if transpose_A: - lora_A = lora_A.T - if transpose_B: - lora_B = lora_B.T + # Values are already transformed and transposed + lora_A = lora_data["lora_A"] + lora_B = lora_data["lora_B"] # Handle alpha scaling alpha_scale = 1.0 if "alpha" in lora_data: - alpha_value, _ = lora_data["alpha"] + alpha_value = lora_data["alpha"] rank = lora_A.shape[1] alpha_scale = float(alpha_value) / rank @@ -167,47 +240,33 @@ class LoRALoader: # Handle fusion: if the current module is already a LoRA layer, fuse them if is_lora_linear: print(f" ๐Ÿ”€ Fusing with existing LoRA at {target_path}") - # Create a temporary LoRA layer from the base linear of the existing LoRA lora_layer = LoRALinear.from_linear(current_module.linear, r=lora_A.shape[1], scale=effective_scale) - # Set the LoRA matrices lora_layer.lora_A = lora_A lora_layer.lora_B = lora_B - # Apply alpha scaling to the matrices if present if "alpha" in lora_data: lora_layer.lora_B = lora_layer.lora_B * alpha_scale - - # Create fused layer with the existing LoRA and the new one fused_layer = FusedLoRALinear(base_linear=current_module.linear, loras=[current_module, lora_layer]) replacement_layer = fused_layer elif is_fused_linear: print(f" ๐Ÿ”€ Adding to existing fusion at {target_path}") - # Create a temporary LoRA layer from the base linear lora_layer = LoRALinear.from_linear( current_module.base_linear, r=lora_A.shape[1], scale=effective_scale ) - # Set the LoRA matrices lora_layer.lora_A = lora_A lora_layer.lora_B = lora_B - # Apply alpha scaling to the matrices if present if "alpha" in lora_data: lora_layer.lora_B = lora_layer.lora_B * alpha_scale - - # Add to existing fusion fused_layer = FusedLoRALinear( base_linear=current_module.base_linear, loras=current_module.loras + [lora_layer] ) replacement_layer = fused_layer else: # First LoRA on this layer - # Create LoRA layer lora_layer = LoRALinear.from_linear(current_module, r=lora_A.shape[1], scale=effective_scale) - # Set the LoRA matrices - use the correct dimensions from the LoRA file lora_layer.lora_A = lora_A lora_layer.lora_B = lora_B - # Apply alpha scaling to the matrices if present if "alpha" in lora_data: lora_layer.lora_B = lora_layer.lora_B * alpha_scale - replacement_layer = lora_layer # Replace the layer in the parent module @@ -228,22 +287,3 @@ class LoRALoader: else: print(f"โŒ Target layer {target_path} is not a linear layer") return False - - @staticmethod - def _get_flat_mapping(targets: list[LoRATarget]) -> Dict[str, Tuple[str, str, bool]]: - flat_mapping = {} - - for target in targets: - # Add up weight patterns (lora_B, transposed) - for pattern in target.possible_up_patterns: - flat_mapping[pattern] = (target.model_path, "lora_B", True) - - # Add down weight patterns (lora_A, transposed) - for pattern in target.possible_down_patterns: - flat_mapping[pattern] = (target.model_path, "lora_A", True) - - # Add alpha patterns (no transpose) - for pattern in target.possible_alpha_patterns: - flat_mapping[pattern] = (target.model_path, "alpha", False) - - return flat_mapping diff --git a/src/mflux/models/common/lora/mapping/lora_mapping.py b/src/mflux/models/common/lora/mapping/lora_mapping.py index 8dde37d..5cdf5ce 100644 --- a/src/mflux/models/common/lora/mapping/lora_mapping.py +++ b/src/mflux/models/common/lora/mapping/lora_mapping.py @@ -1,17 +1,21 @@ -from dataclasses import dataclass +from collections.abc import Callable +from dataclasses import dataclass, field from typing import List, Protocol +import mlx.core as mx + @dataclass class LoRATarget: model_path: str possible_up_patterns: List[str] possible_down_patterns: List[str] - possible_alpha_patterns: List[str] + possible_alpha_patterns: List[str] = field(default_factory=list) + up_transform: Callable[[mx.array], mx.array] | None = None + down_transform: Callable[[mx.array], mx.array] | None = None class LoRAMapping(Protocol): - @staticmethod def get_mapping() -> List[LoRATarget]: return diff --git a/src/mflux/models/common/lora/mapping/lora_transforms.py b/src/mflux/models/common/lora/mapping/lora_transforms.py new file mode 100644 index 0000000..b151133 --- /dev/null +++ b/src/mflux/models/common/lora/mapping/lora_transforms.py @@ -0,0 +1,101 @@ +import mlx.core as mx + + +class LoraTransforms: + @staticmethod + def split_q_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_up(tensor, 0) + + @staticmethod + def split_k_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_up(tensor, 1) + + @staticmethod + def split_v_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_up(tensor, 2) + + @staticmethod + def split_q_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_down(tensor, 0) + + @staticmethod + def split_k_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_down(tensor, 1) + + @staticmethod + def split_v_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_down(tensor, 2) + + @staticmethod + def split_single_q_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_up(tensor, 0) + + @staticmethod + def split_single_k_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_up(tensor, 1) + + @staticmethod + def split_single_v_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_up(tensor, 2) + + @staticmethod + def split_single_mlp_up(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_up(tensor, 3) + + @staticmethod + def split_single_q_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_down(tensor, 0) + + @staticmethod + def split_single_k_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_down(tensor, 1) + + @staticmethod + def split_single_v_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_down(tensor, 2) + + @staticmethod + def split_single_mlp_down(tensor: mx.array) -> mx.array: + return LoraTransforms._split_qkv_mlp_down(tensor, 3) + + @staticmethod + def _transpose(tensor: mx.array) -> mx.array: + return tensor.T + + @staticmethod + def _split_qkv_up(tensor: mx.array, index: int, num_splits: int = 3) -> mx.array: + split_size = tensor.shape[0] // num_splits + start = index * split_size + end = start + split_size + return tensor[start:end, :] + + @staticmethod + def _split_qkv_down(tensor: mx.array, index: int, num_splits: int = 3) -> mx.array: + rank = tensor.shape[0] + if rank % num_splits == 0: + chunk_size = rank // num_splits + start = index * chunk_size + end = start + chunk_size + return tensor[start:end, :] + else: + return tensor + + @staticmethod + def _split_qkv_mlp_up(tensor: mx.array, index: int, dims: list[int] | None = None) -> mx.array: + if dims is None: + dims = [3072, 3072, 3072, 12288] + + start = sum(dims[:index]) + end = start + dims[index] + return tensor[start:end, :] + + @staticmethod + def _split_qkv_mlp_down(tensor: mx.array, index: int, num_splits: int = 4) -> mx.array: + rank = tensor.shape[0] + if rank % num_splits == 0: + chunk_size = rank // num_splits + start = index * chunk_size + end = start + chunk_size + return tensor[start:end, :] + else: + return tensor diff --git a/src/mflux/models/common/quantization/quantization_util.py b/src/mflux/models/common/quantization/quantization_util.py deleted file mode 100644 index 9cb93a4..0000000 --- a/src/mflux/models/common/quantization/quantization_util.py +++ /dev/null @@ -1,112 +0,0 @@ -from typing import TYPE_CHECKING - -import mlx.nn as nn - -if TYPE_CHECKING: - from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler - from mflux.models.flux.variants.controlnet.weight_handler_controlnet import WeightHandlerControlnet - from mflux.models.flux.weights.weight_handler import WeightHandler - from mflux.models.qwen.weights.qwen_weight_handler import QwenWeightHandler - - -class QuantizationUtil: - @staticmethod - def quantize_model( - vae: nn.Module, - transformer: nn.Module, - t5_text_encoder: nn.Module, - clip_text_encoder: nn.Module, - quantize: int, - weights: "WeightHandler", - ) -> None: - q_level = weights.meta_data.quantization_level - if q_level == "None": - q_level = None - if quantize == "None": - quantize = None - - if quantize is not None or q_level is not None: - bits = int(q_level) if q_level is not None else quantize - nn.quantize(vae, bits=bits) - nn.quantize(transformer, bits=bits) - nn.quantize(t5_text_encoder, bits=bits) - nn.quantize(clip_text_encoder, bits=bits) - - @staticmethod - def quantize_controlnet( - quantize: int, - weights: "WeightHandlerControlnet", - transformer_controlnet: nn.Module, - ) -> None: - q_level = weights.meta_data.quantization_level - - if quantize is not None or q_level is not None: - bits = int(q_level) if q_level is not None else quantize - nn.quantize(transformer_controlnet, bits=bits) - - @staticmethod - def quantize_redux_models( - quantize: int, - weights: "WeightHandler", - redux_encoder: nn.Module, - siglip_vision_transformer: nn.Module, - ) -> None: - q_level = weights.meta_data.quantization_level - - if quantize is not None or q_level is not None: - bits = int(q_level) if q_level is not None else quantize - nn.quantize(redux_encoder, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) - nn.quantize(siglip_vision_transformer, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) - - @staticmethod - def quantization_predicate(path, m): - # 1. Skip Conv2d layers - if isinstance(m, nn.Conv2d): - return False - - # 2. Skip any layer with incompatible dimensions - if hasattr(m, "weight") and hasattr(m.weight, "shape"): - if m.weight.shape == (1152, 4304): - return False - - if m.weight.shape[-1] % 64 != 0: - return False - - # Only quantize layers that have to_quantized method - return hasattr(m, "to_quantized") - - @staticmethod - def quantize_qwen_models( - text_encoder: nn.Module, - vae: nn.Module, - transformer: nn.Module, - quantize: int, - weights: "QwenWeightHandler", - ) -> None: - q_level = weights.meta_data.quantization_level - - if quantize is not None or q_level is not None: - bits = int(q_level) if q_level is not None else quantize - nn.quantize(vae, bits=bits) - nn.quantize(transformer, bits=bits) - # nn.quantize(text_encoder, bits=bits) # Quantization of text encoder causes significant semantic degradation - - @staticmethod - def quantize_fibo_models( - vae: nn.Module, - transformer: nn.Module, - text_encoder: nn.Module | None, - quantize: int, - weights: "FIBOWeightHandler", - ) -> None: - q_level = weights.meta_data.quantization_level - if q_level == "None": - q_level = None - if quantize == "None": - quantize = None - - if quantize is not None or q_level is not None: - bits = int(q_level) if q_level is not None else quantize - nn.quantize(vae, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) - nn.quantize(transformer, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) - nn.quantize(text_encoder, class_predicate=QuantizationUtil.quantization_predicate, bits=bits) diff --git a/src/mflux/models/common/resolution/__init__.py b/src/mflux/models/common/resolution/__init__.py new file mode 100644 index 0000000..7410ce7 --- /dev/null +++ b/src/mflux/models/common/resolution/__init__.py @@ -0,0 +1,6 @@ +from mflux.models.common.resolution.config_resolution import ConfigResolution +from mflux.models.common.resolution.lora_resolution import LoraResolution +from mflux.models.common.resolution.path_resolution import PathResolution +from mflux.models.common.resolution.quantization_resolution import QuantizationResolution + +__all__ = ["ConfigResolution", "LoraResolution", "PathResolution", "QuantizationResolution"] diff --git a/src/mflux/models/common/resolution/actions.py b/src/mflux/models/common/resolution/actions.py new file mode 100644 index 0000000..8d3b9c8 --- /dev/null +++ b/src/mflux/models/common/resolution/actions.py @@ -0,0 +1,39 @@ +from enum import Enum +from typing import NamedTuple + + +class QuantizationAction(Enum): + NONE = "No quantization - use full precision" + STORED = "Use quantization level stored in weights" + REQUESTED = "Apply requested quantization on-the-fly" + + +class PathAction(Enum): + LOCAL = "Use local filesystem path" + HUGGINGFACE_CACHED = "Use cached HuggingFace files (no network)" + HUGGINGFACE = "Download from HuggingFace" + ERROR = "Path not found" + + +class LoraAction(Enum): + LOCAL = "Use local filesystem path" + REGISTRY = "Lookup in LORA_LIBRARY_PATH registry" + HUGGINGFACE_COLLECTION_CACHED = "Use cached file from HuggingFace collection (no network)" + HUGGINGFACE_COLLECTION = "Download specific file from HuggingFace collection" + HUGGINGFACE_REPO_CACHED = "Use cached HuggingFace repository (no network)" + HUGGINGFACE_REPO = "Download from HuggingFace repository" + ERROR = "LoRA not found" + + +class ConfigAction(Enum): + EXACT_MATCH = "Model name exactly matches a known alias" + EXPLICIT_BASE = "Use explicitly provided --base-model" + INFER_SUBSTRING = "Infer base model from substring match" + ERROR = "Cannot determine model configuration" + + +class Rule(NamedTuple): + priority: int + name: str + check: str + action: QuantizationAction | PathAction | LoraAction | ConfigAction diff --git a/src/mflux/models/common/resolution/config_resolution.py b/src/mflux/models/common/resolution/config_resolution.py new file mode 100644 index 0000000..d4c77db --- /dev/null +++ b/src/mflux/models/common/resolution/config_resolution.py @@ -0,0 +1,126 @@ +import logging +from typing import TYPE_CHECKING + +from mflux.models.common.resolution.actions import ConfigAction, Rule + +if TYPE_CHECKING: + from mflux.models.common.config.model_config import ModelConfig + +logger = logging.getLogger(__name__) + + +class ConfigResolution: + RULES = frozenset( + { + Rule(priority=0, name="exact_match", check="is_exact_match", action=ConfigAction.EXACT_MATCH), + Rule(priority=1, name="explicit_base", check="has_explicit_base", action=ConfigAction.EXPLICIT_BASE), + Rule(priority=2, name="infer_substring", check="can_infer_substring", action=ConfigAction.INFER_SUBSTRING), + Rule(priority=3, name="error", check="always", action=ConfigAction.ERROR), + } + ) + + @staticmethod + def resolve(model_name: str, base_model: str | None = None) -> "ModelConfig": + from mflux.models.common.config.model_config import AVAILABLE_MODELS, ModelConfig + from mflux.utils.exceptions import InvalidBaseModel, ModelConfigError + + base_models = sorted( + [m for m in AVAILABLE_MODELS.values() if m.base_model is None], + key=lambda x: x.priority, + ) + + ctx = { + "model_name": model_name, + "base_model": base_model, + "base_models": base_models, + "ModelConfig": ModelConfig, + "InvalidBaseModel": InvalidBaseModel, + "ModelConfigError": ModelConfigError, + } + + for rule in sorted(ConfigResolution.RULES, key=lambda r: r.priority): + if ConfigResolution._check(rule.check, ctx): + logger.debug(f"Config resolution: '{model_name}' โ†’ rule '{rule.name}' ({rule.action.value})") + return ConfigResolution._execute(rule.action, ctx) + + raise ValueError(f"No rule matched for model_name: {model_name}") + + @staticmethod + def _check(check: str, ctx: dict) -> bool: + if check == "is_exact_match": + model_name = ctx["model_name"] + for base in ctx["base_models"]: + if model_name == base.model_name or model_name in base.aliases: + return True + return False + if check == "has_explicit_base": + return ctx["base_model"] is not None + if check == "can_infer_substring": + model_name_lower = ctx["model_name"].lower() + for base in ctx["base_models"]: + for alias in base.aliases: + if alias and alias.lower() in model_name_lower: + return True + return False + if check == "always": + return True + return False + + @staticmethod + def _execute(action: ConfigAction, ctx: dict) -> "ModelConfig": + model_name = ctx["model_name"] + base_model = ctx["base_model"] + base_models = ctx["base_models"] + ModelConfig = ctx["ModelConfig"] + InvalidBaseModel = ctx["InvalidBaseModel"] + ModelConfigError = ctx["ModelConfigError"] + + if action == ConfigAction.EXACT_MATCH: + for base in base_models: + if model_name == base.model_name or model_name in base.aliases: + return base + raise ValueError("Exact match check passed but no match found") + + if action == ConfigAction.EXPLICIT_BASE: + allowed_names = [] + for base in base_models: + allowed_names.extend(base.aliases + [base.model_name]) + if base_model not in allowed_names: + raise InvalidBaseModel(f"Invalid base_model. Choose one of {allowed_names}") + + default_base = next( + (b for b in base_models if base_model == b.model_name or base_model in b.aliases), + None, + ) + return ConfigResolution._create_config(model_name, default_base, ModelConfig) + + if action == ConfigAction.INFER_SUBSTRING: + model_name_lower = model_name.lower() + matching_bases = [ + (b, alias) for b in base_models for alias in b.aliases if alias and alias.lower() in model_name_lower + ] + if not matching_bases: + raise ModelConfigError(f"Cannot infer base_model from {model_name}") + + default_base = sorted(matching_bases, key=lambda x: (-len(x[1]), x[0].priority))[0][0] + return ConfigResolution._create_config(model_name, default_base, ModelConfig) + + if action == ConfigAction.ERROR: + raise ModelConfigError(f"Cannot infer base_model from {model_name}") + + raise ValueError(f"Unknown action: {action}") + + @staticmethod + def _create_config(model_name: str, base: "ModelConfig", ModelConfig: type) -> "ModelConfig": + return ModelConfig( + aliases=base.aliases, + model_name=model_name, + base_model=base.model_name, + controlnet_model=base.controlnet_model, + custom_transformer_model=base.custom_transformer_model, + num_train_steps=base.num_train_steps, + max_sequence_length=base.max_sequence_length, + supports_guidance=base.supports_guidance, + requires_sigma_shift=base.requires_sigma_shift, + priority=base.priority, + ) diff --git a/src/mflux/models/common/resolution/lora_resolution.py b/src/mflux/models/common/resolution/lora_resolution.py new file mode 100644 index 0000000..7d2e306 --- /dev/null +++ b/src/mflux/models/common/resolution/lora_resolution.py @@ -0,0 +1,298 @@ +import logging +import os +import shutil +from pathlib import Path + +from huggingface_hub import snapshot_download +from huggingface_hub.utils import LocalEntryNotFoundError + +from mflux.cli.defaults.defaults import MFLUX_LORA_CACHE_DIR +from mflux.models.common.resolution.actions import LoraAction, Rule + +logger = logging.getLogger(__name__) + + +class LoraResolution: + RULES = frozenset( + { + Rule(priority=0, name="local", check="exists_locally", action=LoraAction.LOCAL), + Rule(priority=1, name="registry", check="in_registry", action=LoraAction.REGISTRY), + Rule(priority=2, name="collection_cached", check="is_collection_cached", action=LoraAction.HUGGINGFACE_COLLECTION_CACHED), + Rule(priority=3, name="collection_download", check="is_collection_format", action=LoraAction.HUGGINGFACE_COLLECTION), + Rule(priority=4, name="repo_cached", check="is_repo_cached", action=LoraAction.HUGGINGFACE_REPO_CACHED), + Rule(priority=5, name="repo_download", check="is_hf_format", action=LoraAction.HUGGINGFACE_REPO), + Rule(priority=6, name="error", check="always", action=LoraAction.ERROR), + } + ) # fmt: off + + _registry: dict[str, Path] = {} + + @staticmethod + def resolve(path: str) -> str: + for rule in sorted(LoraResolution.RULES, key=lambda r: r.priority): + if LoraResolution._check(rule.check, path): + logger.debug(f"LoRA resolution: '{path}' โ†’ rule '{rule.name}' ({rule.action.value})") + return LoraResolution._execute(rule.action, path) + + raise ValueError(f"No rule matched for LoRA path: {path}") + + @staticmethod + def resolve_paths(paths: list[str] | None) -> list[str]: + if not paths: + return [] + return [r for path in paths if (r := LoraResolution._try_resolve(path)) is not None] + + @staticmethod + def _try_resolve(path: str) -> str | None: + try: + return LoraResolution.resolve(path) + except FileNotFoundError as e: + print(f"โš ๏ธ {e}") + return None + + @staticmethod + def resolve_scales(scales: list[float] | None, num_paths: int) -> list[float]: + if not scales: + return [1.0] * num_paths + if len(scales) != num_paths: + print( + f"โš ๏ธ Number of LoRA scales ({len(scales)}) doesn't match number of LoRA paths ({num_paths}). " + f"Using provided scales and defaulting remaining to 1.0." + ) + # Pad with 1.0 if too few scales, truncate if too many + if len(scales) < num_paths: + return list(scales) + [1.0] * (num_paths - len(scales)) + return list(scales[:num_paths]) + return scales + + @staticmethod + def _is_collection_format(path: str) -> bool: + return ":" in path and "/" in path + + @staticmethod + def _is_hf_format(path: str) -> bool: + return "/" in path and path.count("/") == 1 and not path.startswith(("./", "../", "~/")) + + @staticmethod + def _check(check: str, path: str) -> bool: + if check == "exists_locally": + return Path(path).expanduser().exists() + if check == "in_registry": + return path in LoraResolution._registry + if check == "is_collection_cached": + if not LoraResolution._is_collection_format(path): + return False + repo_id, filename = path.split(":", 1) + return LoraResolution._is_collection_in_cache(repo_id, filename) + if check == "is_collection_format": + return LoraResolution._is_collection_format(path) + if check == "is_repo_cached": + if not LoraResolution._is_hf_format(path): + return False + return LoraResolution._is_repo_in_cache(path) + if check == "is_hf_format": + return LoraResolution._is_hf_format(path) + if check == "always": + return True + return False + + @staticmethod + def _is_collection_in_cache(repo_id: str, filename: str) -> bool: + cache_path = MFLUX_LORA_CACHE_DIR + # Check mflux cache + cached_file_path = cache_path / filename + if cached_file_path.exists() and cached_file_path.is_file(): + return True + # Check HF cache + try: + snapshot_download( + repo_id=repo_id, + allow_patterns=[f"*{filename}*"], + cache_dir=str(cache_path), + local_files_only=True, + ) + return True + except LocalEntryNotFoundError: + return False + + @staticmethod + def _is_repo_in_cache(repo_id: str) -> bool: + cache_path = MFLUX_LORA_CACHE_DIR + try: + snapshot_download( + repo_id=repo_id, + allow_patterns=["*.safetensors"], + cache_dir=str(cache_path), + local_files_only=True, + ) + return True + except LocalEntryNotFoundError: + return False + + @staticmethod + def _execute(action: LoraAction, path: str) -> str: + if action == LoraAction.LOCAL: + return str(Path(path).expanduser()) + if action == LoraAction.REGISTRY: + return str(LoraResolution._registry[path]) + if action == LoraAction.HUGGINGFACE_COLLECTION_CACHED: + repo_id, filename = path.split(":", 1) + return LoraResolution._load_collection_from_cache(repo_id, filename) + if action == LoraAction.HUGGINGFACE_COLLECTION: + repo_id, filename = path.split(":", 1) + return LoraResolution._download_collection(repo_id, filename) + if action == LoraAction.HUGGINGFACE_REPO_CACHED: + return LoraResolution._load_repo_from_cache(path) + if action == LoraAction.HUGGINGFACE_REPO: + return LoraResolution._download_repo(path) + if action == LoraAction.ERROR: + raise FileNotFoundError( + f"LoRA file not found: '{path}'. File does not exist and is not in the LoRA library." + ) + raise ValueError(f"Unknown action: {action}") + + @staticmethod + def _load_repo_from_cache(repo_id: str) -> str: + cache_path = MFLUX_LORA_CACHE_DIR + cache_path.mkdir(parents=True, exist_ok=True) + + download_path = Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=["*.safetensors"], + cache_dir=str(cache_path), + local_files_only=True, + ) + ) + + safetensor_files = list(download_path.glob("*.safetensors")) + if not safetensor_files: + raise FileNotFoundError(f"No .safetensors file found in cached repo: {repo_id}") + + if len(safetensor_files) > 1: + file_names = [f.name for f in safetensor_files] + raise ValueError( + f"Multiple .safetensors files found in '{repo_id}': {file_names}. " + f"Please specify which file to use with the collection format: '{repo_id}:.safetensors'" + ) + + return str(safetensor_files[0]) + + @staticmethod + def _download_repo(repo_id: str) -> str: + cache_path = MFLUX_LORA_CACHE_DIR + cache_path.mkdir(parents=True, exist_ok=True) + + print(f"Downloading LoRA from HuggingFace: {repo_id}...") + download_path = Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=["*.safetensors"], + cache_dir=str(cache_path), + ) + ) + + safetensor_files = list(download_path.glob("*.safetensors")) + if not safetensor_files: + raise FileNotFoundError(f"No .safetensors file found in HuggingFace repo: {repo_id}") + + if len(safetensor_files) > 1: + file_names = [f.name for f in safetensor_files] + raise ValueError( + f"Multiple .safetensors files found in '{repo_id}': {file_names}. " + f"Please specify which file to use with the collection format: '{repo_id}:.safetensors'" + ) + + return str(safetensor_files[0]) + + @staticmethod + def _load_collection_from_cache(repo_id: str, filename: str) -> str: + cache_path = MFLUX_LORA_CACHE_DIR + cache_path.mkdir(parents=True, exist_ok=True) + + # Check mflux cache first + cached_file_path = cache_path / filename + if cached_file_path.exists() and cached_file_path.is_file(): + try: + with open(cached_file_path, "rb") as f: + f.read(1) + return str(cached_file_path) + except (OSError, IOError): + # File corrupted, fall through to HF cache + pass + + # Load from HF cache + download_path = Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=[f"*{filename}*"], + cache_dir=str(cache_path), + local_files_only=True, + ) + ) + + return LoraResolution._find_and_link_file(download_path, filename, cache_path) + + @staticmethod + def _download_collection(repo_id: str, filename: str) -> str: + cache_path = MFLUX_LORA_CACHE_DIR + cache_path.mkdir(parents=True, exist_ok=True) + + print(f"Downloading LoRA '{filename}' from {repo_id}...") + download_path = Path( + snapshot_download( + repo_id=repo_id, + allow_patterns=[f"*{filename}*"], + cache_dir=str(cache_path), + ) + ) + + return LoraResolution._find_and_link_file(download_path, filename, cache_path) + + @staticmethod + def _find_and_link_file(download_path: Path, filename: str, cache_path: Path) -> str: + found_files = list(download_path.glob(f"**/*{filename}*")) + for file in found_files: + if file.is_file() and file.suffix in [".safetensors", ".bin"]: + if not filename.endswith(file.suffix): + target_name = f"{filename}{file.suffix}" + else: + target_name = filename + + target_path = cache_path / target_name + if not target_path.exists(): + try: + target_path.symlink_to(file) + except (OSError, AttributeError): + shutil.copy2(file, target_path) + + return str(target_path) + + raise FileNotFoundError(f"Could not find LoRA file '{filename}' in downloaded files") + + @staticmethod + def get_registry() -> dict[str, Path]: + return LoraResolution._registry.copy() + + @staticmethod + def discover_files(library_paths: list[Path]) -> dict[str, Path]: + lora_files = {} + for library_path in reversed(library_paths): + if not library_path.exists() or not library_path.is_dir(): + continue + for safetensor_path in library_path.rglob("*.safetensors"): + basename = safetensor_path.stem + if basename.isdigit() and safetensor_path.parent.name == "transformer": + continue + lora_files[basename] = safetensor_path.resolve() + return lora_files + + @staticmethod + def _initialize_registry() -> None: + library_path_env = os.environ.get("LORA_LIBRARY_PATH") + if library_path_env: + library_paths = [Path(p.strip()) for p in library_path_env.split(":") if p.strip()] + LoraResolution._registry = LoraResolution.discover_files(library_paths) + + +LoraResolution._initialize_registry() diff --git a/src/mflux/models/common/resolution/path_resolution.py b/src/mflux/models/common/resolution/path_resolution.py new file mode 100644 index 0000000..663aef9 --- /dev/null +++ b/src/mflux/models/common/resolution/path_resolution.py @@ -0,0 +1,157 @@ +import logging +import os +from pathlib import Path + +from huggingface_hub import snapshot_download +from huggingface_hub.constants import HF_HUB_CACHE + +from mflux.models.common.resolution.actions import PathAction, Rule + +logger = logging.getLogger(__name__) + + +class PathResolution: + RULES = frozenset( + { + Rule(priority=0, name="none", check="is_none", action=PathAction.LOCAL), + Rule(priority=1, name="local", check="exists_locally", action=PathAction.LOCAL), + Rule(priority=2, name="hf_cached", check="is_hf_cached", action=PathAction.HUGGINGFACE_CACHED), + Rule(priority=3, name="hf_download", check="is_hf_format", action=PathAction.HUGGINGFACE), + Rule(priority=4, name="error", check="always", action=PathAction.ERROR), + } + ) + + @staticmethod + def resolve(path: str | None, patterns: list[str] | None = None) -> Path | None: + if patterns is None: + patterns = ["*.safetensors"] + + for rule in sorted(PathResolution.RULES, key=lambda r: r.priority): + if PathResolution._check(rule.check, path, patterns): + logger.debug(f"Path resolution: '{path}' โ†’ rule '{rule.name}' ({rule.action.value})") + return PathResolution._execute(rule.action, path, patterns) + + raise ValueError(f"No rule matched for path: {path}") + + @staticmethod + def _is_hf_format(path: str | None) -> bool: + return path is not None and "/" in path and path.count("/") == 1 and not path.startswith(("./", "../", "~/")) + + @staticmethod + def _check(check: str, path: str | None, patterns: list[str]) -> bool: + if check == "is_none": + return path is None + if check == "exists_locally": + if path is None: + return False + local_path = Path(path).expanduser() + if not local_path.exists(): + return False + # Warn if directory exists but contains no matching files + if local_path.is_dir(): + has_matching_files = any(list(local_path.glob(p)) for p in patterns) + if not has_matching_files: + print( + f"โš ๏ธ Directory '{path}' exists but contains no files matching {patterns}. " + f"Model loading may fail." + ) + return True + if check == "is_hf_cached": + if not PathResolution._is_hf_format(path): + return False + # Check if we have a complete cached snapshot + return PathResolution._find_complete_cached_snapshot(path, patterns) is not None + if check == "is_hf_format": + return PathResolution._is_hf_format(path) + if check == "always": + return True + return False + + @staticmethod + def _execute(action: PathAction, path: str | None, patterns: list[str]) -> Path | None: + if action == PathAction.LOCAL: + return Path(path).expanduser() if path else None + if action == PathAction.HUGGINGFACE_CACHED: + # Find the best complete cached snapshot + cached_path = PathResolution._find_complete_cached_snapshot(path, patterns) + if cached_path: + return cached_path + # Fallback to standard snapshot_download (shouldn't happen if _check passed) + return Path(snapshot_download(repo_id=path, allow_patterns=patterns, local_files_only=True)) + if action == PathAction.HUGGINGFACE: + print(f"Downloading model from HuggingFace: {path}...") + return Path(snapshot_download(repo_id=path, allow_patterns=patterns)) + if action == PathAction.ERROR: + raise FileNotFoundError( + f"Model not found: '{path}'. " + f"If local path, make sure it exists. " + f"If HuggingFace repo, use 'org/model' format." + ) + return None + + @staticmethod + def _find_complete_cached_snapshot(repo_id: str, patterns: list[str]) -> Path | None: + # Build the cache directory path for this repo + # HuggingFace cache structure: {cache_dir}/models--{org}--{model}/snapshots/{revision}/ + repo_cache_name = f"models--{repo_id.replace('/', '--')}" + repo_cache_dir = Path(HF_HUB_CACHE) / repo_cache_name / "snapshots" + + if not repo_cache_dir.exists(): + return None + + # Extract subdirectories that need safetensors files (e.g., "vae/*.safetensors" โ†’ "vae") + required_subdirs = PathResolution._get_required_subdirs_with_safetensors(patterns) + + # Check each snapshot for completeness, prefer more recent ones + snapshots = sorted(repo_cache_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True) + + for snapshot_path in snapshots: + if not snapshot_path.is_dir(): + continue + if PathResolution._is_snapshot_complete(snapshot_path, required_subdirs): + logger.debug(f"Found complete cached snapshot: {snapshot_path}") + return snapshot_path + + return None + + @staticmethod + def _get_required_subdirs_with_safetensors(patterns: list[str]) -> set[str]: + subdirs = set() + for pattern in patterns: + # Only care about safetensors patterns + if "*.safetensors" not in pattern: + continue + # Handle patterns like "vae/*.safetensors" + if "/" in pattern: + subdir = pattern.split("/")[0] + # Only add if it's a real subdir name (not a glob pattern itself) + if "*" not in subdir: + subdirs.add(subdir) + return subdirs + + @staticmethod + def _is_snapshot_complete(snapshot_path: Path, required_subdirs: set[str]) -> bool: + if not required_subdirs: + # No specific subdirs required, just check for any safetensors + return any(snapshot_path.glob("**/*.safetensors")) + + for subdir in required_subdirs: + subdir_path = snapshot_path / subdir + if not subdir_path.exists(): + return False + # Check if subdir has at least one safetensors file (following symlinks) + has_safetensors = False + for f in subdir_path.iterdir(): + if f.name.endswith(".safetensors"): + # Verify the symlink target exists (handles broken symlinks) + if f.is_symlink(): + if os.path.exists(f): + has_safetensors = True + break + else: + has_safetensors = True + break + if not has_safetensors: + return False + + return True diff --git a/src/mflux/models/common/resolution/quantization_resolution.py b/src/mflux/models/common/resolution/quantization_resolution.py new file mode 100644 index 0000000..40e2b71 --- /dev/null +++ b/src/mflux/models/common/resolution/quantization_resolution.py @@ -0,0 +1,52 @@ +import logging + +from mflux.models.common.resolution.actions import QuantizationAction, Rule + +logger = logging.getLogger(__name__) + + +class QuantizationResolution: + RULES = frozenset( + { + Rule(priority=0, name="none", check="none_none", action=QuantizationAction.NONE), + Rule(priority=1, name="on_the_fly", check="none_any", action=QuantizationAction.REQUESTED), + Rule(priority=2, name="pre_quantized", check="any_none", action=QuantizationAction.STORED), + Rule(priority=3, name="conflict", check="any_any", action=QuantizationAction.STORED), + } + ) + + @staticmethod + def resolve(stored: int | None, requested: int | None) -> tuple[int | None, str | None]: + for rule in sorted(QuantizationResolution.RULES, key=lambda r: r.priority): + if QuantizationResolution._check(rule.check, stored, requested): + logger.debug( + f"Quantization resolution: stored={stored}, requested={requested} " + f"โ†’ rule '{rule.name}' ({rule.action.value})" + ) + return QuantizationResolution._execute(rule, stored, requested) + + raise ValueError(f"Unexpected quantization state: stored={stored}, requested={requested}") + + @staticmethod + def _check(check: str, stored: int | None, requested: int | None) -> bool: + if check == "none_none": + return stored is None and requested is None + if check == "none_any": + return stored is None and requested is not None + if check == "any_none": + return stored is not None and requested is None + if check == "any_any": + return stored is not None and requested is not None + return False + + @staticmethod + def _execute(rule: Rule, stored: int | None, requested: int | None) -> tuple[int | None, str | None]: + if rule.action == QuantizationAction.NONE: + return None, None + if rule.action == QuantizationAction.REQUESTED: + return requested, None + if rule.action == QuantizationAction.STORED: + warn = rule.name == "conflict" and stored != requested + warning = f"Model is pre-quantized at {stored}-bit. Ignoring -q {requested} flag." if warn else None + return stored, warning + return None, None diff --git a/src/mflux/models/common/schedulers/base_scheduler.py b/src/mflux/models/common/schedulers/base_scheduler.py index ffacb7e..0fbe2d1 100644 --- a/src/mflux/models/common/schedulers/base_scheduler.py +++ b/src/mflux/models/common/schedulers/base_scheduler.py @@ -4,23 +4,12 @@ import mlx.core as mx class BaseScheduler(ABC): - """ - Abstract base class for all schedulers. - """ - @property @abstractmethod - def sigmas(self) -> mx.array: - """ - The sigma schedule for the diffusion process. - """ - ... + def sigmas(self) -> mx.array: ... @abstractmethod - def step(self, model_output: mx.array, timestep: int, sample: mx.array, **kwargs) -> mx.array: ... + def step(self, noise: mx.array, timestep: int, latents: mx.array, **kwargs) -> mx.array: ... def scale_model_input(self, latents: mx.array, t: int) -> mx.array: - """ - Scale the denoising model input. By default, no scaling applied. - """ return latents diff --git a/src/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.py b/src/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.py index 09b7fe4..56508de 100644 --- a/src/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.py +++ b/src/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.py @@ -4,15 +4,15 @@ from typing import TYPE_CHECKING import mlx.core as mx if TYPE_CHECKING: - from mflux.config.runtime_config import RuntimeConfig + from mflux.models.common.config.config import Config from mflux.models.common.schedulers.base_scheduler import BaseScheduler class FlowMatchEulerDiscreteScheduler(BaseScheduler): - def __init__(self, runtime_config: "RuntimeConfig"): - self.runtime_config = runtime_config - self.model_config = runtime_config.model_config + def __init__(self, config: "Config"): + self.config = config + self.model_config = config.model_config self.num_train_timesteps = 1000 self.shift_terminal = 0.02 self.base_shift = 0.5 @@ -30,8 +30,8 @@ class FlowMatchEulerDiscreteScheduler(BaseScheduler): return self._timesteps def _compute_mu(self) -> float: - h_patches = self.runtime_config.height // 16 - w_patches = self.runtime_config.width // 16 + h_patches = self.config.height // 16 + w_patches = self.config.width // 16 seq_len = h_patches * w_patches m = (self.max_shift - self.base_shift) / (self.max_image_seq_len - self.base_image_seq_len) b = self.base_shift - m * self.base_image_seq_len @@ -49,7 +49,7 @@ class FlowMatchEulerDiscreteScheduler(BaseScheduler): return stretched def _compute_timesteps_and_sigmas(self) -> tuple[mx.array, mx.array]: - num_steps = self.runtime_config.num_inference_steps + num_steps = self.config.num_inference_steps sigma_min = 1.0 / self.num_train_timesteps sigma_max = 1.0 timesteps_linear = [ @@ -66,10 +66,9 @@ class FlowMatchEulerDiscreteScheduler(BaseScheduler): timesteps_arr = mx.array(timesteps, dtype=mx.float32) return sigmas_arr, timesteps_arr - def step(self, model_output: mx.array, timestep: int, sample: mx.array, **kwargs) -> mx.array: + def step(self, noise: mx.array, timestep: int, latents: mx.array, **kwargs) -> mx.array: dt = self._sigmas[timestep + 1] - self._sigmas[timestep] - prev_sample = sample + dt * model_output - return prev_sample + return latents + dt * noise def scale_model_input(self, latents: mx.array, t: int) -> mx.array: return latents diff --git a/src/mflux/models/common/schedulers/linear_scheduler.py b/src/mflux/models/common/schedulers/linear_scheduler.py index 277f606..5d7187d 100644 --- a/src/mflux/models/common/schedulers/linear_scheduler.py +++ b/src/mflux/models/common/schedulers/linear_scheduler.py @@ -3,14 +3,14 @@ from typing import TYPE_CHECKING import mlx.core as mx if TYPE_CHECKING: - from mflux.config.runtime_config import RuntimeConfig + from mflux.models.common.config.config import Config from mflux.models.common.schedulers.base_scheduler import BaseScheduler class LinearScheduler(BaseScheduler): - def __init__(self, runtime_config: "RuntimeConfig"): - self.runtime_config = runtime_config + def __init__(self, config: "Config"): + self.config = config self._sigmas = self._get_sigmas() self._timesteps = self._get_timesteps() @@ -23,11 +23,11 @@ class LinearScheduler(BaseScheduler): return self._timesteps def _get_sigmas(self) -> mx.array: - model_config = self.runtime_config.model_config + model_config = self.config.model_config sigmas = mx.linspace( 1.0, - 1.0 / self.runtime_config.num_inference_steps, - self.runtime_config.num_inference_steps, + 1.0 / self.config.num_inference_steps, + self.config.num_inference_steps, ) sigmas = mx.array(sigmas).astype(mx.float32) sigmas = mx.concatenate([sigmas, mx.zeros(1)]) @@ -36,7 +36,7 @@ class LinearScheduler(BaseScheduler): x1 = 256 m = (1.15 - y1) / (4096 - x1) b = y1 - m * x1 - mu = m * self.runtime_config.width * self.runtime_config.height / 256 + b + mu = m * self.config.width * self.config.height / 256 + b mu = mx.array(mu) shifted_sigmas = mx.exp(mu) / (mx.exp(mu) + (1 / sigmas - 1)) shifted_sigmas[-1] = 0 @@ -45,11 +45,11 @@ class LinearScheduler(BaseScheduler): return sigmas def _get_timesteps(self) -> mx.array: - num_steps = self.runtime_config.num_inference_steps + num_steps = self.config.num_inference_steps timesteps = mx.arange(num_steps, dtype=mx.float32) return timesteps - def step(self, model_output: mx.array, timestep: int, sample: mx.array, **kwargs) -> mx.array: + def step(self, noise: mx.array, timestep: int, latents: mx.array, **kwargs) -> mx.array: dt = self._sigmas[timestep + 1] - self._sigmas[timestep] - return sample + model_output * dt + return latents + noise * dt diff --git a/src/mflux/models/common/tokenizer/__init__.py b/src/mflux/models/common/tokenizer/__init__.py new file mode 100644 index 0000000..e760f9d --- /dev/null +++ b/src/mflux/models/common/tokenizer/__init__.py @@ -0,0 +1,17 @@ +from mflux.models.common.tokenizer.tokenizer import ( + BaseTokenizer, + LanguageTokenizer, + Tokenizer, + VisionLanguageTokenizer, +) +from mflux.models.common.tokenizer.tokenizer_loader import TokenizerLoader +from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput + +__all__ = [ + "Tokenizer", + "BaseTokenizer", + "LanguageTokenizer", + "VisionLanguageTokenizer", + "TokenizerLoader", + "TokenizerOutput", +] diff --git a/src/mflux/models/common/tokenizer/tokenizer.py b/src/mflux/models/common/tokenizer/tokenizer.py new file mode 100644 index 0000000..7c1d6f6 --- /dev/null +++ b/src/mflux/models/common/tokenizer/tokenizer.py @@ -0,0 +1,188 @@ +from abc import ABC, abstractmethod +from typing import Protocol, runtime_checkable + +import mlx.core as mx +import numpy as np +from PIL import Image +from transformers import PreTrainedTokenizer + +from mflux.models.common.tokenizer.tokenizer_output import TokenizerOutput + + +@runtime_checkable +class Tokenizer(Protocol): + tokenizer: PreTrainedTokenizer + + def tokenize( + self, + prompt: str | list[str], + images: list[Image.Image] | None = None, + max_length: int | None = None, + **kwargs, + ) -> TokenizerOutput: ... + + +class BaseTokenizer(ABC): + def __init__(self, tokenizer: PreTrainedTokenizer, max_length: int = 512): + self.tokenizer = tokenizer + self.max_length = max_length + + @abstractmethod + def tokenize( + self, + prompt: str | list[str], + images: list[Image.Image] | None = None, + max_length: int | None = None, + **kwargs, + ) -> TokenizerOutput: ... + + +class LanguageTokenizer(BaseTokenizer): + def __init__( + self, + tokenizer: PreTrainedTokenizer, + max_length: int = 512, + padding: str = "max_length", + return_attention_mask: bool = True, + template: str | None = None, + use_chat_template: bool = False, + chat_template_kwargs: dict | None = None, + add_special_tokens: bool = True, + ): + super().__init__(tokenizer, max_length) + self.padding = padding + self.return_attention_mask = return_attention_mask + self.template = template + self.use_chat_template = use_chat_template + self.chat_template_kwargs = chat_template_kwargs or {} + self.add_special_tokens = add_special_tokens + + def tokenize( + self, + prompt: str | list[str], + images: list[Image.Image] | None = None, + max_length: int | None = None, + **kwargs, + ) -> TokenizerOutput: + max_length = max_length or self.max_length + + if isinstance(prompt, str): + prompts = [prompt] + else: + prompts = list(prompt) + + prompts = [p if p is not None else "" for p in prompts] + if all(p == "" for p in prompts): + batch_size = len(prompts) + input_ids = mx.array(np.empty((batch_size, 0), dtype=np.int32)) + attention_mask = mx.array(np.empty((batch_size, 0), dtype=np.int32)) + return TokenizerOutput(input_ids=input_ids, attention_mask=attention_mask) + + if self.template or self.use_chat_template: + formatted_prompts = [] + for p in prompts: + if self.template: + formatted = self.template.format(p) + elif self.use_chat_template: + formatted = self.tokenizer.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, + add_generation_prompt=True, + **self.chat_template_kwargs, + ) + else: + formatted = p + formatted_prompts.append(formatted) + prompts = formatted_prompts + + tokens = self.tokenizer( + prompts, + padding=self.padding, + max_length=max_length, + truncation=True, + add_special_tokens=self.add_special_tokens, + return_length=False, + return_overflowing_tokens=False, + return_tensors="np", + ) + + input_ids = mx.array(tokens["input_ids"]) + if self.return_attention_mask: + attention_mask = mx.array(tokens["attention_mask"]) + else: + attention_mask = mx.ones_like(input_ids) + + return TokenizerOutput( + input_ids=input_ids, + attention_mask=attention_mask, + ) + + +class VisionLanguageTokenizer(BaseTokenizer): + def __init__( + self, + tokenizer: PreTrainedTokenizer, + processor, + max_length: int = 1024, + template: str | None = None, + image_token: str = "<|image_pad|>", + ): + super().__init__(tokenizer, max_length) + self.processor = processor + self.template = template + self.image_token = image_token + + def tokenize( + self, + prompt: str | list[str], + images: list[Image.Image] | None = None, + max_length: int | None = None, + **kwargs, + ) -> TokenizerOutput: + max_length = max_length or self.max_length + + if isinstance(prompt, str): + prompt = [prompt] + + if self.template and images: + img_prompt = "" + for i in range(len(images)): + img_prompt += f"Picture {i + 1}: <|vision_start|>{self.image_token}<|vision_end|>" + formatted_text = self.template.format(img_prompt + prompt[0]) + elif self.template: + formatted_text = self.template.format(prompt[0]) + else: + formatted_text = prompt[0] + + pixel_values = None + image_grid_thw = None + + if images: + model_inputs = self.processor( + text=[formatted_text], + images=images, + padding=True, + return_tensors=None, + ) + + input_ids = model_inputs["input_ids"] + attention_mask = model_inputs["attention_mask"] + pixel_values = mx.array(model_inputs["pixel_values"]) + image_grid_thw = mx.array(model_inputs["image_grid_thw"]) + else: + tokens = self.tokenizer( + [formatted_text], + max_length=max_length, + padding=True, + truncation=True, + return_tensors="np", + ) + input_ids = mx.array(tokens["input_ids"]) + attention_mask = mx.array(tokens["attention_mask"]) + + return TokenizerOutput( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + ) diff --git a/src/mflux/models/common/tokenizer/tokenizer_loader.py b/src/mflux/models/common/tokenizer/tokenizer_loader.py new file mode 100644 index 0000000..319cf9c --- /dev/null +++ b/src/mflux/models/common/tokenizer/tokenizer_loader.py @@ -0,0 +1,144 @@ +from pathlib import Path +from typing import TYPE_CHECKING + +import transformers +from huggingface_hub import snapshot_download + +from mflux.models.common.tokenizer.tokenizer import ( + BaseTokenizer, + LanguageTokenizer, + VisionLanguageTokenizer, +) + +if TYPE_CHECKING: + from mflux.models.common.weights.loading.weight_definition import TokenizerDefinition + + +class TokenizerLoader: + @staticmethod + def load( + definition: "TokenizerDefinition", + model_path: str, + ) -> BaseTokenizer: + tokenizer_path = TokenizerLoader._resolve_path( + model_path=model_path, + hf_subdir=definition.hf_subdir, + fallback_subdirs=definition.fallback_subdirs, + download_patterns=definition.download_patterns, + ) + + raw_tokenizer = TokenizerLoader._load_raw_tokenizer( + tokenizer_path=tokenizer_path, + tokenizer_class=definition.tokenizer_class, + ) + + return TokenizerLoader._create_tokenizer( + raw_tokenizer=raw_tokenizer, + definition=definition, + ) + + @staticmethod + def load_all( + definitions: list["TokenizerDefinition"], + model_path: str, + max_length_overrides: dict[str, int] | None = None, + ) -> dict[str, BaseTokenizer]: + max_length_overrides = max_length_overrides or {} + result = {} + for d in definitions: + tokenizer = TokenizerLoader.load(d, model_path) + if d.name in max_length_overrides: + tokenizer.max_length = max_length_overrides[d.name] + result[d.name] = tokenizer + return result + + @staticmethod + def _resolve_path( + model_path: str, + hf_subdir: str, + fallback_subdirs: list[str] | None, + download_patterns: list[str] | None, + ) -> Path: + expanded = Path(model_path).expanduser() + if expanded.exists(): + root_path = expanded + elif "/" in model_path and model_path.count("/") == 1 and not model_path.startswith(("./", "../")): + patterns = download_patterns or [f"{hf_subdir}/**"] + root_path = Path( + snapshot_download( + repo_id=model_path, + allow_patterns=patterns, + ) + ) + else: + raise FileNotFoundError( + f"Model not found: '{model_path}'. " + f"If local path, make sure it exists. " + f"If HuggingFace repo, use 'org/model' format." + ) + + tokenizer_path = root_path / hf_subdir + if tokenizer_path.exists(): + return tokenizer_path + + if fallback_subdirs: + for subdir in fallback_subdirs: + if subdir == ".": + if TokenizerLoader._has_tokenizer_files(root_path): + return root_path + else: + fallback_path = root_path / subdir + if fallback_path.exists(): + return fallback_path + + return tokenizer_path + + @staticmethod + def _has_tokenizer_files(path: Path) -> bool: + tokenizer_indicators = ["vocab.json", "tokenizer.json", "tokenizer_config.json"] + return any((path / f).exists() for f in tokenizer_indicators) + + @staticmethod + def _load_raw_tokenizer( + tokenizer_path: Path, + tokenizer_class: str, + ): + if hasattr(transformers, tokenizer_class): + cls = getattr(transformers, tokenizer_class) + else: + raise ValueError(f"Unknown tokenizer class: {tokenizer_class}") + + return cls.from_pretrained( + pretrained_model_name_or_path=str(tokenizer_path), + local_files_only=True, + ) + + @staticmethod + def _create_tokenizer( + raw_tokenizer, + definition: "TokenizerDefinition", + ) -> BaseTokenizer: + encoder_class = definition.encoder_class + + if encoder_class is VisionLanguageTokenizer: + if definition.processor_class is None: + raise ValueError("VisionLanguageTokenizer requires processor_class in definition") + processor = definition.processor_class(tokenizer=raw_tokenizer) + return VisionLanguageTokenizer( + tokenizer=raw_tokenizer, + processor=processor, + max_length=definition.max_length, + template=definition.template, + image_token=definition.image_token, + ) + else: + # Default to LanguageTokenizer for all text-only cases + return LanguageTokenizer( + tokenizer=raw_tokenizer, + max_length=definition.max_length, + padding=definition.padding, + template=definition.template, + use_chat_template=definition.use_chat_template, + chat_template_kwargs=definition.chat_template_kwargs or {}, + add_special_tokens=definition.add_special_tokens, + ) diff --git a/src/mflux/models/common/tokenizer/tokenizer_output.py b/src/mflux/models/common/tokenizer/tokenizer_output.py new file mode 100644 index 0000000..efb44b1 --- /dev/null +++ b/src/mflux/models/common/tokenizer/tokenizer_output.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass + +import mlx.core as mx + + +@dataclass +class TokenizerOutput: + input_ids: mx.array + attention_mask: mx.array + pixel_values: mx.array | None = None + image_grid_thw: mx.array | None = None diff --git a/src/mflux/models/common/weights/__init__.py b/src/mflux/models/common/weights/__init__.py index d602926..3c68bee 100644 --- a/src/mflux/models/common/weights/__init__.py +++ b/src/mflux/models/common/weights/__init__.py @@ -1,5 +1,14 @@ -"""Common weight mapping utilities.""" +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights, MetaData +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition +from mflux.models.common.weights.loading.weight_loader import WeightLoader +from mflux.models.common.weights.saving.model_saver import ModelSaver -from mflux.models.common.weights.model_saver import ModelSaver - -__all__ = ["ModelSaver"] +__all__ = [ + "ComponentDefinition", + "LoadedWeights", + "MetaData", + "ModelSaver", + "WeightApplier", + "WeightLoader", +] diff --git a/src/mflux/ui/__init__.py b/src/mflux/models/common/weights/loading/__init__.py similarity index 100% rename from src/mflux/ui/__init__.py rename to src/mflux/models/common/weights/loading/__init__.py diff --git a/src/mflux/models/common/weights/loading/loaded_weights.py b/src/mflux/models/common/weights/loading/loaded_weights.py new file mode 100644 index 0000000..b45e0f2 --- /dev/null +++ b/src/mflux/models/common/weights/loading/loaded_weights.py @@ -0,0 +1,40 @@ +from dataclasses import dataclass + + +@dataclass +class MetaData: + quantization_level: int | None = None + mflux_version: str | None = None + + +@dataclass +class LoadedWeights: + components: dict[str, dict] + meta_data: MetaData + + def __getattr__(self, name: str) -> dict | None: + if name in ("components", "meta_data"): + return object.__getattribute__(self, name) + return self.components.get(name) + + def num_transformer_blocks(self, component_name: str = "transformer") -> int: + transformer = self.components.get(component_name) + if transformer is None: + for comp in self.components.values(): + if isinstance(comp, dict) and "transformer_blocks" in comp: + transformer = comp + break + if transformer and "transformer_blocks" in transformer: + return len(transformer["transformer_blocks"]) + return 0 + + def num_single_transformer_blocks(self, component_name: str = "transformer") -> int: + transformer = self.components.get(component_name) + if transformer is None: + for comp in self.components.values(): + if isinstance(comp, dict) and "single_transformer_blocks" in comp: + transformer = comp + break + if transformer and "single_transformer_blocks" in transformer: + return len(transformer["single_transformer_blocks"]) + return 0 diff --git a/src/mflux/models/common/weights/loading/weight_applier.py b/src/mflux/models/common/weights/loading/weight_applier.py new file mode 100644 index 0000000..4b267c8 --- /dev/null +++ b/src/mflux/models/common/weights/loading/weight_applier.py @@ -0,0 +1,103 @@ +from typing import TYPE_CHECKING + +import mlx.nn as nn + +from mflux.models.common.resolution.quantization_resolution import QuantizationResolution +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition + +if TYPE_CHECKING: + from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType + + +class WeightApplier: + @staticmethod + def apply_and_quantize_single( + weights: LoadedWeights, + model: nn.Module, + component: ComponentDefinition, + quantize_arg: int | None, + quantization_predicate=None, + ) -> int | None: + stored_q = weights.meta_data.quantization_level + component_weights = weights.components.get(component.name) + + if component_weights is None: + raise ValueError(f"No weights found for component: {component.name}") + + if quantization_predicate is None: + + def quantization_predicate(path, module): + return hasattr(module, "to_quantized") + + bits, warning = QuantizationResolution.resolve(stored=stored_q, requested=quantize_arg) + + if warning: + print(f"โš ๏ธ {warning}") + + if bits is None: + model.update(component_weights, strict=False) + elif stored_q is None: + model.update(component_weights, strict=False) + if not component.skip_quantization: + nn.quantize(model, class_predicate=quantization_predicate, bits=bits) + else: + if not component.skip_quantization: + nn.quantize(model, class_predicate=quantization_predicate, bits=bits) + model.update(component_weights, strict=False) + + return bits + + @staticmethod + def apply_and_quantize( + weights: LoadedWeights, + models: dict[str, nn.Module], + quantize_arg: int | None, + weight_definition: "WeightDefinitionType", + ) -> int | None: + stored_q = weights.meta_data.quantization_level + components = {c.name: c for c in weight_definition.get_components()} + + bits, warning = QuantizationResolution.resolve(stored=stored_q, requested=quantize_arg) + + if warning: + print(f"โš ๏ธ {warning}") + + if bits is None: + WeightApplier._set_weights(weights, models, components) + elif stored_q is None: + WeightApplier._set_weights(weights, models, components) + WeightApplier._quantize(models, bits, components, weight_definition) + else: + WeightApplier._quantize(models, bits, components, weight_definition) + WeightApplier._set_weights(weights, models, components) + + return bits + + @staticmethod + def _set_weights( + weights: LoadedWeights, + models: dict[str, nn.Module], + components: dict | None = None, + ) -> None: + for name, model in models.items(): + component_weights = weights.components.get(name) + if component_weights is not None: + if components is not None: + component = components.get(name) + if component is not None and component.weight_subkey is not None: + component_weights = component_weights.get(component.weight_subkey, component_weights) + model.update(component_weights, strict=False) + + @staticmethod + def _quantize( + models: dict[str, nn.Module], + bits: int, + components: dict, + weight_definition: "WeightDefinitionType", + ) -> None: + for name, model in models.items(): + component = components.get(name) + if component and component.skip_quantization: + continue + nn.quantize(model, class_predicate=weight_definition.quantization_predicate, bits=bits) diff --git a/src/mflux/models/common/weights/loading/weight_definition.py b/src/mflux/models/common/weights/loading/weight_definition.py new file mode 100644 index 0000000..ba86034 --- /dev/null +++ b/src/mflux/models/common/weights/loading/weight_definition.py @@ -0,0 +1,59 @@ +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, List, TypeAlias + +import mlx.core as mx + +from mflux.models.common.weights.mapping.weight_mapping import WeightTarget + +if TYPE_CHECKING: + from mflux.models.common.tokenizer.tokenizer import BaseTokenizer + from mflux.models.depth_pro.weights.depth_pro_weight_definition import DepthProWeightDefinition + from mflux.models.fibo.weights.fibo_weight_definition import FIBOWeightDefinition + from mflux.models.fibo_vlm.weights.fibo_vlm_weight_definition import FIBOVLMWeightDefinition + from mflux.models.flux.weights.flux_weight_definition import FluxWeightDefinition + from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition + from mflux.models.z_image.weights.z_image_weight_definition import ZImageWeightDefinition + + WeightDefinitionType: TypeAlias = type[ + FluxWeightDefinition + | FIBOWeightDefinition + | FIBOVLMWeightDefinition + | QwenWeightDefinition + | ZImageWeightDefinition + | DepthProWeightDefinition + ] + + +@dataclass +class ComponentDefinition: + name: str + hf_subdir: str + mapping_getter: Callable[[], List[WeightTarget]] | None = None + model_attr: str | None = None + num_blocks: int | None = None + num_layers: int | None = None + loading_mode: str = "mlx_native" + precision: mx.Dtype | None = None + skip_quantization: bool = False + bulk_transform: Callable[[mx.array], mx.array] | None = None + weight_subkey: str | None = None + download_url: str | None = None + weight_prefix_filters: List[str] | None = None + + +@dataclass +class TokenizerDefinition: + name: str + hf_subdir: str + tokenizer_class: str = "AutoTokenizer" + fallback_subdirs: List[str] | None = None + download_patterns: List[str] | None = None + encoder_class: type["BaseTokenizer"] | None = None + max_length: int = 512 + padding: str = "max_length" + template: str | None = None + use_chat_template: bool = False + chat_template_kwargs: dict | None = field(default_factory=dict) + add_special_tokens: bool = True + processor_class: type | None = None + image_token: str = "<|image_pad|>" diff --git a/src/mflux/models/common/weights/loading/weight_loader.py b/src/mflux/models/common/weights/loading/weight_loader.py new file mode 100644 index 0000000..c483dda --- /dev/null +++ b/src/mflux/models/common/weights/loading/weight_loader.py @@ -0,0 +1,317 @@ +import json +import logging +import urllib.error +import urllib.request +from pathlib import Path +from typing import TYPE_CHECKING + +import mlx.core as mx +import torch +from huggingface_hub import snapshot_download +from mlx.utils import tree_unflatten +from safetensors.torch import load_file as torch_load_file + +from mflux.cli.defaults.defaults import MFLUX_CACHE_DIR +from mflux.models.common.resolution.path_resolution import PathResolution +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights, MetaData +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition +from mflux.models.common.weights.mapping.weight_mapper import WeightMapper + +if TYPE_CHECKING: + from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType + +logger = logging.getLogger(__name__) + + +class WeightLoader: + @staticmethod + def load_single( + component: ComponentDefinition, + repo_id: str, + file_pattern: str = "*.safetensors", + ) -> LoadedWeights: + root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=[file_pattern, "config.json"])) + weights, q_level, version = WeightLoader._load_component(root_path, component) + return LoadedWeights( + components={component.name: weights}, + meta_data=MetaData(quantization_level=q_level, mflux_version=version), + ) + + @staticmethod + def load( + weight_definition: "WeightDefinitionType", + model_path: str | None = None, + ) -> LoadedWeights: + root_path = PathResolution.resolve( + path=model_path, + patterns=weight_definition.get_download_patterns(), + ) + + # 2. Load each component (with caching for shared sources) + components = {} + quantization_level = None + mflux_version = None + raw_weights_cache: dict[tuple, dict] = {} # Cache by (path, loading_mode) + + for component in weight_definition.get_components(): + weights, q_level, version = WeightLoader._load_component(root_path, component, raw_weights_cache) + components[component.name] = weights + + # Track metadata from first component that has it + if quantization_level is None and q_level is not None: + quantization_level = q_level + mflux_version = version + + return LoadedWeights( + components=components, + meta_data=MetaData( + quantization_level=quantization_level, + mflux_version=mflux_version, + ), + ) + + @staticmethod + def _load_component( + root_path: Path | None, + component: ComponentDefinition, + raw_weights_cache: dict[tuple, dict] | None = None, + ) -> tuple[dict, int | None, str | None]: + # Handle direct URL downloads (e.g., Apple CDN for DepthPro) + if component.download_url is not None: + file_path = WeightLoader._download_from_url(component.download_url, component.name) + raw_weights = WeightLoader._load_weights_file(file_path, component.loading_mode) + else: + if root_path is None: + raise ValueError(f"No root_path and no download_url for component: {component.name}") + component_path = root_path / component.hf_subdir + + # Try mflux saved format first + weights, q_level, version = WeightLoader._try_load_mflux_format(component_path) + if weights is not None: + return weights, q_level, version + + # Check cache for shared loading (e.g., FIBO VLM decoder + visual from same source) + cache_key = (str(component_path), component.loading_mode) + if raw_weights_cache is not None and cache_key in raw_weights_cache: + raw_weights = raw_weights_cache[cache_key] + else: + # Fall back to HuggingFace format with mapping + raw_weights = WeightLoader._load_safetensors(component_path, component.loading_mode) + # Cache for potential reuse by other components + if raw_weights_cache is not None: + raw_weights_cache[cache_key] = raw_weights + + # Apply prefix filtering if specified (e.g., filter "model.language_model" vs "model.visual") + if component.weight_prefix_filters is not None: + raw_weights = { + k: v + for k, v in raw_weights.items() + if any(k.startswith(prefix) for prefix in component.weight_prefix_filters) + } + + # Apply precision conversion if specified + if component.precision is not None: + raw_weights = WeightLoader._convert_precision(raw_weights, component.precision) + + # Passthrough mode: apply bulk transform and unflatten (no key mapping) + if component.mapping_getter is None: + if component.bulk_transform is not None: + raw_weights = {k: component.bulk_transform(v) for k, v in raw_weights.items()} + return tree_unflatten(list(raw_weights.items())), None, None + + # Standard mode: apply declarative weight mapping + mapped_weights = WeightMapper.apply_mapping( + hf_weights=raw_weights, + mapping=component.mapping_getter(), + num_blocks=component.num_blocks, + num_layers=component.num_layers, + ) + return mapped_weights, None, None + + @staticmethod + def _try_load_mflux_format(path: Path) -> tuple[dict | None, int | None, str | None]: + if not path.exists(): + return None, None, None + + shard_files = sorted(f for f in path.glob("*.safetensors") if not f.name.startswith("._")) + if not shard_files: + return None, None, None + + # Check metadata on first file + data = mx.load(str(shard_files[0]), return_metadata=True) + if len(data) <= 1: + return None, None, None + + quantization_level_str = data[1].get("quantization_level") + mflux_version = data[1].get("mflux_version") + + # If no mflux metadata, this isn't our format + if quantization_level_str is None and mflux_version is None: + return None, None, None + + # Convert quantization level from string to int + quantization_level = int(quantization_level_str) if quantization_level_str is not None else None + + # Load all shards + all_weights: dict[str, mx.array] = {} + for shard in shard_files: + shard_data = mx.load(str(shard), return_metadata=True) + all_weights.update(dict(shard_data[0].items())) + + unflattened = tree_unflatten(list(all_weights.items())) + return unflattened, quantization_level, mflux_version + + @staticmethod + def _download_from_url(url: str, component_name: str) -> Path: + cache_dir = MFLUX_CACHE_DIR / component_name + cache_dir.mkdir(parents=True, exist_ok=True) + + # Extract filename from URL + filename = url.split("/")[-1] + file_path = cache_dir / filename + + if not file_path.exists(): + logger.info(f"Downloading {component_name} weights from {url}...") + try: + urllib.request.urlretrieve(url, file_path) + logger.info(f"Downloaded to {file_path}") + except (urllib.error.URLError, urllib.error.HTTPError) as e: + logger.error(f"Failed to download: {e}") + logger.info(f"Please manually download from: {url}") + raise FileNotFoundError(f"Model file not found at {file_path}") from e + + return file_path + + @staticmethod + def _load_weights_file(file_path: Path, loading_mode: str) -> dict[str, mx.array]: + if loading_mode == "torch_checkpoint": + return WeightLoader._load_torch_checkpoint(file_path) + elif loading_mode in ("mlx_native", "single"): + data = mx.load(str(file_path), return_metadata=True) + return dict(data[0].items()) + else: + raise ValueError(f"Unsupported loading mode for single file: {loading_mode}") + + @staticmethod + def _load_torch_checkpoint(file_path: Path) -> dict[str, mx.array]: + pt_weights = torch.load(file_path, map_location="cpu", weights_only=False) + return {k: mx.array(v.numpy()) for k, v in pt_weights.items() if isinstance(v, torch.Tensor)} + + @staticmethod + def _load_safetensors(path: Path, loading_mode: str) -> dict[str, mx.array]: + if loading_mode == "mlx_native": + return WeightLoader._load_mlx_native(path) + elif loading_mode == "torch_convert": + return WeightLoader._load_torch_convert(path) + elif loading_mode == "multi_json": + return WeightLoader._load_multi_json(path) + elif loading_mode == "torch_bfloat16": + return WeightLoader._load_torch_bfloat16(path) + elif loading_mode == "single": + return WeightLoader._load_single(path) + elif loading_mode == "multi_glob": + return WeightLoader._load_multi_glob(path) + else: + raise ValueError(f"Unknown loading mode: {loading_mode}") + + @staticmethod + def _load_mlx_native(path: Path) -> dict[str, mx.array]: + shard_files = sorted(f for f in path.glob("*.safetensors") if not f.name.startswith("._")) + if not shard_files: + raise FileNotFoundError(f"No safetensors files found in {path}") + + all_weights: dict[str, mx.array] = {} + for shard in shard_files: + weights = mx.load(str(shard)) + all_weights.update(weights) + + return all_weights + + @staticmethod + def _load_torch_convert(path: Path) -> dict[str, mx.array]: + shard_files = sorted(f for f in path.glob("*.safetensors") if not f.name.startswith("._")) + if not shard_files: + raise FileNotFoundError(f"No safetensors files found in {path}") + + all_weights: dict[str, mx.array] = {} + for shard in shard_files: + torch_weights = torch_load_file(str(shard)) + for key, tensor in torch_weights.items(): + if tensor.dtype == torch.bfloat16: + tensor = tensor.to(torch.float16) + all_weights[key] = mx.array(tensor.numpy()) + + return all_weights + + @staticmethod + def _load_multi_json(path: Path) -> dict[str, mx.array]: + index_path = path / "model.safetensors.index.json" + with open(index_path) as f: + index = json.load(f) + + # Group weights by file + files_to_load: dict[str, list[str]] = {} + for param_name, file_name in index["weight_map"].items(): + if file_name not in files_to_load: + files_to_load[file_name] = [] + files_to_load[file_name].append(param_name) + + all_weights: dict[str, mx.array] = {} + for file_name, param_names in files_to_load.items(): + file_path = path / file_name + + # Use mx.load which handles bfloat16 natively + file_weights = mx.load(str(file_path)) + + for param_name in param_names: + if param_name in file_weights: + all_weights[param_name] = file_weights[param_name] + + return all_weights + + @staticmethod + def _load_torch_bfloat16(path: Path) -> dict[str, mx.array]: + index_path = path / "model.safetensors.index.json" + with open(index_path) as f: + index = json.load(f) + + weight_files = sorted(set(index["weight_map"].values())) + + all_weights: dict[str, mx.array] = {} + for wf in weight_files: + file_path = path / wf + data = torch_load_file(str(file_path)) + for k, v in data.items(): + if v.dtype == torch.bfloat16: + v = v.to(torch.float16) + np_arr = v.detach().cpu().numpy() + all_weights[k] = mx.array(np_arr) + + return all_weights + + @staticmethod + def _load_single(path: Path) -> dict[str, mx.array]: + safetensors_files = [f for f in path.glob("*.safetensors") if not f.name.startswith("._")] + if not safetensors_files: + raise FileNotFoundError(f"No safetensors files found in {path}") + + weights_file = safetensors_files[0] + data = mx.load(str(weights_file), return_metadata=True) + return dict(data[0].items()) + + @staticmethod + def _load_multi_glob(path: Path) -> dict[str, mx.array]: + shard_files = sorted(f for f in path.glob("*.safetensors") if not f.name.startswith("._")) + if not shard_files: + raise FileNotFoundError(f"No safetensors files found in {path}") + + all_weights: dict[str, mx.array] = {} + for shard in shard_files: + data, _ = mx.load(str(shard), return_metadata=True) + all_weights.update(dict(data.items())) + + return all_weights + + @staticmethod + def _convert_precision(weights: dict[str, mx.array], precision: mx.Dtype) -> dict[str, mx.array]: + return {k: v.astype(precision) for k, v in weights.items()} diff --git a/src/mflux/models/common/weights/mapping/__init__.py b/src/mflux/models/common/weights/mapping/__init__.py index ded25f5..d5d2c98 100644 --- a/src/mflux/models/common/weights/mapping/__init__.py +++ b/src/mflux/models/common/weights/mapping/__init__.py @@ -1,6 +1,10 @@ -"""Weight mapping utilities for transforming HuggingFace weights to MLX structure.""" - from mflux.models.common.weights.mapping.weight_mapper import WeightMapper from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms -__all__ = ["WeightMapping", "WeightTarget", "WeightMapper"] +__all__ = [ + "WeightMapping", + "WeightTarget", + "WeightMapper", + "WeightTransforms", +] diff --git a/src/mflux/models/common/weights/mapping/weight_mapper.py b/src/mflux/models/common/weights/mapping/weight_mapper.py index 02a7160..718f06a 100644 --- a/src/mflux/models/common/weights/mapping/weight_mapper.py +++ b/src/mflux/models/common/weights/mapping/weight_mapper.py @@ -1,9 +1,3 @@ -""" -Weight Mapper - Applies declarative weight mappings to transform HF weights to MLX structure. - -Similar to LoRALoader, but for weight mapping instead of LoRA application. -""" - import re from typing import Callable, Dict, List, Optional @@ -13,8 +7,6 @@ from mflux.models.common.weights.mapping.weight_mapping import WeightTarget class WeightMapper: - """Maps HuggingFace weights to MLX nested structure using declarative mappings.""" - @staticmethod def apply_mapping( hf_weights: Dict[str, mx.array], @@ -22,18 +14,6 @@ class WeightMapper: num_blocks: Optional[int] = None, num_layers: Optional[int] = None, ) -> Dict: - """ - Apply weight mapping to transform HF weights to MLX structure. - - Args: - hf_weights: Raw HuggingFace weights (flat dict with dot-notation keys) - mapping: List of WeightTarget mappings - num_blocks: Number of transformer blocks (auto-detected if None) - num_layers: Number of text encoder layers (auto-detected if None) - - Returns: - Nested dict structure matching MLX model - """ # Auto-detect number of blocks if not provided if num_blocks is None: num_blocks = WeightMapper._detect_num_blocks(hf_weights) @@ -42,7 +22,7 @@ class WeightMapper: if num_layers is None: num_layers = WeightMapper._detect_num_layers(hf_weights) - # Build flat mapping: HF pattern -> (MLX path, transform) + # Build flat mapping: HF pattern -> [(MLX path, transform), ...] (supports one-to-many) flat_mapping = WeightMapper._build_flat_mapping(mapping, num_blocks, num_layers) # Map weights @@ -51,18 +31,19 @@ class WeightMapper: skipped_count = 0 for hf_key, hf_tensor in hf_weights.items(): - # Try to find matching mapping - mlx_path, transform = WeightMapper._find_mapping(hf_key, flat_mapping) + # Try to find matching mappings (can be multiple targets for one source) + targets = flat_mapping.get(hf_key, []) - if mlx_path: - # Apply transform if specified - tensor = hf_tensor - if transform: - tensor = transform(tensor) + if targets: + for mlx_path, transform in targets: + # Apply transform if specified + tensor = hf_tensor + if transform: + tensor = transform(tensor) - # Build nested structure - WeightMapper._set_nested_value(mapped_weights, mlx_path, tensor) - mapped_count += 1 + # Build nested structure + WeightMapper._set_nested_value(mapped_weights, mlx_path, tensor) + mapped_count += 1 else: # Weight not in mapping - might be intentionally skipped (e.g., lm_head) # or optional weight (e.g., conv_shortcut) - that's OK @@ -75,7 +56,6 @@ class WeightMapper: @staticmethod def _detect_num_blocks(hf_weights: Dict[str, mx.array]) -> int: - """Detect number of transformer blocks from weight keys.""" block_numbers = set() for key in hf_weights.keys(): # Match pattern: transformer_blocks.{number}.something @@ -89,7 +69,6 @@ class WeightMapper: @staticmethod def _detect_num_layers(hf_weights: Dict[str, mx.array]) -> int: - """Detect number of text encoder layers from weight keys.""" layer_numbers = set() for key in hf_weights.keys(): # Match pattern: model.layers.{number}.something @@ -104,89 +83,85 @@ class WeightMapper: @staticmethod def _build_flat_mapping( mapping: List[WeightTarget], num_blocks: int = 0, num_layers: int = 28 - ) -> Dict[str, tuple[str, Optional[Callable[[mx.array], mx.array]]]]: - """ - Build flat mapping from declarative targets. + ) -> Dict[str, List[tuple[str, Optional[Callable[[mx.array], mx.array]]]]]: + flat: Dict[str, List[tuple[str, Optional[Callable[[mx.array], mx.array]]]]] = {} - Returns: - Dict mapping HF pattern -> (MLX path, transform) - """ - flat = {} + def add_mapping(hf_key: str, mlx_path: str, transform: Optional[Callable[[mx.array], mx.array]]): + if hf_key not in flat: + flat[hf_key] = [] + flat[hf_key].append((mlx_path, transform)) for target in mapping: # Expand placeholders for each pattern - for hf_pattern in target.hf_patterns: - # Check which placeholders are present - has_block = "{block}" in hf_pattern or "{block}" in target.mlx_path - has_i = "{i}" in hf_pattern or "{i}" in target.mlx_path - has_res = "{res}" in hf_pattern or "{res}" in target.mlx_path - has_layer = "{layer}" in hf_pattern or "{layer}" in target.mlx_path + for hf_pattern in target.from_pattern: + # Check which placeholders are present in BOTH patterns + hf_has_block = "{block}" in hf_pattern + to_has_block = "{block}" in target.to_pattern + has_i = "{i}" in hf_pattern or "{i}" in target.to_pattern + has_res = "{res}" in hf_pattern or "{res}" in target.to_pattern + has_layer = "{layer}" in hf_pattern or "{layer}" in target.to_pattern # Handle multiple placeholders together - if has_block and has_res: + if (hf_has_block or to_has_block) and has_res: # Up blocks: expand both {block} and {res} max_blocks = num_blocks if num_blocks > 0 else 4 # Default 4 for up_blocks for block_num in range(max_blocks): for res in range(3): # 3 resnets per up_block concrete_hf = hf_pattern.replace("{block}", str(block_num)).replace("{res}", str(res)) - concrete_mlx = target.mlx_path.replace("{block}", str(block_num)).replace("{res}", str(res)) - flat[concrete_hf] = (concrete_mlx, target.transform) - elif has_block: - # Expand {block} only (for transformer blocks or visual blocks) - # Check if target has max_blocks override + concrete_mlx = target.to_pattern.replace("{block}", str(block_num)).replace( + "{res}", str(res) + ) + add_mapping(concrete_hf, concrete_mlx, target.transform) + elif hf_has_block and to_has_block: + # Both have {block} - standard one-to-one expansion if target.max_blocks is not None: max_blocks = target.max_blocks - # Check if this is for visual blocks (32 blocks) or transformer blocks - elif "visual.blocks" in hf_pattern or "visual.blocks" in target.mlx_path: + elif "visual.blocks" in hf_pattern or "visual.blocks" in target.to_pattern: max_blocks = 32 # Visual blocks are always 32 else: - max_blocks = num_blocks if num_blocks > 0 else 4 # Default 4 for up_blocks + max_blocks = num_blocks if num_blocks > 0 else 4 for block_num in range(max_blocks): concrete_hf = hf_pattern.replace("{block}", str(block_num)) - concrete_mlx = target.mlx_path.replace("{block}", str(block_num)) - flat[concrete_hf] = (concrete_mlx, target.transform) + concrete_mlx = target.to_pattern.replace("{block}", str(block_num)) + add_mapping(concrete_hf, concrete_mlx, target.transform) + elif to_has_block and not hf_has_block: + # One-to-many: single HF key maps to multiple MLX targets (e.g., relative_attention_bias) + if target.max_blocks is not None: + max_blocks = target.max_blocks + else: + max_blocks = num_blocks if num_blocks > 0 else 24 # Default for T5 + for block_num in range(max_blocks): + concrete_mlx = target.to_pattern.replace("{block}", str(block_num)) + add_mapping(hf_pattern, concrete_mlx, target.transform) elif has_layer: # Expand {layer} for text encoder layers or visual blocks max_layers = num_layers if num_layers > 0 else 28 # Default 28 for text encoder for layer_num in range(max_layers): concrete_hf = hf_pattern.replace("{layer}", str(layer_num)) - concrete_mlx = target.mlx_path.replace("{layer}", str(layer_num)) - flat[concrete_hf] = (concrete_mlx, target.transform) + concrete_mlx = target.to_pattern.replace("{layer}", str(layer_num)) + add_mapping(concrete_hf, concrete_mlx, target.transform) elif has_i: # Expand {i} only (for mid_block resnets) for i in range(2): # 2 resnets in mid_block concrete_hf = hf_pattern.replace("{i}", str(i)) - concrete_mlx = target.mlx_path.replace("{i}", str(i)) - flat[concrete_hf] = (concrete_mlx, target.transform) + concrete_mlx = target.to_pattern.replace("{i}", str(i)) + add_mapping(concrete_hf, concrete_mlx, target.transform) elif has_res: # This shouldn't happen for VAE (encoder down_blocks are explicit) # But handle it just in case if "up_block" in hf_pattern: for res in range(3): concrete_hf = hf_pattern.replace("{res}", str(res)) - concrete_mlx = target.mlx_path.replace("{res}", str(res)) - flat[concrete_hf] = (concrete_mlx, target.transform) + concrete_mlx = target.to_pattern.replace("{res}", str(res)) + add_mapping(concrete_hf, concrete_mlx, target.transform) else: # No placeholder, use as-is - flat[hf_pattern] = (target.mlx_path, target.transform) + add_mapping(hf_pattern, target.to_pattern, target.transform) return flat - @staticmethod - def _find_mapping( - hf_key: str, flat_mapping: Dict[str, tuple[str, Optional[Callable[[mx.array], mx.array]]]] - ) -> tuple[Optional[str], Optional[Callable[[mx.array], mx.array]]]: - """Find MLX path and transform for a given HF key.""" - return flat_mapping.get(hf_key, (None, None)) - @staticmethod def _set_nested_value(d: Dict, path: str, value: mx.array): - """ - Set value in nested dict using dot-notation path. - - Creates nested structure as needed. - Handles both dict keys and list indices. - """ parts = path.split(".") current = d i = 0 diff --git a/src/mflux/models/common/weights/mapping/weight_mapping.py b/src/mflux/models/common/weights/mapping/weight_mapping.py index 501c37a..c7ed8a8 100644 --- a/src/mflux/models/common/weights/mapping/weight_mapping.py +++ b/src/mflux/models/common/weights/mapping/weight_mapping.py @@ -1,9 +1,3 @@ -""" -Base classes for declarative weight mapping. - -Similar to LoRA mapping, but for transforming HuggingFace weights to MLX structure. -""" - from dataclasses import dataclass from typing import Callable, List, Optional, Protocol @@ -12,23 +6,14 @@ import mlx.core as mx @dataclass class WeightTarget: - """ - Declarative weight mapping target. - - Maps HuggingFace weight names to MLX nested structure paths. - """ - - mlx_path: str # MLX nested path, e.g., "transformer_blocks.{block}.attn.to_q.weight" - hf_patterns: List[str] # HuggingFace naming patterns, e.g., ["transformer_blocks.{block}.attn.to_q.weight"] - transform: Optional[Callable[[mx.array], mx.array]] = None # Optional transform (reshape, transpose, etc.) - required: bool = True # If False, weight is optional (may not exist in all models) - max_blocks: Optional[int] = None # Override num_blocks for this target (e.g., for fixed-size lists) + to_pattern: str + from_pattern: List[str] + transform: Optional[Callable[[mx.array], mx.array]] = None + required: bool = True + max_blocks: Optional[int] = None class WeightMapping(Protocol): - """Protocol for weight mapping classes.""" - @staticmethod def get_mapping() -> List[WeightTarget]: - """Return list of weight mapping targets.""" return [] diff --git a/src/mflux/models/common/weights/mapping/weight_transforms.py b/src/mflux/models/common/weights/mapping/weight_transforms.py new file mode 100644 index 0000000..9f64bcb --- /dev/null +++ b/src/mflux/models/common/weights/mapping/weight_transforms.py @@ -0,0 +1,33 @@ +import mlx.core as mx + + +class WeightTransforms: + @staticmethod + def reshape_gamma_to_1d(tensor: mx.array) -> mx.array: + if len(tensor.shape) > 1: + return mx.reshape(tensor, (tensor.shape[0],)) + return tensor + + @staticmethod + def transpose_patch_embed(tensor: mx.array) -> mx.array: + if len(tensor.shape) == 5: + return tensor.transpose(0, 2, 3, 4, 1) + return tensor + + @staticmethod + def transpose_conv3d_weight(tensor: mx.array) -> mx.array: + if len(tensor.shape) == 5: + return tensor.transpose(0, 2, 3, 4, 1) + return tensor + + @staticmethod + def transpose_conv2d_weight(tensor: mx.array) -> mx.array: + if len(tensor.shape) == 4: + return tensor.transpose(0, 2, 3, 1) + return tensor + + @staticmethod + def transpose_conv_transpose2d_weight(tensor: mx.array) -> mx.array: + if len(tensor.shape) == 4: + return tensor.transpose(1, 2, 3, 0) + return tensor diff --git a/src/mflux/models/common/weights/model_saver.py b/src/mflux/models/common/weights/model_saver.py deleted file mode 100644 index 44ee146..0000000 --- a/src/mflux/models/common/weights/model_saver.py +++ /dev/null @@ -1,114 +0,0 @@ -from pathlib import Path -from typing import Any - -import mlx.core as mx -from mlx import nn -from mlx.utils import tree_flatten -from transformers import PreTrainedTokenizer - -from mflux.utils.version_util import VersionUtil - - -class ModelSaver: - @staticmethod - def save_model( - model: Any, - bits: int, - base_path: str, - tokenizers: list[tuple[str, str]] | None = None, - components: list[tuple[str, str]] | None = None, - ) -> None: - # Default tokenizers: try common patterns - if tokenizers is None: - tokenizers = ModelSaver._detect_tokenizers(model) - - # Default components: try common patterns - if components is None: - components = ModelSaver._detect_components(model) - - # Save tokenizers - for attr_path, subdir in tokenizers: - tokenizer = ModelSaver._get_nested_attr(model, attr_path) - if tokenizer is not None: - ModelSaver._save_tokenizer(base_path, tokenizer, subdir) - - # Save model components - for attr_name, subdir in components: - component = getattr(model, attr_name, None) - if component is not None: - ModelSaver._save_weights(base_path, bits, component, subdir) - - @staticmethod - def _detect_tokenizers(model: Any) -> list[tuple[str, str]]: - tokenizers = [] - if hasattr(model, "clip_tokenizer") and hasattr(model.clip_tokenizer, "tokenizer"): - tokenizers.append(("clip_tokenizer.tokenizer", "tokenizer")) - if hasattr(model, "t5_tokenizer") and hasattr(model.t5_tokenizer, "tokenizer"): - tokenizers.append(("t5_tokenizer.tokenizer", "tokenizer_2")) - if hasattr(model, "qwen_tokenizer") and hasattr(model.qwen_tokenizer, "tokenizer"): - tokenizers.append(("qwen_tokenizer.tokenizer", "tokenizer")) - if hasattr(model, "fibo_tokenizer") and hasattr(model.fibo_tokenizer, "tokenizer"): - tokenizers.append(("fibo_tokenizer.tokenizer", "tokenizer")) - return tokenizers - - @staticmethod - def _detect_components(model: Any) -> list[tuple[str, str]]: - components = [] - if hasattr(model, "vae"): - components.append(("vae", "vae")) - if hasattr(model, "transformer"): - components.append(("transformer", "transformer")) - if hasattr(model, "clip_text_encoder"): - components.append(("clip_text_encoder", "text_encoder")) - if hasattr(model, "t5_text_encoder"): - components.append(("t5_text_encoder", "text_encoder_2")) - if hasattr(model, "text_encoder"): - # Only add if we haven't already added clip_text_encoder or t5_text_encoder - if not any(c[1] == "text_encoder" for c in components): - components.append(("text_encoder", "text_encoder")) - return components - - @staticmethod - def _get_nested_attr(obj: Any, attr_path: str) -> Any: - attrs = attr_path.split(".") - result = obj - for attr in attrs: - if not hasattr(result, attr): - return None - result = getattr(result, attr) - return result - - @staticmethod - def _save_tokenizer(base_path: str, tokenizer: PreTrainedTokenizer, subdir: str) -> None: - path = Path(base_path) / subdir - path.mkdir(parents=True, exist_ok=True) - tokenizer.save_pretrained(path) - - @staticmethod - def _save_weights(base_path: str, bits: int, model: nn.Module, subdir: str) -> None: - path = Path(base_path) / subdir - path.mkdir(parents=True, exist_ok=True) - weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters()))) - for i, weight in enumerate(weights): - mx.save_safetensors( - str(path / f"{i}.safetensors"), - weight, - { - "quantization_level": str(bits), - "mflux_version": VersionUtil.get_mflux_version(), - }, - ) - - @staticmethod - def _split_weights(base_path: str, weights: dict, max_file_size_gb: int = 2) -> list: - max_file_size_bytes = max_file_size_gb << 30 - shards = [] - shard, shard_size = {}, 0 - for k, v in weights.items(): - if shard_size + v.nbytes > max_file_size_bytes: - shards.append(shard) - shard, shard_size = {}, 0 - shard[k] = v - shard_size += v.nbytes - shards.append(shard) - return shards diff --git a/src/mflux/ui/cli/__init__.py b/src/mflux/models/common/weights/saving/__init__.py similarity index 100% rename from src/mflux/ui/cli/__init__.py rename to src/mflux/models/common/weights/saving/__init__.py diff --git a/src/mflux/models/common/weights/saving/model_saver.py b/src/mflux/models/common/weights/saving/model_saver.py new file mode 100644 index 0000000..6ebeba8 --- /dev/null +++ b/src/mflux/models/common/weights/saving/model_saver.py @@ -0,0 +1,96 @@ +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import mlx.core as mx +from mlx import nn +from mlx.utils import tree_flatten +from tqdm import tqdm +from transformers import PreTrainedTokenizer + +from mflux.utils.version_util import VersionUtil + +if TYPE_CHECKING: + from mflux.models.common.weights.loading.weight_definition import WeightDefinitionType + + +class ModelSaver: + @staticmethod + def save_model( + model: Any, + bits: int, + base_path: str, + weight_definition: "WeightDefinitionType", + ) -> None: + # Save tokenizers from model.tokenizers dict + tokenizer_defs = weight_definition.get_tokenizers() + for t in tokenizer_defs: + if hasattr(model, "tokenizers") and t.name in model.tokenizers: + tokenizer_wrapper = model.tokenizers[t.name] + if hasattr(tokenizer_wrapper, "tokenizer"): + ModelSaver._save_tokenizer(base_path, tokenizer_wrapper.tokenizer, t.hf_subdir) + + # Save model components with progress bar + components = [(c.model_attr or c.name, c.hf_subdir) for c in weight_definition.get_components()] + for attr_name, subdir in tqdm(components, desc="Saving components", unit="component"): + component = getattr(model, attr_name, None) + if component is not None: + ModelSaver._save_weights(base_path, bits, component, subdir) + + @staticmethod + def _save_tokenizer(base_path: str, tokenizer: PreTrainedTokenizer, subdir: str) -> None: + path = Path(base_path) / subdir + path.mkdir(parents=True, exist_ok=True) + tokenizer.save_pretrained(path) + + @staticmethod + def _save_weights(base_path: str, bits: int, model: nn.Module, subdir: str) -> None: + path = Path(base_path) / subdir + path.mkdir(parents=True, exist_ok=True) + weights = dict(tree_flatten(model.parameters())) + shards = ModelSaver._split_weights(weights) + + # Build weight_map for index.json (maps each weight key to its shard file) + weight_map = {} + shard_iter = tqdm(enumerate(shards), total=len(shards), desc=f" {subdir}", unit="shard", leave=False) + for i, shard in shard_iter: + shard_filename = f"{i}.safetensors" + mx.save_safetensors( + str(path / shard_filename), + shard, + { + "quantization_level": str(bits), + "mflux_version": VersionUtil.get_mflux_version(), + }, + ) + # Record which file each weight belongs to + for key in shard.keys(): + weight_map[key] = shard_filename + + # Write model.safetensors.index.json for HuggingFace compatibility + # This ensures the saved model works even if custom metadata is stripped + index_data = { + "metadata": { + "quantization_level": str(bits), + "mflux_version": VersionUtil.get_mflux_version(), + }, + "weight_map": weight_map, + } + with open(path / "model.safetensors.index.json", "w") as f: + json.dump(index_data, f, indent=2) + + @staticmethod + def _split_weights(weights: dict, max_file_size_gb: int = 2) -> list[dict]: + max_file_size_bytes = max_file_size_gb << 30 + shards: list[dict] = [] + shard: dict = {} + shard_size = 0 + for k, v in weights.items(): + if shard_size + v.nbytes > max_file_size_bytes: + shards.append(shard) + shard, shard_size = {}, 0 + shard[k] = v + shard_size += v.nbytes + if shard: # Don't append empty shard + shards.append(shard) + return shards diff --git a/tests/model_config/__init__.py b/src/mflux/models/depth_pro/cli/__init__.py similarity index 100% rename from tests/model_config/__init__.py rename to src/mflux/models/depth_pro/cli/__init__.py diff --git a/src/mflux/save_depth.py b/src/mflux/models/depth_pro/cli/save_depth.py similarity index 86% rename from src/mflux/save_depth.py rename to src/mflux/models/depth_pro/cli/save_depth.py index cd54517..bed9d1b 100644 --- a/src/mflux/save_depth.py +++ b/src/mflux/models/depth_pro/cli/save_depth.py @@ -1,7 +1,7 @@ from pathlib import Path -from mflux.models.depth_pro.depth_pro import DepthPro -from mflux.ui.cli.parsers import CommandLineParser +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.depth_pro.model.depth_pro import DepthPro def main(): diff --git a/src/mflux/models/depth_pro/depth_pro_initializer.py b/src/mflux/models/depth_pro/depth_pro_initializer.py index d3edb92..56d56d6 100644 --- a/src/mflux/models/depth_pro/depth_pro_initializer.py +++ b/src/mflux/models/depth_pro/depth_pro_initializer.py @@ -1,25 +1,21 @@ -import mlx.nn as nn - +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_loader import WeightLoader from mflux.models.depth_pro.model.depth_pro_model import DepthProModel -from mflux.models.depth_pro.weights.weight_handler_depth_pro import WeightHandlerDepthPro +from mflux.models.depth_pro.weights.depth_pro_weight_definition import DepthProWeightDefinition class DepthProInitializer: @staticmethod - def init(depth_pro_model: DepthProModel, quantize: int | None = None) -> None: - # 1. Load the weights - depth_pro_weights = WeightHandlerDepthPro.load_weights() - WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent0") - WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent1") - WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample0") - WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample1") - WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample2") - WeightHandlerDepthPro.reposition_head_weights(depth_pro_weights) - WeightHandlerDepthPro.reshape_transposed_convolution_weights(depth_pro_weights) + def init(model: DepthProModel, quantize: int | None = None) -> None: + # 1. Load weights using unified loader (handles download from Apple CDN) + weights = WeightLoader.load(weight_definition=DepthProWeightDefinition) - # 2. Assign the weights to the model - depth_pro_model.update(depth_pro_weights.weights, strict=False) - - # 3. Optionally quantize the model - if quantize: - nn.quantize(depth_pro_model, bits=quantize) + # 2. Apply weights and quantize using unified applier + WeightApplier.apply_and_quantize( + weights=weights, + quantize_arg=quantize, + weight_definition=DepthProWeightDefinition, + models={ + "depth_pro": model, + }, + ) diff --git a/src/mflux/models/depth_pro/model/conv_utils.py b/src/mflux/models/depth_pro/model/conv_utils.py deleted file mode 100644 index 2f6be4f..0000000 --- a/src/mflux/models/depth_pro/model/conv_utils.py +++ /dev/null @@ -1,23 +0,0 @@ -import mlx.core as mx -import mlx.nn as nn - - -class ConvUtils: - @staticmethod - def apply_conv(x: mx.array, conv_module: nn.Module) -> mx.array: - """Apply a convolution with channel format conversion. - - MLX expects channels-last format (B,H,W,C) for convolutions, - but tensors are generally in channels-first format (B,C,H,W). - This helper handles the conversion automatically. - - Args: - x: Input tensor in channels-first format (B,C,H,W) - conv_module: Convolution module to apply - - Returns: - Output tensor in channels-first format (B,C,H,W) - """ - x = mx.transpose(x, (0, 2, 3, 1)) - x = conv_module(x) - return mx.transpose(x, (0, 3, 1, 2)) diff --git a/tests/ui/__init__.py b/src/mflux/models/depth_pro/model/decoder/__init__.py similarity index 100% rename from tests/ui/__init__.py rename to src/mflux/models/depth_pro/model/decoder/__init__.py diff --git a/src/mflux/models/depth_pro/model/feature_fusion_block_2d.py b/src/mflux/models/depth_pro/model/decoder/feature_fusion_block_2d.py similarity index 77% rename from src/mflux/models/depth_pro/model/feature_fusion_block_2d.py rename to src/mflux/models/depth_pro/model/decoder/feature_fusion_block_2d.py index 5dba7d9..8a19e38 100644 --- a/src/mflux/models/depth_pro/model/feature_fusion_block_2d.py +++ b/src/mflux/models/depth_pro/model/decoder/feature_fusion_block_2d.py @@ -1,8 +1,8 @@ import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.conv_utils import ConvUtils -from mflux.models.depth_pro.model.residual_block import ResidualBlock +from mflux.models.depth_pro.model.decoder.residual_block import ResidualBlock +from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil class FeatureFusionBlock2d(nn.Module): @@ -21,6 +21,6 @@ class FeatureFusionBlock2d(nn.Module): x = x + res x = self.resnet2(x) if self.use_deconv: - x = ConvUtils.apply_conv(x, self.deconv) - x = ConvUtils.apply_conv(x, self.out_conv) + x = DepthProUtil.apply_conv(x, self.deconv) + x = DepthProUtil.apply_conv(x, self.out_conv) return x diff --git a/src/mflux/models/depth_pro/model/multires_conv_decoder.py b/src/mflux/models/depth_pro/model/decoder/multires_conv_decoder.py similarity index 74% rename from src/mflux/models/depth_pro/model/multires_conv_decoder.py rename to src/mflux/models/depth_pro/model/decoder/multires_conv_decoder.py index 655650b..b3db71d 100644 --- a/src/mflux/models/depth_pro/model/multires_conv_decoder.py +++ b/src/mflux/models/depth_pro/model/decoder/multires_conv_decoder.py @@ -1,8 +1,8 @@ import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.conv_utils import ConvUtils -from mflux.models.depth_pro.model.feature_fusion_block_2d import FeatureFusionBlock2d +from mflux.models.depth_pro.model.decoder.feature_fusion_block_2d import FeatureFusionBlock2d +from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil class MultiresConvDecoder(nn.Module): @@ -32,20 +32,20 @@ class MultiresConvDecoder(nn.Module): x_global_features: mx.array, ) -> mx.array: # Process global features: - features = ConvUtils.apply_conv(x_global_features, self.convs[4]) + features = DepthProUtil.apply_conv(x_global_features, self.convs[4]) features = self.fusions[4](features) # Process remaining levels with skip connections: - x1_skip_features = ConvUtils.apply_conv(x1_features, self.convs[3]) + x1_skip_features = DepthProUtil.apply_conv(x1_features, self.convs[3]) features = self.fusions[3](features, x1_skip_features) - x0_skip_features = ConvUtils.apply_conv(x0_features, self.convs[2]) + x0_skip_features = DepthProUtil.apply_conv(x0_features, self.convs[2]) features = self.fusions[2](features, x0_skip_features) - x1_skip_latents = ConvUtils.apply_conv(x1_latent, self.convs[1]) + x1_skip_latents = DepthProUtil.apply_conv(x1_latent, self.convs[1]) features = self.fusions[1](features, x1_skip_latents) - x0_skip_latents = ConvUtils.apply_conv(x0_latent, self.convs[0]) + x0_skip_latents = DepthProUtil.apply_conv(x0_latent, self.convs[0]) features = self.fusions[0](features, x0_skip_latents) return features diff --git a/src/mflux/models/depth_pro/model/residual_block.py b/src/mflux/models/depth_pro/model/decoder/residual_block.py similarity index 80% rename from src/mflux/models/depth_pro/model/residual_block.py rename to src/mflux/models/depth_pro/model/decoder/residual_block.py index 9f4ca42..f1cd4d8 100644 --- a/src/mflux/models/depth_pro/model/residual_block.py +++ b/src/mflux/models/depth_pro/model/decoder/residual_block.py @@ -1,7 +1,7 @@ import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.conv_utils import ConvUtils +from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil class ResidualBlock(nn.Module): @@ -30,7 +30,7 @@ class ResidualBlock(nn.Module): def __call__(self, x: mx.array) -> mx.array: delta_x = nn.relu(x) - delta_x = ConvUtils.apply_conv(delta_x, self.residual[1]) + delta_x = DepthProUtil.apply_conv(delta_x, self.residual[1]) delta_x = nn.relu(delta_x) - delta_x = ConvUtils.apply_conv(delta_x, self.residual[3]) + delta_x = DepthProUtil.apply_conv(delta_x, self.residual[3]) return x + delta_x diff --git a/src/mflux/models/depth_pro/depth_pro.py b/src/mflux/models/depth_pro/model/depth_pro.py similarity index 100% rename from src/mflux/models/depth_pro/depth_pro.py rename to src/mflux/models/depth_pro/model/depth_pro.py diff --git a/src/mflux/models/depth_pro/model/depth_pro_model.py b/src/mflux/models/depth_pro/model/depth_pro_model.py index d3bb45c..af3f664 100644 --- a/src/mflux/models/depth_pro/model/depth_pro_model.py +++ b/src/mflux/models/depth_pro/model/depth_pro_model.py @@ -1,9 +1,9 @@ import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.depth_pro_encoder import DepthProEncoder -from mflux.models.depth_pro.model.fov_head import FOVHead -from mflux.models.depth_pro.model.multires_conv_decoder import MultiresConvDecoder +from mflux.models.depth_pro.model.decoder.multires_conv_decoder import MultiresConvDecoder +from mflux.models.depth_pro.model.encoder.depth_pro_encoder import DepthProEncoder +from mflux.models.depth_pro.model.head.fov_head import FOVHead class DepthProModel(nn.Module): diff --git a/src/mflux/models/depth_pro/model/depth_pro_util.py b/src/mflux/models/depth_pro/model/depth_pro_util.py index c1c2024..99785c8 100644 --- a/src/mflux/models/depth_pro/model/depth_pro_util.py +++ b/src/mflux/models/depth_pro/model/depth_pro_util.py @@ -1,6 +1,7 @@ import math import mlx.core as mx +import mlx.nn as nn import numpy as np from PIL import Image @@ -64,3 +65,9 @@ class DepthProUtil: final_result_np = result_proc return mx.array(final_result_np) + + @staticmethod + def apply_conv(x: mx.array, conv_module: nn.Module) -> mx.array: + x = mx.transpose(x, (0, 2, 3, 1)) + x = conv_module(x) + return mx.transpose(x, (0, 3, 1, 2)) diff --git a/tests/weights/__init__.py b/src/mflux/models/depth_pro/model/encoder/__init__.py similarity index 100% rename from tests/weights/__init__.py rename to src/mflux/models/depth_pro/model/encoder/__init__.py diff --git a/src/mflux/models/depth_pro/model/depth_pro_encoder.py b/src/mflux/models/depth_pro/model/encoder/depth_pro_encoder.py similarity index 93% rename from src/mflux/models/depth_pro/model/depth_pro_encoder.py rename to src/mflux/models/depth_pro/model/encoder/depth_pro_encoder.py index fb0439d..c447fa9 100644 --- a/src/mflux/models/depth_pro/model/depth_pro_encoder.py +++ b/src/mflux/models/depth_pro/model/encoder/depth_pro_encoder.py @@ -3,9 +3,9 @@ import math import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.conv_utils import ConvUtils +from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil from mflux.models.depth_pro.model.dino_v2.dino_vision_transformer import DinoVisionTransformer -from mflux.models.depth_pro.model.upsample_block import UpSampleBlock +from mflux.models.depth_pro.model.encoder.upsample_block import UpSampleBlock class DepthProEncoder(nn.Module): @@ -56,9 +56,9 @@ class DepthProEncoder(nn.Module): # 4. Apply the image encoder model. x_global_features, _, _ = self.image_encoder(x2) x_global_features = DepthProEncoder._reshape_feature(embeddings=x_global_features, width=24, height=24) - x_global_features = ConvUtils.apply_conv(x_global_features, self.upsample_lowres) + x_global_features = DepthProUtil.apply_conv(x_global_features, self.upsample_lowres) x_global_features = mx.concatenate((x2_features, x_global_features), axis=1) - x_global_features = ConvUtils.apply_conv(x_global_features, self.fuse_lowres) + x_global_features = DepthProUtil.apply_conv(x_global_features, self.fuse_lowres) return ( x_latent0_features, diff --git a/src/mflux/models/depth_pro/model/upsample_block.py b/src/mflux/models/depth_pro/model/encoder/upsample_block.py similarity index 92% rename from src/mflux/models/depth_pro/model/upsample_block.py rename to src/mflux/models/depth_pro/model/encoder/upsample_block.py index 3a31621..4d29a42 100644 --- a/src/mflux/models/depth_pro/model/upsample_block.py +++ b/src/mflux/models/depth_pro/model/encoder/upsample_block.py @@ -1,7 +1,7 @@ import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.conv_utils import ConvUtils +from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil class UpSampleBlock(nn.Module): @@ -47,5 +47,5 @@ class UpSampleBlock(nn.Module): def __call__(self, x: mx.array) -> mx.array: for layer in self.layers: - x = ConvUtils.apply_conv(x, layer) + x = DepthProUtil.apply_conv(x, layer) return x diff --git a/src/mflux/models/depth_pro/model/head/__init__.py b/src/mflux/models/depth_pro/model/head/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/models/depth_pro/model/fov_head.py b/src/mflux/models/depth_pro/model/head/fov_head.py similarity index 70% rename from src/mflux/models/depth_pro/model/fov_head.py rename to src/mflux/models/depth_pro/model/head/fov_head.py index 1d94074..488175e 100644 --- a/src/mflux/models/depth_pro/model/fov_head.py +++ b/src/mflux/models/depth_pro/model/head/fov_head.py @@ -1,7 +1,7 @@ import mlx.core as mx import mlx.nn as nn -from mflux.models.depth_pro.model.conv_utils import ConvUtils +from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil class FOVHead(nn.Module): @@ -16,10 +16,10 @@ class FOVHead(nn.Module): ] def __call__(self, x: mx.array) -> mx.array: - x = ConvUtils.apply_conv(x, self.convs[0]) - x = ConvUtils.apply_conv(x, self.convs[1]) - x = ConvUtils.apply_conv(x, self.convs[2]) + x = DepthProUtil.apply_conv(x, self.convs[0]) + x = DepthProUtil.apply_conv(x, self.convs[1]) + x = DepthProUtil.apply_conv(x, self.convs[2]) x = nn.relu(x) - x = ConvUtils.apply_conv(x, self.convs[4]) + x = DepthProUtil.apply_conv(x, self.convs[4]) x = nn.relu(x) return x diff --git a/src/mflux/models/depth_pro/weights/depth_pro_weight_definition.py b/src/mflux/models/depth_pro/weights/depth_pro_weight_definition.py new file mode 100644 index 0000000..27f72f8 --- /dev/null +++ b/src/mflux/models/depth_pro/weights/depth_pro_weight_definition.py @@ -0,0 +1,34 @@ +from typing import List + +import mlx.nn as nn + +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition, TokenizerDefinition +from mflux.models.depth_pro.weights.depth_pro_weight_mapping import DepthProWeightMapping + + +class DepthProWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="depth_pro", + hf_subdir="", + loading_mode="torch_checkpoint", + mapping_getter=DepthProWeightMapping.get_mapping, + download_url="https://ml-site.cdn-apple.com/models/depth-pro/depth_pro.pt", + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [] + + @staticmethod + def get_download_patterns() -> List[str]: + return [] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + if isinstance(module, nn.Conv2d): + return False + return hasattr(module, "to_quantized") diff --git a/src/mflux/models/depth_pro/weights/depth_pro_weight_mapping.py b/src/mflux/models/depth_pro/weights/depth_pro_weight_mapping.py new file mode 100644 index 0000000..e73db3c --- /dev/null +++ b/src/mflux/models/depth_pro/weights/depth_pro_weight_mapping.py @@ -0,0 +1,307 @@ +from typing import List + +from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms + + +class DepthProWeightMapping(WeightMapping): + @staticmethod + def get_mapping() -> List[WeightTarget]: + return ( + DepthProWeightMapping._get_dino_encoder_mapping("patch_encoder") + + DepthProWeightMapping._get_dino_encoder_mapping("image_encoder") + + DepthProWeightMapping._get_upsample_block_mapping("upsample_latent0", num_layers=4) + + DepthProWeightMapping._get_upsample_block_mapping("upsample_latent1", num_layers=3) + + DepthProWeightMapping._get_upsample_block_mapping("upsample0", num_layers=2) + + DepthProWeightMapping._get_upsample_block_mapping("upsample1", num_layers=2) + + DepthProWeightMapping._get_upsample_block_mapping("upsample2", num_layers=2) + + DepthProWeightMapping._get_encoder_conv_mapping() + + DepthProWeightMapping._get_decoder_mapping() + + DepthProWeightMapping._get_head_mapping() + ) + + @staticmethod + def _get_dino_encoder_mapping(encoder_name: str) -> List[WeightTarget]: + prefix = f"encoder.{encoder_name}" + return [ + WeightTarget( + to_pattern=f"{prefix}.cls_token", + from_pattern=[f"{prefix}.cls_token"], + ), + WeightTarget( + to_pattern=f"{prefix}.pos_embed", + from_pattern=[f"{prefix}.pos_embed"], + ), + WeightTarget( + to_pattern=f"{prefix}.patch_embed.proj.weight", + from_pattern=[f"{prefix}.patch_embed.proj.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern=f"{prefix}.patch_embed.proj.bias", + from_pattern=[f"{prefix}.patch_embed.proj.bias"], + ), + WeightTarget( + to_pattern=f"{prefix}.norm.weight", + from_pattern=[f"{prefix}.norm.weight"], + ), + WeightTarget( + to_pattern=f"{prefix}.norm.bias", + from_pattern=[f"{prefix}.norm.bias"], + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.norm1.weight", + from_pattern=[f"{prefix}.blocks.{{block}}.norm1.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.norm1.bias", + from_pattern=[f"{prefix}.blocks.{{block}}.norm1.bias"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.attn.qkv.weight", + from_pattern=[f"{prefix}.blocks.{{block}}.attn.qkv.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.attn.qkv.bias", + from_pattern=[f"{prefix}.blocks.{{block}}.attn.qkv.bias"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.attn.proj.weight", + from_pattern=[f"{prefix}.blocks.{{block}}.attn.proj.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.attn.proj.bias", + from_pattern=[f"{prefix}.blocks.{{block}}.attn.proj.bias"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.ls1.gamma", + from_pattern=[f"{prefix}.blocks.{{block}}.ls1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.norm2.weight", + from_pattern=[f"{prefix}.blocks.{{block}}.norm2.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.norm2.bias", + from_pattern=[f"{prefix}.blocks.{{block}}.norm2.bias"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.mlp.fc1.weight", + from_pattern=[f"{prefix}.blocks.{{block}}.mlp.fc1.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.mlp.fc1.bias", + from_pattern=[f"{prefix}.blocks.{{block}}.mlp.fc1.bias"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.mlp.fc2.weight", + from_pattern=[f"{prefix}.blocks.{{block}}.mlp.fc2.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.mlp.fc2.bias", + from_pattern=[f"{prefix}.blocks.{{block}}.mlp.fc2.bias"], + max_blocks=24, + ), + WeightTarget( + to_pattern=f"{prefix}.blocks.{{block}}.ls2.gamma", + from_pattern=[f"{prefix}.blocks.{{block}}.ls2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + max_blocks=24, + ), + ] + + @staticmethod + def _get_upsample_block_mapping(block_name: str, num_layers: int) -> List[WeightTarget]: + prefix = f"encoder.{block_name}" + targets = [] + + targets.append( + WeightTarget( + to_pattern=f"{prefix}.layers.0.weight", + from_pattern=[f"{prefix}.0.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ) + ) + + targets.extend( + WeightTarget( + to_pattern=f"{prefix}.layers.{layer}.weight", + from_pattern=[f"{prefix}.{layer}.weight"], + transform=WeightTransforms.transpose_conv_transpose2d_weight, + ) + for layer in range(1, num_layers) + ) + + return targets + + @staticmethod + def _get_encoder_conv_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="encoder.upsample_lowres.weight", + from_pattern=["encoder.upsample_lowres.weight"], + transform=WeightTransforms.transpose_conv_transpose2d_weight, + ), + WeightTarget( + to_pattern="encoder.upsample_lowres.bias", + from_pattern=["encoder.upsample_lowres.bias"], + ), + WeightTarget( + to_pattern="encoder.fuse_lowres.weight", + from_pattern=["encoder.fuse_lowres.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.fuse_lowres.bias", + from_pattern=["encoder.fuse_lowres.bias"], + ), + ] + + @staticmethod + def _get_decoder_mapping() -> List[WeightTarget]: + return ( + DepthProWeightMapping._get_decoder_convs_mapping() + + DepthProWeightMapping._get_fusion_block_mapping(0) + + DepthProWeightMapping._get_fusion_block_mapping(1) + + DepthProWeightMapping._get_fusion_block_mapping(2) + + DepthProWeightMapping._get_fusion_block_mapping(3) + + DepthProWeightMapping._get_fusion_block_mapping(4) + ) + + @staticmethod + def _get_decoder_convs_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="decoder.convs.1.weight", + from_pattern=["decoder.convs.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.convs.2.weight", + from_pattern=["decoder.convs.2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.convs.3.weight", + from_pattern=["decoder.convs.3.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.convs.4.weight", + from_pattern=["decoder.convs.4.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + ] + + @staticmethod + def _get_fusion_block_mapping(i: int) -> List[WeightTarget]: + targets = [ + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet1.residual.1.weight", + from_pattern=[f"decoder.fusions.{i}.resnet1.residual.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet1.residual.1.bias", + from_pattern=[f"decoder.fusions.{i}.resnet1.residual.1.bias"], + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet1.residual.3.weight", + from_pattern=[f"decoder.fusions.{i}.resnet1.residual.3.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet1.residual.3.bias", + from_pattern=[f"decoder.fusions.{i}.resnet1.residual.3.bias"], + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet2.residual.1.weight", + from_pattern=[f"decoder.fusions.{i}.resnet2.residual.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet2.residual.1.bias", + from_pattern=[f"decoder.fusions.{i}.resnet2.residual.1.bias"], + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet2.residual.3.weight", + from_pattern=[f"decoder.fusions.{i}.resnet2.residual.3.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.resnet2.residual.3.bias", + from_pattern=[f"decoder.fusions.{i}.resnet2.residual.3.bias"], + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.out_conv.weight", + from_pattern=[f"decoder.fusions.{i}.out_conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern=f"decoder.fusions.{i}.out_conv.bias", + from_pattern=[f"decoder.fusions.{i}.out_conv.bias"], + ), + ] + if i > 0: + targets.append( + WeightTarget( + to_pattern=f"decoder.fusions.{i}.deconv.weight", + from_pattern=[f"decoder.fusions.{i}.deconv.weight"], + transform=WeightTransforms.transpose_conv_transpose2d_weight, + ) + ) + return targets + + @staticmethod + def _get_head_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="head.convs.0.weight", + from_pattern=["head.0.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="head.convs.0.bias", + from_pattern=["head.0.bias"], + ), + WeightTarget( + to_pattern="head.convs.1.weight", + from_pattern=["head.1.weight"], + transform=WeightTransforms.transpose_conv_transpose2d_weight, + ), + WeightTarget( + to_pattern="head.convs.1.bias", + from_pattern=["head.1.bias"], + ), + WeightTarget( + to_pattern="head.convs.2.weight", + from_pattern=["head.2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="head.convs.2.bias", + from_pattern=["head.2.bias"], + ), + WeightTarget( + to_pattern="head.convs.4.weight", + from_pattern=["head.4.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="head.convs.4.bias", + from_pattern=["head.4.bias"], + ), + ] diff --git a/src/mflux/models/depth_pro/weights/weight_handler_depth_pro.py b/src/mflux/models/depth_pro/weights/weight_handler_depth_pro.py deleted file mode 100644 index 244c853..0000000 --- a/src/mflux/models/depth_pro/weights/weight_handler_depth_pro.py +++ /dev/null @@ -1,131 +0,0 @@ -import logging -import urllib.error -import urllib.request - -import mlx.core as mx -import torch -from mlx.utils import tree_unflatten - -from mflux.models.flux.weights.weight_handler import MetaData -from mflux.models.flux.weights.weight_util import WeightUtil -from mflux.ui.defaults import MFLUX_CACHE_DIR - -logger = logging.getLogger(__name__) - - -class WeightHandlerDepthPro: - def __init__(self, weights: dict, meta_data: MetaData): - self.weights = weights - self.meta_data = meta_data - - @staticmethod - def load_weights() -> "WeightHandlerDepthPro": - model_path = WeightHandlerDepthPro._download_or_get_cached_weights() - pt_weights = torch.load(model_path, map_location="cpu") - weights = WeightHandlerDepthPro._to_mlx_weights(pt_weights) - weights = [WeightUtil.reshape_weights(k, v) for k, v in weights.items()] - weights = WeightUtil.flatten(weights) - weights = tree_unflatten(weights) - return WeightHandlerDepthPro( - weights=weights, - meta_data=MetaData(quantization_level=None) - ) # fmt:off - - @staticmethod - def _to_mlx_weights(pt_weights) -> dict: - mlx_weights = {} - for key, value in pt_weights.items(): - if isinstance(value, torch.Tensor): - mlx_weights[key] = mx.array(value.numpy()) - else: - mlx_weights[key] = value - return mlx_weights - - @staticmethod - def _download_or_get_cached_weights(): - APPLE_MODEL_URL = "https://ml-site.cdn-apple.com/models/depth-pro/depth_pro.pt" - - # 1. Create cache directory for the model - cache_dir = MFLUX_CACHE_DIR / "depth_pro" - - cache_dir.mkdir(parents=True, exist_ok=True) - model_path = cache_dir / "depth_pro.pt" - - # 2. Download if model doesn't exist - if not model_path.exists(): - logger.info("Downloading Depth Pro model from Apple...") - try: - urllib.request.urlretrieve(APPLE_MODEL_URL, model_path) - logger.info(f"Downloaded model to {model_path}") - except (urllib.error.URLError, urllib.error.HTTPError) as e: - logger.error(f"Failed to download model: {e}") - logger.info(f"Please manually download from: {APPLE_MODEL_URL}") - if not model_path.exists(): - raise FileNotFoundError(f"Model file not found at {model_path}") - - return model_path - - @staticmethod - def reposition_encoder_weights(depth_pro_weights, name): - tmp = depth_pro_weights.weights["encoder"][name] - depth_pro_weights.weights["encoder"][name] = {} - depth_pro_weights.weights["encoder"][name]["layers"] = tmp - - @staticmethod - def reposition_head_weights(depth_pro_weights): - tmp = depth_pro_weights.weights["head"] - depth_pro_weights.weights["head"] = {} - depth_pro_weights.weights["head"]["convs"] = tmp - - @staticmethod - def reshape_transposed_convolution_weights(depth_pro_weights): - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 1) - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 2) - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent0", 3) - - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent1", 1) - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample_latent1", 2) - - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample0", 1) - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample1", 1) - WeightHandlerDepthPro._reshape_upsample(depth_pro_weights, "upsample2", 1) - - WeightHandlerDepthPro._reshape_upsample_lowres(depth_pro_weights) - - WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 1) - WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 2) - WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 3) - WeightHandlerDepthPro._reshape_deconv(depth_pro_weights, 4) - - WeightHandlerDepthPro._reshape_head(depth_pro_weights, 1) - - @staticmethod - def _reshape_upsample(depth_pro_weights, name, layer): - tmp = depth_pro_weights.weights["encoder"][name]["layers"][layer]["weight"] - tmp = WeightHandlerDepthPro._reshape(tmp) - depth_pro_weights.weights["encoder"][name]["layers"][layer]["weight"] = tmp - - @staticmethod - def _reshape_upsample_lowres(depth_pro_weights): - tmp = depth_pro_weights.weights["encoder"]["upsample_lowres"]["weight"] - tmp = WeightHandlerDepthPro._reshape(tmp) - depth_pro_weights.weights["encoder"]["upsample_lowres"]["weight"] = tmp - - @staticmethod - def _reshape_deconv(depth_pro_weights, layer): - tmp = depth_pro_weights.weights["decoder"]["fusions"][layer]["deconv"]["weight"] - tmp = WeightHandlerDepthPro._reshape(tmp) - depth_pro_weights.weights["decoder"]["fusions"][layer]["deconv"]["weight"] = tmp - - @staticmethod - def _reshape_head(depth_pro_weights, layer): - tmp = depth_pro_weights.weights["head"]["convs"][layer]["weight"] - tmp = WeightHandlerDepthPro._reshape(tmp) - depth_pro_weights.weights["head"]["convs"][layer]["weight"] = tmp - - @staticmethod - def _reshape(tensor): - tensor = tensor.transpose(0, 3, 1, 2) - tensor = tensor.transpose(1, 0, 2, 3) - tensor = tensor.transpose(0, 2, 3, 1) - return tensor diff --git a/src/mflux/models/fibo/cli/__init__.py b/src/mflux/models/fibo/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/generate_fibo.py b/src/mflux/models/fibo/cli/fibo_generate.py similarity index 52% rename from src/mflux/generate_fibo.py rename to src/mflux/models/fibo/cli/fibo_generate.py index ac1696b..8feac1b 100644 --- a/src/mflux/generate_fibo.py +++ b/src/mflux/models/fibo/cli/fibo_generate.py @@ -1,17 +1,12 @@ -import gc -import json - -import mlx.core as mx - from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator from mflux.models.fibo.variants.txt2img.fibo import FIBO -from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.models.fibo.variants.txt2img.util import FiboUtil +from mflux.utils.dimension_resolver import DimensionResolver from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -20,7 +15,7 @@ def main(): parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) parser.add_lora_arguments() - parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_generator_arguments(supports_metadata_config=True, supports_dimension_scale_factor=True) parser.add_image_to_image_arguments(required=False) parser.add_output_arguments() args = parser.parse_args() @@ -29,34 +24,42 @@ def main(): if args.guidance is None: args.guidance = ui_defaults.GUIDANCE_SCALE - json_prompt = _get_json_prompt(args) + json_prompt = FiboUtil.get_json_prompt(args, quantize=args.quantize) # 1. Load the FIBO model fibo = FIBO( - model_config=ModelConfig.fibo(), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=fibo) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=fibo, + latent_creator=FiboLatentCreator, + ) try: + # Resolve dimensions (supports ScaleFactor like "2x" when --image-path is provided) + width, height = DimensionResolver.resolve( + width=args.width, + height=args.height, + reference_image_path=args.image_path, + ) + for seed in args.seed: # 3. Generate an image for each seed value image = fibo.generate_image( seed=seed, prompt=json_prompt, - negative_prompt=PromptUtils.get_effective_negative_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - image_strength=args.image_strength, - scheduler="flow_match_euler_discrete", - ), + width=width, + height=height, + guidance=args.guidance, + image_path=args.image_path, + num_inference_steps=args.steps, + image_strength=args.image_strength, + scheduler="flow_match_euler_discrete", + negative_prompt=PromptUtil.read_negative_prompt(args), ) # 4. Save the image @@ -68,20 +71,5 @@ def main(): print(memory_saver.memory_stats()) -def _get_json_prompt(args): - prompt = PromptUtils.get_effective_prompt(args) - - try: - json.loads(prompt) - json_prompt = prompt - except json.JSONDecodeError: - vlm = FiboVLM() - json_prompt = vlm.generate(prompt=prompt, seed=42) - del vlm - gc.collect() - mx.clear_cache() - return json_prompt - - if __name__ == "__main__": main() diff --git a/src/mflux/models/fibo/fibo_initializer.py b/src/mflux/models/fibo/fibo_initializer.py index b41f526..c73640b 100644 --- a/src/mflux/models/fibo/fibo_initializer.py +++ b/src/mflux/models/fibo/fibo_initializer.py @@ -1,47 +1,66 @@ -from mflux.config.model_config import ModelConfig +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.models.common.config import ModelConfig +from mflux.models.common.tokenizer import TokenizerLoader +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_loader import WeightLoader from mflux.models.fibo.model.fibo_text_encoder import SmolLM3_3B_TextEncoder from mflux.models.fibo.model.fibo_transformer import FiboTransformer from mflux.models.fibo.model.fibo_vae.wan_2_2_vae import Wan2_2_VAE -from mflux.models.fibo.tokenizer import FiboTokenizerHandler -from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler -from mflux.models.fibo.weights.fibo_weight_util import FIBOWeightUtil +from mflux.models.fibo.weights.fibo_weight_definition import FIBOWeightDefinition class FIBOInitializer: @staticmethod def init( - fibo_model, - model_config: ModelConfig | None = None, + model, + model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, + lora_paths: list[str] | None = None, # noqa: ARG004 + lora_scales: list[float] | None = None, # noqa: ARG004 ) -> None: - # 1. Load VAE weights - weights = FIBOWeightHandler.load_regular_weights( - repo_id=model_config.model_name, - local_path=local_path, + path = model_path if model_path else model_config.model_name + FIBOInitializer._init_config(model, model_config) + weights = FIBOInitializer._load_weights(path) + FIBOInitializer._init_tokenizers(model, path) + FIBOInitializer._init_models(model) + FIBOInitializer._apply_weights(model, weights, quantize) + + @staticmethod + def _init_config(model, model_config: ModelConfig) -> None: + model.model_config = model_config + model.callbacks = CallbackRegistry() + + @staticmethod + def _load_weights(model_path: str) -> LoadedWeights: + return WeightLoader.load( + weight_definition=FIBOWeightDefinition, + model_path=model_path, ) - # 2. Initialize tokenizers - tokenizer_handler = FiboTokenizerHandler( - repo_id=model_config.model_name, - local_path=local_path, - ) - fibo_model.fibo_tokenizer = tokenizer_handler.fibo - - # 3. Initialize all models - fibo_model.vae = Wan2_2_VAE() - fibo_model.text_encoder = SmolLM3_3B_TextEncoder() - fibo_model.transformer = FiboTransformer( - in_channels=48, - num_layers=8, - num_single_layers=38, + @staticmethod + def _init_tokenizers(model, model_path: str) -> None: + model.tokenizers = TokenizerLoader.load_all( + definitions=FIBOWeightDefinition.get_tokenizers(), + model_path=model_path, ) - # 4. Apply weights and quantize VAE, transformer, and text encoder - fibo_model.bits = FIBOWeightUtil.set_weights_and_quantize( - quantize_arg=quantize, + @staticmethod + def _init_models(model) -> None: + model.vae = Wan2_2_VAE() + model.text_encoder = SmolLM3_3B_TextEncoder() + model.transformer = FiboTransformer() + + @staticmethod + def _apply_weights(model, weights: LoadedWeights, quantize: int | None) -> None: + model.bits = WeightApplier.apply_and_quantize( weights=weights, - vae=fibo_model.vae, - transformer=fibo_model.transformer, - text_encoder=fibo_model.text_encoder, + quantize_arg=quantize, + weight_definition=FIBOWeightDefinition, + models={ + "vae": model.vae, + "transformer": model.transformer, + "text_encoder": model.text_encoder, + }, ) diff --git a/src/mflux/models/fibo/latent_creator/fibo_latent_creator.py b/src/mflux/models/fibo/latent_creator/fibo_latent_creator.py index d07f9f2..687d840 100644 --- a/src/mflux/models/fibo/latent_creator/fibo_latent_creator.py +++ b/src/mflux/models/fibo/latent_creator/fibo_latent_creator.py @@ -17,3 +17,13 @@ class FiboLatentCreator: batch_size, channels, latent_height, latent_width = latents.shape latents = mx.transpose(latents, (0, 2, 3, 1)) return mx.reshape(latents, (batch_size, latent_height * latent_width, channels)) + + @staticmethod + def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: + batch_size, seq_len, channels = latents.shape + vae_scale_factor = 16 + latent_height = height // vae_scale_factor + latent_width = width // vae_scale_factor + latents = mx.reshape(latents, (batch_size, latent_height, latent_width, channels)) + latents = mx.transpose(latents, (0, 3, 1, 2)) + return latents diff --git a/src/mflux/models/fibo/model/__init__.py b/src/mflux/models/fibo/model/__init__.py index 55a5bd5..e69de29 100644 --- a/src/mflux/models/fibo/model/__init__.py +++ b/src/mflux/models/fibo/model/__init__.py @@ -1 +0,0 @@ -"""FIBO model components.""" diff --git a/src/mflux/models/fibo/model/fibo_text_encoder/prompt_encoder.py b/src/mflux/models/fibo/model/fibo_text_encoder/prompt_encoder.py index b61d7a9..0564222 100644 --- a/src/mflux/models/fibo/model/fibo_text_encoder/prompt_encoder.py +++ b/src/mflux/models/fibo/model/fibo_text_encoder/prompt_encoder.py @@ -3,8 +3,8 @@ from typing import List, Union import mlx.core as mx +from mflux.models.common.tokenizer import Tokenizer from mflux.models.fibo.model.fibo_text_encoder.smol_lm3_3b_text_encoder import SmolLM3_3B_TextEncoder -from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo class PromptEncoder: @@ -12,7 +12,7 @@ class PromptEncoder: def encode_prompt( prompt: str, negative_prompt: str | None, - tokenizer: TokenizerFibo, + tokenizer: Tokenizer, text_encoder: SmolLM3_3B_TextEncoder, ) -> tuple[str, mx.array, List[mx.array]]: # 0. Set default negative prompt if not provided @@ -87,7 +87,7 @@ class PromptEncoder: def _get_prompt_embeds( prompt: Union[str, List[str]], text_encoder: SmolLM3_3B_TextEncoder, - tokenizer: TokenizerFibo, + tokenizer: Tokenizer, num_images_per_prompt: int = 1, max_sequence_length: int = 2048, tokenization_prefix: str | None = None, @@ -97,13 +97,9 @@ class PromptEncoder: raise ValueError("`prompt` must be a non-empty string or list of strings.") # 1) Tokenize and convert to MX - input_ids_mx, attention_mask_mx = tokenizer.tokenize( - prompts=prompts, - max_length=max_sequence_length, - padding="longest", - truncation=True, - add_special_tokens=True, - ) + tokenizer_output = tokenizer.tokenize(prompt=prompts, max_length=max_sequence_length) + input_ids_mx = tokenizer_output.input_ids + attention_mask_mx = tokenizer_output.attention_mask # 2) Run MLX text encoder and collect hidden states hidden_states_list = text_encoder( diff --git a/src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_attention.py b/src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_attention.py index 5e779dd..b4f9276 100644 --- a/src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_attention.py +++ b/src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_attention.py @@ -58,9 +58,6 @@ class SmolLM3_3B_SelfAttention(nn.Module): k = k.astype(hidden_states.dtype) v = v.astype(hidden_states.dtype) - # Force evaluation to ensure computation is complete - mx.eval(q, k, v) - # Reshape to (batch, heads, seq, head_dim) q = q.reshape(batch_size, seq_len, self.num_attention_heads, self.head_dim).transpose(0, 2, 1, 3) k = k.reshape(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(0, 2, 1, 3) diff --git a/src/mflux/models/fibo/model/fibo_transformer/transformer.py b/src/mflux/models/fibo/model/fibo_transformer/transformer.py index 17ba14a..ed70564 100644 --- a/src/mflux/models/fibo/model/fibo_transformer/transformer.py +++ b/src/mflux/models/fibo/model/fibo_transformer/transformer.py @@ -1,7 +1,7 @@ import mlx.core as mx from mlx import nn -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config from mflux.models.fibo.model.fibo_transformer.fibo_embed_nd import FiboEmbedND from mflux.models.fibo.model.fibo_transformer.joint_transformer_block import FiboJointTransformerBlock from mflux.models.fibo.model.fibo_transformer.single_transformer_block import FiboSingleTransformerBlock @@ -31,7 +31,7 @@ class FiboTransformer(nn.Module): def __call__( self, t: int, - config: RuntimeConfig, + config: Config, hidden_states: mx.array, encoder_hidden_states: mx.array, text_encoder_layers: list[mx.array], @@ -150,7 +150,7 @@ class FiboTransformer(nn.Module): @staticmethod def _compute_time_embeddings( t: int, - config: RuntimeConfig, + config: Config, batch_size: int, dtype: mx.Dtype, time_embed: BriaFiboTimestepProjEmbeddings, @@ -163,7 +163,7 @@ class FiboTransformer(nn.Module): def _compute_rotary_embeddings( encoder_hidden_states: mx.array, pos_embed: FiboEmbedND, - config: RuntimeConfig, + config: Config, dtype: mx.Dtype, ) -> mx.array: max_tokens = encoder_hidden_states.shape[1] @@ -209,7 +209,7 @@ class FiboTransformer(nn.Module): @staticmethod def _compute_attention_mask( batch_size: int, - config: RuntimeConfig, + config: Config, encoder_hidden_states: mx.array, max_tokens: int, ) -> mx.array: diff --git a/src/mflux/models/fibo/model/fibo_vae/__init__.py b/src/mflux/models/fibo/model/fibo_vae/__init__.py index 2136841..e69de29 100644 --- a/src/mflux/models/fibo/model/fibo_vae/__init__.py +++ b/src/mflux/models/fibo/model/fibo_vae/__init__.py @@ -1 +0,0 @@ -"""FIBO VAE components.""" diff --git a/src/mflux/models/fibo/model/fibo_vae/common/__init__.py b/src/mflux/models/fibo/model/fibo_vae/common/__init__.py index 0683107..db586b8 100644 --- a/src/mflux/models/fibo/model/fibo_vae/common/__init__.py +++ b/src/mflux/models/fibo/model/fibo_vae/common/__init__.py @@ -1,5 +1,3 @@ -"""FIBO VAE common/shared components.""" - from mflux.models.fibo.model.fibo_vae.common.wan_2_2_attention_block import Wan2_2_AttentionBlock from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d from mflux.models.fibo.model.fibo_vae.common.wan_2_2_mid_block import Wan2_2_MidBlock diff --git a/src/mflux/models/fibo/model/fibo_vae/decoder/__init__.py b/src/mflux/models/fibo/model/fibo_vae/decoder/__init__.py index 76a4c21..c298475 100644 --- a/src/mflux/models/fibo/model/fibo_vae/decoder/__init__.py +++ b/src/mflux/models/fibo/model/fibo_vae/decoder/__init__.py @@ -1,5 +1,3 @@ -"""FIBO VAE decoder components.""" - from mflux.models.fibo.model.fibo_vae.decoder.wan_2_2_decoder_3d import Wan2_2_Decoder3d __all__ = ["Wan2_2_Decoder3d"] diff --git a/src/mflux/models/fibo/model/fibo_vae/encoder/__init__.py b/src/mflux/models/fibo/model/fibo_vae/encoder/__init__.py index 1ba10c4..ba1ab80 100644 --- a/src/mflux/models/fibo/model/fibo_vae/encoder/__init__.py +++ b/src/mflux/models/fibo/model/fibo_vae/encoder/__init__.py @@ -1,5 +1,3 @@ -"""FIBO VAE encoder components.""" - from mflux.models.fibo.model.fibo_vae.encoder.wan_2_2_encoder_3d import Wan2_2_Encoder3d __all__ = ["Wan2_2_Encoder3d"] diff --git a/src/mflux/models/fibo/tokenizer/__init__.py b/src/mflux/models/fibo/tokenizer/__init__.py deleted file mode 100644 index 8c086d5..0000000 --- a/src/mflux/models/fibo/tokenizer/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo -from mflux.models.fibo.tokenizer.smol_lm3_3b_tokenizer import FiboTokenizerHandler - -__all__ = [ - "FiboTokenizerHandler", - "TokenizerFibo", -] diff --git a/src/mflux/models/fibo/tokenizer/fibo_tokenizer.py b/src/mflux/models/fibo/tokenizer/fibo_tokenizer.py deleted file mode 100644 index 461a49d..0000000 --- a/src/mflux/models/fibo/tokenizer/fibo_tokenizer.py +++ /dev/null @@ -1,37 +0,0 @@ -import mlx.core as mx -import numpy as np -from transformers import PreTrainedTokenizer - - -class TokenizerFibo: - def __init__(self, tokenizer: PreTrainedTokenizer, bot_token_id: int = 128000): - self.tokenizer = tokenizer - self.bot_token_id = bot_token_id - - def tokenize( - self, - prompts: list[str], - max_length: int = 2048, - padding: str = "longest", - truncation: bool = True, - add_special_tokens: bool = True, - ) -> tuple[mx.array, mx.array]: - prompts = [p if p is not None else "" for p in prompts] - - if all(p == "" for p in prompts): - batch_size = len(prompts) - input_ids_mx = mx.array(np.empty((batch_size, 0), dtype=np.int32)) - attention_mask_mx = mx.array(np.empty((batch_size, 0), dtype=np.int32)) - else: - tokenized = self.tokenizer( - prompts, - padding=padding, - max_length=max_length, - truncation=truncation, - add_special_tokens=add_special_tokens, - return_tensors="mlx", - ) - input_ids_mx = tokenized["input_ids"] - attention_mask_mx = tokenized["attention_mask"] - - return input_ids_mx, attention_mask_mx diff --git a/src/mflux/models/fibo/tokenizer/smol_lm3_3b_tokenizer.py b/src/mflux/models/fibo/tokenizer/smol_lm3_3b_tokenizer.py deleted file mode 100644 index b7aa04e..0000000 --- a/src/mflux/models/fibo/tokenizer/smol_lm3_3b_tokenizer.py +++ /dev/null @@ -1,42 +0,0 @@ -from pathlib import Path - -import transformers - -from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo -from mflux.utils.download import snapshot_download - - -class FiboTokenizerHandler: - def __init__( - self, - repo_id: str, - bot_token_id: int = 128000, - local_path: str | None = None, - ): - root_path = Path(local_path) if local_path else FiboTokenizerHandler._download_or_get_cached_tokenizer(repo_id) - - # Try different possible tokenizer paths - tokenizer_path = root_path / "tokenizer" - if not tokenizer_path.exists(): - tokenizer_path = root_path / "text_encoder" - if not tokenizer_path.exists(): - tokenizer_path = root_path - - # Load the raw tokenizer - self.tokenizer_raw = transformers.AutoTokenizer.from_pretrained( - pretrained_model_name_or_path=str(tokenizer_path), - local_files_only=True, - fix_mistral_regex=True, # Fix Mistral regex pattern warning - ) - - # Wrap in our TokenizerFibo class - self.fibo = TokenizerFibo(self.tokenizer_raw, bot_token_id=bot_token_id) - - @staticmethod - def _download_or_get_cached_tokenizer(repo_id: str) -> Path: - return Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=["tokenizer/**", "text_encoder/**"], - ) - ) diff --git a/src/mflux/models/fibo/variants/__init__.py b/src/mflux/models/fibo/variants/__init__.py index e27de1f..e69de29 100644 --- a/src/mflux/models/fibo/variants/__init__.py +++ b/src/mflux/models/fibo/variants/__init__.py @@ -1 +0,0 @@ -"""FIBO model variants.""" diff --git a/src/mflux/models/fibo/variants/txt2img/__init__.py b/src/mflux/models/fibo/variants/txt2img/__init__.py index 38b2180..e69de29 100644 --- a/src/mflux/models/fibo/variants/txt2img/__init__.py +++ b/src/mflux/models/fibo/variants/txt2img/__init__.py @@ -1 +0,0 @@ -"""FIBO text-to-image variant.""" diff --git a/src/mflux/models/fibo/variants/txt2img/fibo.py b/src/mflux/models/fibo/variants/txt2img/fibo.py index 7babb7c..1423744 100644 --- a/src/mflux/models/fibo/variants/txt2img/fibo.py +++ b/src/mflux/models/fibo/variants/txt2img/fibo.py @@ -1,20 +1,19 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator -from mflux.models.common.weights.model_saver import ModelSaver +from mflux.models.common.weights.saving.model_saver import ModelSaver from mflux.models.fibo.fibo_initializer import FIBOInitializer from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator from mflux.models.fibo.model.fibo_text_encoder.prompt_encoder import PromptEncoder from mflux.models.fibo.model.fibo_text_encoder.smol_lm3_3b_text_encoder import SmolLM3_3B_TextEncoder from mflux.models.fibo.model.fibo_transformer import FiboTransformer from mflux.models.fibo.model.fibo_vae.wan_2_2_vae import Wan2_2_VAE -from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo +from mflux.models.fibo.weights.fibo_weight_definition import FIBOWeightDefinition from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -24,48 +23,61 @@ class FIBO(nn.Module): vae: Wan2_2_VAE transformer: FiboTransformer text_encoder: SmolLM3_3B_TextEncoder - fibo_tokenizer: TokenizerFibo def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.fibo(), ): super().__init__() - self.model_config = model_config - self.bits = quantize - self.local_path = local_path - FIBOInitializer.init( - fibo_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", negative_prompt: str | None = None, ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - runtime_config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(runtime_config.init_time_step, runtime_config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = LatentCreator.create_for_txt2img_or_img2img( seed=seed, - height=runtime_config.height, - width=runtime_config.width, + width=config.width, + height=config.height, img2img=Img2Img( vae=self.vae, latent_creator=FiboLatentCreator, - image_path=runtime_config.image_path, - sigmas=runtime_config.scheduler.sigmas, - init_time_step=runtime_config.init_time_step, + image_path=config.image_path, + sigmas=config.scheduler.sigmas, + init_time_step=config.init_time_step, ), ) @@ -73,85 +85,58 @@ class FIBO(nn.Module): json_prompt, encoder_hidden_states, text_encoder_layers = PromptEncoder.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, - tokenizer=self.fibo_tokenizer, + tokenizer=self.tokenizers["fibo"], text_encoder=self.text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=json_prompt, - latents=latents, - config=runtime_config, - ) + # 3. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=json_prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: - # 3.t Predict the noise + # 4.t Predict the noise noise = self.transformer( t=t, - config=runtime_config, + config=config, hidden_states=latents, - encoder_hidden_states=encoder_hidden_states, text_encoder_layers=text_encoder_layers, + encoder_hidden_states=encoder_hidden_states, ) - noise = FIBO._apply_classifier_free_guidance(noise, runtime_config.guidance) + noise = FIBO._apply_classifier_free_guidance(noise, config.guidance) - # 4.t Take one denoise step - latents = runtime_config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 5.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=json_prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, - ) + # 6.t Call subscribers in-loop + ctx.in_loop(t, latents) # (Optional) Evaluate to enable progress tracking mx.eval(latents) except KeyboardInterrupt: # noqa: PERF203 - Callbacks.interruption( - t=t, - seed=seed, - prompt=json_prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, - ) + ctx.interruption(t, latents) raise StopImageGenerationException( - f"Stopping image generation at step {t + 1}/{runtime_config.num_inference_steps}" + f"Stopping image generation at step {t + 1}/{config.num_inference_steps}" ) - # (Optional) Call subscribers after loop - Callbacks.after_loop( - seed=seed, - prompt=json_prompt, - latents=latents, - config=runtime_config, - ) + # 7. Call subscribers after loop + ctx.after_loop(latents) - # 5. Decode the latent array and return the image - latents = FIBO._unpack_latents(latents, runtime_config.height, runtime_config.width) + # 8. Decode the latent array and return the image + latents = FiboLatentCreator.unpack_latents(latents, config.height, config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, - config=runtime_config, + config=config, seed=seed, prompt=json_prompt, quantization=self.bits, lora_paths=None, lora_scales=None, - image_path=runtime_config.image_path, - image_strength=runtime_config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + image_path=config.image_path, + image_strength=config.image_strength, + generation_time=config.time_steps.format_dict["elapsed"], ) @staticmethod @@ -161,27 +146,10 @@ class FIBO(nn.Module): noise_text = noise[half:] return noise_uncond + guidance * (noise_text - noise_uncond) - @staticmethod - def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: - batch_size, seq_len, channels = latents.shape - vae_scale_factor = 16 - latent_height = height // vae_scale_factor - latent_width = width // vae_scale_factor - latents = mx.reshape(latents, (batch_size, latent_height, latent_width, channels)) - latents = mx.transpose(latents, (0, 3, 1, 2)) - return latents - def save_model(self, base_path: str) -> None: ModelSaver.save_model( model=self, bits=self.bits, base_path=base_path, - tokenizers=[ - ("fibo_tokenizer.tokenizer", "tokenizer"), - ], - components=[ - ("vae", "vae"), - ("transformer", "transformer"), - ("text_encoder", "text_encoder"), - ], + weight_definition=FIBOWeightDefinition, ) diff --git a/src/mflux/models/fibo/variants/txt2img/util.py b/src/mflux/models/fibo/variants/txt2img/util.py new file mode 100644 index 0000000..775d810 --- /dev/null +++ b/src/mflux/models/fibo/variants/txt2img/util.py @@ -0,0 +1,24 @@ +import gc +import json + +import mlx.core as mx + +from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM +from mflux.utils.prompt_util import PromptUtil + + +class FiboUtil: + @staticmethod + def get_json_prompt(args, quantize: int | None): + prompt = PromptUtil.read_prompt(args) + + try: + json.loads(prompt) + json_prompt = prompt + except json.JSONDecodeError: + vlm = FiboVLM(quantize=quantize) + json_prompt = vlm.generate(prompt=prompt, seed=42) + del vlm + gc.collect() + mx.clear_cache() + return json_prompt diff --git a/src/mflux/models/fibo/weights/__init__.py b/src/mflux/models/fibo/weights/__init__.py index 5ec85eb..001dd3e 100644 --- a/src/mflux/models/fibo/weights/__init__.py +++ b/src/mflux/models/fibo/weights/__init__.py @@ -1,7 +1,4 @@ -"""FIBO weight handling.""" - -from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler +from mflux.models.fibo.weights.fibo_weight_definition import FIBOWeightDefinition from mflux.models.fibo.weights.fibo_weight_mapping import FIBOWeightMapping -from mflux.models.fibo.weights.fibo_weight_util import FIBOWeightUtil -__all__ = ["FIBOWeightHandler", "FIBOWeightMapping", "FIBOWeightUtil"] +__all__ = ["FIBOWeightDefinition", "FIBOWeightMapping"] diff --git a/src/mflux/models/fibo/weights/fibo_weight_definition.py b/src/mflux/models/fibo/weights/fibo_weight_definition.py new file mode 100644 index 0000000..906629c --- /dev/null +++ b/src/mflux/models/fibo/weights/fibo_weight_definition.py @@ -0,0 +1,78 @@ +from typing import List + +import mlx.nn as nn + +from mflux.models.common.tokenizer import LanguageTokenizer +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition, TokenizerDefinition +from mflux.models.fibo.weights.fibo_weight_mapping import FIBOWeightMapping + + +class FIBOWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="vae", + hf_subdir="vae", + num_blocks=4, + loading_mode="torch_convert", + mapping_getter=FIBOWeightMapping.get_vae_mapping, + ), + ComponentDefinition( + name="transformer", + hf_subdir="transformer", + num_blocks=38, + num_layers=46, + loading_mode="torch_convert", + mapping_getter=FIBOWeightMapping.get_transformer_mapping, + ), + ComponentDefinition( + name="text_encoder", + hf_subdir="text_encoder", + num_blocks=36, + loading_mode="torch_convert", + mapping_getter=FIBOWeightMapping.get_text_encoder_mapping, + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [ + TokenizerDefinition( + name="fibo", + hf_subdir="tokenizer", + tokenizer_class="AutoTokenizer", + encoder_class=LanguageTokenizer, + max_length=2048, + padding="longest", + fallback_subdirs=["text_encoder", "."], + download_patterns=["tokenizer/**", "text_encoder/**"], + ), + ] + + @staticmethod + def get_download_patterns() -> List[str]: + return [ + "vae/*.safetensors", + "vae/*.json", + "transformer/*.safetensors", + "transformer/*.json", + "text_encoder/*.safetensors", + "text_encoder/*.json", + ] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + # 1. Skip Conv2d layers + if isinstance(module, nn.Conv2d): + return False + + # 2. Skip any layer with incompatible dimensions + if hasattr(module, "weight") and hasattr(module.weight, "shape"): + if module.weight.shape == (1152, 4304): + return False + if module.weight.shape[-1] % 64 != 0: + return False + + # Only quantize layers that have to_quantized method + return hasattr(module, "to_quantized") diff --git a/src/mflux/models/fibo/weights/fibo_weight_handler.py b/src/mflux/models/fibo/weights/fibo_weight_handler.py deleted file mode 100644 index d3b8a10..0000000 --- a/src/mflux/models/fibo/weights/fibo_weight_handler.py +++ /dev/null @@ -1,179 +0,0 @@ -from pathlib import Path - -import mlx.core as mx -import torch -from mlx.utils import tree_unflatten -from safetensors.torch import load_file as torch_load_file - -from mflux.models.common.weights.mapping.weight_mapper import WeightMapper -from mflux.models.fibo.weights.fibo_weight_mapping import FIBOWeightMapping -from mflux.models.flux.weights.weight_handler import ( - MetaData, - WeightHandler as FluxWeightHandler, -) - - -class FIBOWeightHandler: - def __init__( - self, - meta_data: MetaData, - vae: dict | None = None, - transformer: dict | None = None, - text_encoder: dict | None = None, - decoder: dict | None = None, - visual: dict | None = None, - config: dict | None = None, - ): - self.vae = vae - self.transformer = transformer - self.text_encoder = text_encoder - self.decoder = decoder - self.visual = visual - self.config = config - self.meta_data = meta_data - - @staticmethod - def load_regular_weights( - repo_id: str | None = None, - local_path: str | None = None, - ) -> "FIBOWeightHandler": - root_path: Path | None = None - if local_path: - root_path = Path(local_path) - elif repo_id: - root_path = FluxWeightHandler.download_or_get_cached_weights(repo_id) - - vae_weights = None - transformer_weights = None - text_encoder_weights = None - quantization_level: int | None = None - mflux_version: str | None = None - - if root_path is not None: - vae_weights, _, _ = FIBOWeightHandler._try_load_saved_component(root_path, "vae") - text_encoder_weights, _, _ = FIBOWeightHandler._try_load_saved_component(root_path, "text_encoder") # fmt: off - transformer_weights, quantization_level, mflux_version = FIBOWeightHandler._try_load_saved_component(root_path, "transformer") # fmt: off - - if vae_weights is None or transformer_weights is None or text_encoder_weights is None: - vae_weights = FIBOWeightHandler._load_vae_weights(repo_id, local_path) - transformer_weights = FIBOWeightHandler._load_transformer_weights(repo_id, local_path) - text_encoder_weights = FIBOWeightHandler._load_text_encoder_weights(repo_id, local_path) - quantization_level = None - mflux_version = None - - return FIBOWeightHandler( - vae=vae_weights, - transformer=transformer_weights, - text_encoder=text_encoder_weights, - meta_data=MetaData( - quantization_level=quantization_level, - scale=None, - is_lora=False, - mflux_version=mflux_version, - ), - ) - - @staticmethod - def _load_vae_weights( - repo_id: str | None = None, - local_path: str | None = None, - ) -> dict: - root_path = FIBOWeightHandler._get_root_path(repo_id, local_path) - vae_path = root_path / "vae" - raw_weights = FIBOWeightHandler._load_safetensors_shards(vae_path) - mapping = FIBOWeightMapping.get_vae_mapping() - mapped_weights = WeightMapper.apply_mapping(raw_weights, mapping, num_blocks=4) - return mapped_weights - - @staticmethod - def _load_transformer_weights( - repo_id: str | None = None, - local_path: str | None = None, - ) -> dict: - root_path = FIBOWeightHandler._get_root_path(repo_id, local_path) - transformer_path = root_path / "transformer" - if transformer_path.exists() and list(transformer_path.glob("*.safetensors")): - raw_weights = FIBOWeightHandler._load_safetensors_shards(transformer_path) - else: - raw_weights = FIBOWeightHandler._load_safetensors_shards(root_path) - mapping = FIBOWeightMapping.get_transformer_mapping() - mapped_weights = WeightMapper.apply_mapping( - raw_weights, - mapping, - num_blocks=38, - num_layers=46, - ) - return mapped_weights - - @staticmethod - def _load_text_encoder_weights( - repo_id: str | None = None, - local_path: str | None = None, - ) -> dict: - root_path = FIBOWeightHandler._get_root_path(repo_id, local_path) - text_encoder_path = root_path / "text_encoder" - raw_weights = FIBOWeightHandler._load_safetensors_shards(text_encoder_path) - mapping = FIBOWeightMapping.get_text_encoder_mapping() - mapped_weights = WeightMapper.apply_mapping( - raw_weights, - mapping, - num_blocks=36, - ) - return mapped_weights - - @staticmethod - def _get_root_path( - repo_id: str | None = None, - local_path: str | None = None, - ) -> Path: - if local_path: - return Path(local_path) - return Path(FluxWeightHandler.download_or_get_cached_weights(repo_id or "briaai/FIBO")) - - @staticmethod - def _load_safetensors_shards(path: Path) -> dict[str, mx.array]: - shard_files = sorted(f for f in path.glob("*.safetensors") if not f.name.startswith("._")) - if not shard_files: - raise FileNotFoundError(f"No safetensors files found in {path}") - - all_weights: dict[str, mx.array] = {} - for shard in shard_files: - torch_weights = torch_load_file(str(shard)) - for key, tensor in torch_weights.items(): - if tensor.dtype == torch.bfloat16: - tensor = tensor.to(torch.float16) - all_weights[key] = mx.array(tensor.numpy()) - - return all_weights - - @staticmethod - def _try_load_saved_component( - root_path: Path, - component_name: str, - ) -> tuple[dict | None, int | None, str | None]: - component_path = root_path / component_name - if not component_path.exists(): - return None, None, None - - shard_files = sorted(f for f in component_path.glob("*.safetensors") if not f.name.startswith("._")) - if not shard_files: - return None, None, None - - all_weights: dict[str, mx.array] = {} - quantization_level: int | None = None - mflux_version: str | None = None - - for idx, shard in enumerate(shard_files): - data = mx.load(str(shard), return_metadata=True) - weights_dict = data[0] - all_weights.update(dict(weights_dict.items())) - - if idx == 0 and len(data) > 1: - quantization_level = data[1].get("quantization_level") - mflux_version = data[1].get("mflux_version") - - if quantization_level is None and mflux_version is None: - return None, None, None - - unflattened = tree_unflatten(list(all_weights.items())) - return unflattened, quantization_level, mflux_version diff --git a/src/mflux/models/fibo/weights/fibo_weight_mapping.py b/src/mflux/models/fibo/weights/fibo_weight_mapping.py index 5032579..9bbce49 100644 --- a/src/mflux/models/fibo/weights/fibo_weight_mapping.py +++ b/src/mflux/models/fibo/weights/fibo_weight_mapping.py @@ -1,278 +1,264 @@ from typing import List from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget -from mflux.models.qwen.weights.qwen_weight_mapping import ( - reshape_gamma_to_1d, - transpose_conv2d_weight, - transpose_conv3d_weight, -) +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms class FIBOWeightMapping(WeightMapping): @staticmethod def get_transformer_mapping() -> List[WeightTarget]: return [ - # ========== Global projections ========== WeightTarget( - mlx_path="time_embed.timestep_embedder.linear_1.weight", - hf_patterns=["time_embed.timestep_embedder.linear_1.weight"], + to_pattern="time_embed.timestep_embedder.linear_1.weight", + from_pattern=["time_embed.timestep_embedder.linear_1.weight"], ), WeightTarget( - mlx_path="time_embed.timestep_embedder.linear_1.bias", - hf_patterns=["time_embed.timestep_embedder.linear_1.bias"], + to_pattern="time_embed.timestep_embedder.linear_1.bias", + from_pattern=["time_embed.timestep_embedder.linear_1.bias"], ), WeightTarget( - mlx_path="time_embed.timestep_embedder.linear_2.weight", - hf_patterns=["time_embed.timestep_embedder.linear_2.weight"], + to_pattern="time_embed.timestep_embedder.linear_2.weight", + from_pattern=["time_embed.timestep_embedder.linear_2.weight"], ), WeightTarget( - mlx_path="time_embed.timestep_embedder.linear_2.bias", - hf_patterns=["time_embed.timestep_embedder.linear_2.bias"], + to_pattern="time_embed.timestep_embedder.linear_2.bias", + from_pattern=["time_embed.timestep_embedder.linear_2.bias"], ), WeightTarget( - mlx_path="context_embedder.weight", - hf_patterns=["context_embedder.weight"], + to_pattern="context_embedder.weight", + from_pattern=["context_embedder.weight"], ), WeightTarget( - mlx_path="context_embedder.bias", - hf_patterns=["context_embedder.bias"], + to_pattern="context_embedder.bias", + from_pattern=["context_embedder.bias"], ), WeightTarget( - mlx_path="x_embedder.weight", - hf_patterns=["x_embedder.weight"], + to_pattern="x_embedder.weight", + from_pattern=["x_embedder.weight"], ), WeightTarget( - mlx_path="x_embedder.bias", - hf_patterns=["x_embedder.bias"], - ), - # ========== Joint transformer blocks ========== - # AdaLayerNormZero (image + context streams) - WeightTarget( - mlx_path="transformer_blocks.{block}.norm1.linear.weight", - hf_patterns=["transformer_blocks.{block}.norm1.linear.weight"], + to_pattern="x_embedder.bias", + from_pattern=["x_embedder.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.norm1.linear.bias", - hf_patterns=["transformer_blocks.{block}.norm1.linear.bias"], + to_pattern="transformer_blocks.{block}.norm1.linear.weight", + from_pattern=["transformer_blocks.{block}.norm1.linear.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.norm1_context.linear.weight", - hf_patterns=["transformer_blocks.{block}.norm1_context.linear.weight"], + to_pattern="transformer_blocks.{block}.norm1.linear.bias", + from_pattern=["transformer_blocks.{block}.norm1.linear.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.norm1_context.linear.bias", - hf_patterns=["transformer_blocks.{block}.norm1_context.linear.bias"], - ), - # Attention weights (BriaFiboAttention) - WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_q.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_q.weight"], + to_pattern="transformer_blocks.{block}.norm1_context.linear.weight", + from_pattern=["transformer_blocks.{block}.norm1_context.linear.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_k.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_k.weight"], + to_pattern="transformer_blocks.{block}.norm1_context.linear.bias", + from_pattern=["transformer_blocks.{block}.norm1_context.linear.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_q.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_q.weight"], + to_pattern="transformer_blocks.{block}.attn.norm_q.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_q.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_q.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_q.bias"], + to_pattern="transformer_blocks.{block}.attn.norm_k.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_k.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_k.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_k.weight"], + to_pattern="transformer_blocks.{block}.attn.to_q.weight", + from_pattern=["transformer_blocks.{block}.attn.to_q.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_k.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_k.bias"], + to_pattern="transformer_blocks.{block}.attn.to_q.bias", + from_pattern=["transformer_blocks.{block}.attn.to_q.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_v.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_v.weight"], + to_pattern="transformer_blocks.{block}.attn.to_k.weight", + from_pattern=["transformer_blocks.{block}.attn.to_k.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_v.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_v.bias"], + to_pattern="transformer_blocks.{block}.attn.to_k.bias", + from_pattern=["transformer_blocks.{block}.attn.to_k.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_out.0.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_out.0.weight"], + to_pattern="transformer_blocks.{block}.attn.to_v.weight", + from_pattern=["transformer_blocks.{block}.attn.to_v.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_out.0.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_out.0.bias"], + to_pattern="transformer_blocks.{block}.attn.to_v.bias", + from_pattern=["transformer_blocks.{block}.attn.to_v.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_added_q.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_added_q.weight"], + to_pattern="transformer_blocks.{block}.attn.to_out.0.weight", + from_pattern=["transformer_blocks.{block}.attn.to_out.0.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_added_k.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_added_k.weight"], + to_pattern="transformer_blocks.{block}.attn.to_out.0.bias", + from_pattern=["transformer_blocks.{block}.attn.to_out.0.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_q_proj.weight", - hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.weight"], + to_pattern="transformer_blocks.{block}.attn.norm_added_q.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_added_q.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_q_proj.bias", - hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.bias"], + to_pattern="transformer_blocks.{block}.attn.norm_added_k.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_added_k.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_k_proj.weight", - hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.weight"], + to_pattern="transformer_blocks.{block}.attn.add_q_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_q_proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_k_proj.bias", - hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.bias"], + to_pattern="transformer_blocks.{block}.attn.add_q_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_q_proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_v_proj.weight", - hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.weight"], + to_pattern="transformer_blocks.{block}.attn.add_k_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_k_proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_v_proj.bias", - hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.bias"], + to_pattern="transformer_blocks.{block}.attn.add_k_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_k_proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_add_out.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_add_out.weight"], + to_pattern="transformer_blocks.{block}.attn.add_v_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_v_proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_add_out.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_add_out.bias"], - ), - # LayerNorm / FFN for image stream - WeightTarget( - mlx_path="transformer_blocks.{block}.norm2.weight", - hf_patterns=["transformer_blocks.{block}.norm2.weight"], + to_pattern="transformer_blocks.{block}.attn.add_v_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_v_proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.norm2.bias", - hf_patterns=["transformer_blocks.{block}.norm2.bias"], + to_pattern="transformer_blocks.{block}.attn.to_add_out.weight", + from_pattern=["transformer_blocks.{block}.attn.to_add_out.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff.net.0.proj.weight", - hf_patterns=["transformer_blocks.{block}.ff.net.0.proj.weight"], + to_pattern="transformer_blocks.{block}.attn.to_add_out.bias", + from_pattern=["transformer_blocks.{block}.attn.to_add_out.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff.net.0.proj.bias", - hf_patterns=["transformer_blocks.{block}.ff.net.0.proj.bias"], + to_pattern="transformer_blocks.{block}.norm2.weight", + from_pattern=["transformer_blocks.{block}.norm2.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff.net.2.weight", - hf_patterns=["transformer_blocks.{block}.ff.net.2.weight"], + to_pattern="transformer_blocks.{block}.norm2.bias", + from_pattern=["transformer_blocks.{block}.norm2.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff.net.2.bias", - hf_patterns=["transformer_blocks.{block}.ff.net.2.bias"], - ), - # LayerNorm / FFN for context stream - WeightTarget( - mlx_path="transformer_blocks.{block}.norm2_context.weight", - hf_patterns=["transformer_blocks.{block}.norm2_context.weight"], + to_pattern="transformer_blocks.{block}.ff.net.0.proj.weight", + from_pattern=["transformer_blocks.{block}.ff.net.0.proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.norm2_context.bias", - hf_patterns=["transformer_blocks.{block}.norm2_context.bias"], + to_pattern="transformer_blocks.{block}.ff.net.0.proj.bias", + from_pattern=["transformer_blocks.{block}.ff.net.0.proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff_context.net.0.proj.weight", - hf_patterns=["transformer_blocks.{block}.ff_context.net.0.proj.weight"], + to_pattern="transformer_blocks.{block}.ff.net.2.weight", + from_pattern=["transformer_blocks.{block}.ff.net.2.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff_context.net.0.proj.bias", - hf_patterns=["transformer_blocks.{block}.ff_context.net.0.proj.bias"], + to_pattern="transformer_blocks.{block}.ff.net.2.bias", + from_pattern=["transformer_blocks.{block}.ff.net.2.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff_context.net.2.weight", - hf_patterns=["transformer_blocks.{block}.ff_context.net.2.weight"], + to_pattern="transformer_blocks.{block}.norm2_context.weight", + from_pattern=["transformer_blocks.{block}.norm2_context.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.ff_context.net.2.bias", - hf_patterns=["transformer_blocks.{block}.ff_context.net.2.bias"], - ), - # ========== Single transformer blocks ========== - WeightTarget( - mlx_path="single_transformer_blocks.{block}.norm.linear.weight", - hf_patterns=["single_transformer_blocks.{block}.norm.linear.weight"], + to_pattern="transformer_blocks.{block}.norm2_context.bias", + from_pattern=["transformer_blocks.{block}.norm2_context.bias"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.norm.linear.bias", - hf_patterns=["single_transformer_blocks.{block}.norm.linear.bias"], + to_pattern="transformer_blocks.{block}.ff_context.net.0.proj.weight", + from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.norm_q.weight", - hf_patterns=["single_transformer_blocks.{block}.attn.norm_q.weight"], + to_pattern="transformer_blocks.{block}.ff_context.net.0.proj.bias", + from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.bias"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.norm_k.weight", - hf_patterns=["single_transformer_blocks.{block}.attn.norm_k.weight"], + to_pattern="transformer_blocks.{block}.ff_context.net.2.weight", + from_pattern=["transformer_blocks.{block}.ff_context.net.2.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.to_q.weight", - hf_patterns=["single_transformer_blocks.{block}.attn.to_q.weight"], + to_pattern="transformer_blocks.{block}.ff_context.net.2.bias", + from_pattern=["transformer_blocks.{block}.ff_context.net.2.bias"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.to_q.bias", - hf_patterns=["single_transformer_blocks.{block}.attn.to_q.bias"], + to_pattern="single_transformer_blocks.{block}.norm.linear.weight", + from_pattern=["single_transformer_blocks.{block}.norm.linear.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.to_k.weight", - hf_patterns=["single_transformer_blocks.{block}.attn.to_k.weight"], + to_pattern="single_transformer_blocks.{block}.norm.linear.bias", + from_pattern=["single_transformer_blocks.{block}.norm.linear.bias"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.to_k.bias", - hf_patterns=["single_transformer_blocks.{block}.attn.to_k.bias"], + to_pattern="single_transformer_blocks.{block}.attn.norm_q.weight", + from_pattern=["single_transformer_blocks.{block}.attn.norm_q.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.to_v.weight", - hf_patterns=["single_transformer_blocks.{block}.attn.to_v.weight"], + to_pattern="single_transformer_blocks.{block}.attn.norm_k.weight", + from_pattern=["single_transformer_blocks.{block}.attn.norm_k.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.attn.to_v.bias", - hf_patterns=["single_transformer_blocks.{block}.attn.to_v.bias"], + to_pattern="single_transformer_blocks.{block}.attn.to_q.weight", + from_pattern=["single_transformer_blocks.{block}.attn.to_q.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.proj_mlp.weight", - hf_patterns=["single_transformer_blocks.{block}.proj_mlp.weight"], + to_pattern="single_transformer_blocks.{block}.attn.to_q.bias", + from_pattern=["single_transformer_blocks.{block}.attn.to_q.bias"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.proj_mlp.bias", - hf_patterns=["single_transformer_blocks.{block}.proj_mlp.bias"], + to_pattern="single_transformer_blocks.{block}.attn.to_k.weight", + from_pattern=["single_transformer_blocks.{block}.attn.to_k.weight"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.proj_out.weight", - hf_patterns=["single_transformer_blocks.{block}.proj_out.weight"], + to_pattern="single_transformer_blocks.{block}.attn.to_k.bias", + from_pattern=["single_transformer_blocks.{block}.attn.to_k.bias"], ), WeightTarget( - mlx_path="single_transformer_blocks.{block}.proj_out.bias", - hf_patterns=["single_transformer_blocks.{block}.proj_out.bias"], - ), - # ========== Caption projection & output head ========== - WeightTarget( - mlx_path="norm_out.linear.weight", - hf_patterns=["norm_out.linear.weight"], + to_pattern="single_transformer_blocks.{block}.attn.to_v.weight", + from_pattern=["single_transformer_blocks.{block}.attn.to_v.weight"], ), WeightTarget( - mlx_path="norm_out.linear.bias", - hf_patterns=["norm_out.linear.bias"], + to_pattern="single_transformer_blocks.{block}.attn.to_v.bias", + from_pattern=["single_transformer_blocks.{block}.attn.to_v.bias"], ), WeightTarget( - mlx_path="proj_out.weight", - hf_patterns=["proj_out.weight"], + to_pattern="single_transformer_blocks.{block}.proj_mlp.weight", + from_pattern=["single_transformer_blocks.{block}.proj_mlp.weight"], ), WeightTarget( - mlx_path="proj_out.bias", - hf_patterns=["proj_out.bias"], + to_pattern="single_transformer_blocks.{block}.proj_mlp.bias", + from_pattern=["single_transformer_blocks.{block}.proj_mlp.bias"], ), - # Caption_projection layers: we rely on explicit layer index to be - # expanded via num_layers; we pass num_layers explicitly from the loader. WeightTarget( - mlx_path="caption_projection.{layer}.linear.weight", - hf_patterns=["caption_projection.{layer}.linear.weight"], + to_pattern="single_transformer_blocks.{block}.proj_out.weight", + from_pattern=["single_transformer_blocks.{block}.proj_out.weight"], + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.proj_out.bias", + from_pattern=["single_transformer_blocks.{block}.proj_out.bias"], + ), + WeightTarget( + to_pattern="norm_out.linear.weight", + from_pattern=["norm_out.linear.weight"], + ), + WeightTarget( + to_pattern="norm_out.linear.bias", + from_pattern=["norm_out.linear.bias"], + ), + WeightTarget( + to_pattern="proj_out.weight", + from_pattern=["proj_out.weight"], + ), + WeightTarget( + to_pattern="proj_out.bias", + from_pattern=["proj_out.bias"], + ), + WeightTarget( + to_pattern="caption_projection.{layer}.linear.weight", + from_pattern=["caption_projection.{layer}.linear.weight"], ), ] @@ -280,356 +266,340 @@ class FIBOWeightMapping(WeightMapping): def get_text_encoder_mapping() -> List[WeightTarget]: return [ WeightTarget( - mlx_path="embed_tokens.weight", - hf_patterns=["model.embed_tokens.weight"], + to_pattern="embed_tokens.weight", + from_pattern=["model.embed_tokens.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.q_proj.weight", - hf_patterns=["model.layers.{block}.self_attn.q_proj.weight"], + to_pattern="layers.{block}.self_attn.q_proj.weight", + from_pattern=["model.layers.{block}.self_attn.q_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.k_proj.weight", - hf_patterns=["model.layers.{block}.self_attn.k_proj.weight"], + to_pattern="layers.{block}.self_attn.k_proj.weight", + from_pattern=["model.layers.{block}.self_attn.k_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.v_proj.weight", - hf_patterns=["model.layers.{block}.self_attn.v_proj.weight"], + to_pattern="layers.{block}.self_attn.v_proj.weight", + from_pattern=["model.layers.{block}.self_attn.v_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.o_proj.weight", - hf_patterns=["model.layers.{block}.self_attn.o_proj.weight"], + to_pattern="layers.{block}.self_attn.o_proj.weight", + from_pattern=["model.layers.{block}.self_attn.o_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.gate_proj.weight", - hf_patterns=["model.layers.{block}.mlp.gate_proj.weight"], + to_pattern="layers.{block}.mlp.gate_proj.weight", + from_pattern=["model.layers.{block}.mlp.gate_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.up_proj.weight", - hf_patterns=["model.layers.{block}.mlp.up_proj.weight"], + to_pattern="layers.{block}.mlp.up_proj.weight", + from_pattern=["model.layers.{block}.mlp.up_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.down_proj.weight", - hf_patterns=["model.layers.{block}.mlp.down_proj.weight"], + to_pattern="layers.{block}.mlp.down_proj.weight", + from_pattern=["model.layers.{block}.mlp.down_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.input_layernorm.weight", - hf_patterns=["model.layers.{block}.input_layernorm.weight"], + to_pattern="layers.{block}.input_layernorm.weight", + from_pattern=["model.layers.{block}.input_layernorm.weight"], ), WeightTarget( - mlx_path="layers.{block}.post_attention_layernorm.weight", - hf_patterns=["model.layers.{block}.post_attention_layernorm.weight"], + to_pattern="layers.{block}.post_attention_layernorm.weight", + from_pattern=["model.layers.{block}.post_attention_layernorm.weight"], ), WeightTarget( - mlx_path="norm.weight", - hf_patterns=["model.norm.weight"], + to_pattern="norm.weight", + from_pattern=["model.norm.weight"], ), ] @staticmethod def get_vae_mapping() -> List[WeightTarget]: return [ - # ========== Encoder conv_in ========== WeightTarget( - mlx_path="encoder.conv_in.conv3d.weight", - hf_patterns=["encoder.conv_in.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.conv_in.conv3d.weight", + from_pattern=["encoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.conv_in.conv3d.bias", - hf_patterns=["encoder.conv_in.bias"], - ), - # ========== Encoder down_blocks ========== - WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.norm1.weight", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.conv_in.conv3d.bias", + from_pattern=["encoder.conv_in.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv1.bias"], + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.norm2.weight", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"], + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv3d_weight, required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias", - hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"], + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"], required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.downsampler.resample_conv.weight", - hf_patterns=["encoder.down_blocks.{block}.downsampler.resample.1.weight"], - transform=transpose_conv2d_weight, + to_pattern="encoder.down_blocks.{block}.downsampler.resample_conv.weight", + from_pattern=["encoder.down_blocks.{block}.downsampler.resample.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.downsampler.resample_conv.bias", - hf_patterns=["encoder.down_blocks.{block}.downsampler.resample.1.bias"], + to_pattern="encoder.down_blocks.{block}.downsampler.resample_conv.bias", + from_pattern=["encoder.down_blocks.{block}.downsampler.resample.1.bias"], required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.weight", - hf_patterns=["encoder.down_blocks.{block}.downsampler.time_conv.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.weight", + from_pattern=["encoder.down_blocks.{block}.downsampler.time_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.bias", - hf_patterns=["encoder.down_blocks.{block}.downsampler.time_conv.bias"], + to_pattern="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.bias", + from_pattern=["encoder.down_blocks.{block}.downsampler.time_conv.bias"], required=False, ), - # ========== Encoder mid_block ========== WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.norm1.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.mid_block.resnets.{i}.conv1.conv3d.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.bias", - hf_patterns=["encoder.mid_block.resnets.{i}.conv1.bias"], + to_pattern="encoder.mid_block.resnets.{i}.conv1.conv3d.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.bias"], ), WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.norm2.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.mid_block.resnets.{i}.conv2.conv3d.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.bias", - hf_patterns=["encoder.mid_block.resnets.{i}.conv2.bias"], + to_pattern="encoder.mid_block.resnets.{i}.conv2.conv3d.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.bias"], ), WeightTarget( - mlx_path="encoder.mid_block.attentions.{i}.norm.weight", - hf_patterns=["encoder.mid_block.attentions.{i}.norm.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.mid_block.attentions.{i}.norm.weight", + from_pattern=["encoder.mid_block.attentions.{i}.norm.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.mid_block.attentions.{i}.to_qkv.weight", - hf_patterns=["encoder.mid_block.attentions.{i}.to_qkv.weight"], - transform=transpose_conv2d_weight, + to_pattern="encoder.mid_block.attentions.{i}.to_qkv.weight", + from_pattern=["encoder.mid_block.attentions.{i}.to_qkv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="encoder.mid_block.attentions.{i}.to_qkv.bias", - hf_patterns=["encoder.mid_block.attentions.{i}.to_qkv.bias"], + to_pattern="encoder.mid_block.attentions.{i}.to_qkv.bias", + from_pattern=["encoder.mid_block.attentions.{i}.to_qkv.bias"], ), WeightTarget( - mlx_path="encoder.mid_block.attentions.{i}.proj.weight", - hf_patterns=["encoder.mid_block.attentions.{i}.proj.weight"], - transform=transpose_conv2d_weight, + to_pattern="encoder.mid_block.attentions.{i}.proj.weight", + from_pattern=["encoder.mid_block.attentions.{i}.proj.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="encoder.mid_block.attentions.{i}.proj.bias", - hf_patterns=["encoder.mid_block.attentions.{i}.proj.bias"], - ), - # ========== Encoder output ========== - WeightTarget( - mlx_path="encoder.norm_out.weight", - hf_patterns=["encoder.norm_out.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.mid_block.attentions.{i}.proj.bias", + from_pattern=["encoder.mid_block.attentions.{i}.proj.bias"], ), WeightTarget( - mlx_path="encoder.conv_out.conv3d.weight", - hf_patterns=["encoder.conv_out.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.norm_out.weight", + from_pattern=["encoder.norm_out.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.conv_out.conv3d.bias", - hf_patterns=["encoder.conv_out.bias"], - ), - # ========== Decoder conv_in ========== - WeightTarget( - mlx_path="decoder.conv_in.conv3d.weight", - hf_patterns=["decoder.conv_in.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.conv_out.conv3d.weight", + from_pattern=["encoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.conv_in.conv3d.bias", - hf_patterns=["decoder.conv_in.bias"], - ), - # ========== Decoder mid_block ========== - # Mid block resnets - WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.norm1.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.conv_out.conv3d.bias", + from_pattern=["encoder.conv_out.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.conv_in.conv3d.weight", + from_pattern=["decoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.bias", - hf_patterns=["decoder.mid_block.resnets.{i}.conv1.bias"], + to_pattern="decoder.conv_in.conv3d.bias", + from_pattern=["decoder.conv_in.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.norm2.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.mid_block.resnets.{i}.conv1.conv3d.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.bias", - hf_patterns=["decoder.mid_block.resnets.{i}.conv2.bias"], - ), - # Mid block attention - WeightTarget( - mlx_path="decoder.mid_block.attentions.{i}.norm.weight", - hf_patterns=["decoder.mid_block.attentions.{i}.norm.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.resnets.{i}.conv1.conv3d.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.attentions.{i}.to_qkv.weight", - hf_patterns=["decoder.mid_block.attentions.{i}.to_qkv.weight"], - transform=transpose_conv2d_weight, + to_pattern="decoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.mid_block.attentions.{i}.to_qkv.bias", - hf_patterns=["decoder.mid_block.attentions.{i}.to_qkv.bias"], + to_pattern="decoder.mid_block.resnets.{i}.conv2.conv3d.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.attentions.{i}.proj.weight", - hf_patterns=["decoder.mid_block.attentions.{i}.proj.weight"], - transform=transpose_conv2d_weight, + to_pattern="decoder.mid_block.resnets.{i}.conv2.conv3d.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.attentions.{i}.proj.bias", - hf_patterns=["decoder.mid_block.attentions.{i}.proj.bias"], - ), - # ========== Decoder up_blocks ========== - # Up blocks resnets - WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.norm1.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.attentions.{i}.norm.weight", + from_pattern=["decoder.mid_block.attentions.{i}.norm.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv1.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.mid_block.attentions.{i}.to_qkv.weight", + from_pattern=["decoder.mid_block.attentions.{i}.to_qkv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv1.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"], + to_pattern="decoder.mid_block.attentions.{i}.to_qkv.bias", + from_pattern=["decoder.mid_block.attentions.{i}.to_qkv.bias"], ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.norm2.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.attentions.{i}.proj.weight", + from_pattern=["decoder.mid_block.attentions.{i}.proj.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv2.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.mid_block.attentions.{i}.proj.bias", + from_pattern=["decoder.mid_block.attentions.{i}.proj.bias"], ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv2.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"], - ), - # Up blocks resnets - conv_shortcut (optional, when in_dim != out_dim) - WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"], - transform=transpose_conv3d_weight, - required=False, # Optional - only exists when dimensions differ + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"], - required=False, # Optional - only exists when dimensions differ - ), - # Up blocks upsamplers - time_conv (blocks 0, 1 only) - WeightTarget( - mlx_path="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.upsampler.time_conv.weight"], - transform=transpose_conv3d_weight, - required=False, # Only exists for blocks 0, 1 + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.upsampler.time_conv.bias"], - required=False, # Only exists for blocks 0, 1 - ), - # Up blocks upsamplers - resample_conv (blocks 0, 1, 2) - WeightTarget( - mlx_path="decoder.up_blocks.{block}.upsampler.resample_conv.weight", - hf_patterns=["decoder.up_blocks.{block}.upsampler.resample.1.weight"], - transform=transpose_conv2d_weight, - required=False, # Only exists for blocks 0, 1, 2 + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"], ), WeightTarget( - mlx_path="decoder.up_blocks.{block}.upsampler.resample_conv.bias", - hf_patterns=["decoder.up_blocks.{block}.upsampler.resample.1.bias"], - required=False, # Only exists for blocks 0, 1, 2 - ), - # ========== Decoder output ========== - WeightTarget( - mlx_path="decoder.norm_out.weight", - hf_patterns=["decoder.norm_out.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.conv_out.conv3d.weight", - hf_patterns=["decoder.conv_out.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.conv_out.conv3d.bias", - hf_patterns=["decoder.conv_out.bias"], - ), - # ========== Quant conv ========== - WeightTarget( - mlx_path="quant_conv.conv3d.weight", - hf_patterns=["quant_conv.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"], ), WeightTarget( - mlx_path="quant_conv.conv3d.bias", - hf_patterns=["quant_conv.bias"], - ), - # ========== Post quant conv ========== - WeightTarget( - mlx_path="post_quant_conv.conv3d.weight", - hf_patterns=["post_quant_conv.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + required=False, ), WeightTarget( - mlx_path="post_quant_conv.conv3d.bias", - hf_patterns=["post_quant_conv.bias"], + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"], + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.upsampler.time_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.upsampler.time_conv.bias"], + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsampler.resample_conv.weight", + from_pattern=["decoder.up_blocks.{block}.upsampler.resample.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsampler.resample_conv.bias", + from_pattern=["decoder.up_blocks.{block}.upsampler.resample.1.bias"], + required=False, + ), + WeightTarget( + to_pattern="decoder.norm_out.weight", + from_pattern=["decoder.norm_out.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="decoder.conv_out.conv3d.weight", + from_pattern=["decoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="decoder.conv_out.conv3d.bias", + from_pattern=["decoder.conv_out.bias"], + ), + WeightTarget( + to_pattern="quant_conv.conv3d.weight", + from_pattern=["quant_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="quant_conv.conv3d.bias", + from_pattern=["quant_conv.bias"], + ), + WeightTarget( + to_pattern="post_quant_conv.conv3d.weight", + from_pattern=["post_quant_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="post_quant_conv.conv3d.bias", + from_pattern=["post_quant_conv.bias"], ), ] diff --git a/src/mflux/models/fibo/weights/fibo_weight_util.py b/src/mflux/models/fibo/weights/fibo_weight_util.py deleted file mode 100644 index d386ed9..0000000 --- a/src/mflux/models/fibo/weights/fibo_weight_util.py +++ /dev/null @@ -1,47 +0,0 @@ -from typing import TYPE_CHECKING - -import mlx.nn as nn - -from mflux.models.common.quantization.quantization_util import QuantizationUtil - -if TYPE_CHECKING: - from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler - - -class FIBOWeightUtil: - @staticmethod - def set_weights_and_quantize( - quantize_arg: int | None, - weights: "FIBOWeightHandler", - vae: nn.Module, - transformer: nn.Module, - text_encoder: nn.Module | None = None, - ) -> int | None: - if weights.meta_data.quantization_level is None and quantize_arg is None: - FIBOWeightUtil._set_model_weights(weights, vae, transformer, text_encoder) - return None - - if weights.meta_data.quantization_level is None and quantize_arg is not None: - bits = quantize_arg - FIBOWeightUtil._set_model_weights(weights, vae, transformer, text_encoder) - QuantizationUtil.quantize_fibo_models(vae, transformer, text_encoder, bits, weights) - return bits - - if weights.meta_data.quantization_level is not None: - bits = weights.meta_data.quantization_level - QuantizationUtil.quantize_fibo_models(vae, transformer, text_encoder, bits, weights) - FIBOWeightUtil._set_model_weights(weights, vae, transformer, text_encoder) - return bits - - raise Exception("Error setting weights") - - @staticmethod - def _set_model_weights( - weights: "FIBOWeightHandler", - vae: nn.Module, - transformer: nn.Module, - text_encoder: nn.Module | None = None, - ): - vae.update(weights.vae, strict=False) - transformer.update(weights.transformer, strict=False) - text_encoder.update(weights.text_encoder, strict=False) diff --git a/src/mflux/models/fibo_vlm/cli/__init__.py b/src/mflux/models/fibo_vlm/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/inspire_fibo.py b/src/mflux/models/fibo_vlm/cli/fibo_inspire.py similarity index 61% rename from src/mflux/inspire_fibo.py rename to src/mflux/models/fibo_vlm/cli/fibo_inspire.py index a3ae008..f2bfbd9 100644 --- a/src/mflux/inspire_fibo.py +++ b/src/mflux/models/fibo_vlm/cli/fibo_inspire.py @@ -1,10 +1,8 @@ -import json from pathlib import Path -from PIL import Image - +from mflux.cli.parser.parsers import CommandLineParser from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM -from mflux.ui.cli.parsers import CommandLineParser +from mflux.models.fibo_vlm.model.util import FiboVLMUtil from mflux.utils.exceptions import PromptFileReadError @@ -15,7 +13,7 @@ def main(): parser.add_argument("--image-path", type=Path, required=True, help="Path to image file to inspire from") parser.add_argument("--prompt", type=str, default=None, help="Optional text prompt to blend with the image (e.g., 'Make futuristic')") parser.add_argument("--output", type=Path, default=Path("inspired.json"), help="Output path for generated JSON prompt (default: inspired.json)") - parser.add_argument("--path", type=str, default=None, help="Local path for loading the VLM model from disk") + parser.add_argument("-q", "--quantize", type=int, choices=[3, 4, 5, 6, 8], default=None, help="Quantize the VLM model (3, 4, 5, 6, or 8 bits)") parser.add_argument("--top-p", type=float, default=0.9, help="Top-p sampling for VLM (default: 0.9)") parser.add_argument("--temperature", type=float, default=0.2, help="Temperature for VLM (default: 0.2)") parser.add_argument("--max-tokens", type=int, default=4096, help="Max tokens for VLM generation (default: 4096)") @@ -25,44 +23,24 @@ def main(): try: # 1. Load the image - image = _load_image(args.image_path) + image = FiboVLMUtil.load_image(args.image_path) # 2. Generate JSON prompt from image - vlm = FiboVLM(local_path=args.path) + vlm = FiboVLM(quantize=args.quantize) inspired_json = vlm.inspire( + seed=args.seed, image=image, prompt=args.prompt, top_p=args.top_p, temperature=args.temperature, max_tokens=args.max_tokens, - seed=args.seed, ) # 3. Parse and save generated JSON prompt - _save_prompt(args, inspired_json) + FiboVLMUtil.save_json_prompt(args.output, inspired_json) except (PromptFileReadError, ValueError) as exc: print(exc) -def _save_prompt(args, inspired_json): - try: - inspired_json_parsed = json.loads(inspired_json) - except json.JSONDecodeError as e: - raise ValueError(f"VLM did not return valid JSON: {e}") - args.output.parent.mkdir(parents=True, exist_ok=True) - with open(args.output, "w") as f: - json.dump(inspired_json_parsed, f, indent=2, ensure_ascii=False) - - -def _load_image(image_path: Path): - if not image_path.exists(): - raise PromptFileReadError(f"Image file does not exist: {image_path}") - try: - image = Image.open(image_path).convert("RGB") - except (OSError, IOError, ValueError) as e: - raise PromptFileReadError(f"Failed to load image: {e}") - return image - - if __name__ == "__main__": main() diff --git a/src/mflux/refine_fibo.py b/src/mflux/models/fibo_vlm/cli/fibo_refine.py similarity index 56% rename from src/mflux/refine_fibo.py rename to src/mflux/models/fibo_vlm/cli/fibo_refine.py index eb4b689..ce182a0 100644 --- a/src/mflux/refine_fibo.py +++ b/src/mflux/models/fibo_vlm/cli/fibo_refine.py @@ -1,8 +1,8 @@ -import json from pathlib import Path +from mflux.cli.parser.parsers import CommandLineParser from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM -from mflux.ui.cli.parsers import CommandLineParser +from mflux.models.fibo_vlm.model.util import FiboVLMUtil from mflux.utils.exceptions import PromptFileReadError @@ -13,7 +13,7 @@ def main(): parser.add_argument("--prompt-file", type=Path, required=True, help="Path to JSON prompt file to refine") parser.add_argument("--instructions", type=str, required=True, help="Text instructions for how to refine the prompt (e.g., 'make the dragon blue')") parser.add_argument("--output", type=Path, default=Path("refined.json"), help="Output path for refined JSON prompt (default: refined.json)") - parser.add_argument("--path", type=str, default=None, help="Local path for loading the VLM model from disk") + parser.add_argument("-q", "--quantize", type=int, choices=[3, 4, 5, 6, 8], default=None, help="Quantize the VLM model (3, 4, 5, 6, or 8 bits)") parser.add_argument("--top-p", type=float, default=0.9, help="Top-p sampling for VLM (default: 0.9)") parser.add_argument("--temperature", type=float, default=0.2, help="Temperature for VLM (default: 0.2)") parser.add_argument("--max-tokens", type=int, default=4096, help="Max tokens for VLM generation (default: 4096)") @@ -23,46 +23,21 @@ def main(): try: # 1. Refine the JSON prompt - vlm = FiboVLM(local_path=args.path) - + vlm = FiboVLM(quantize=args.quantize) refined_json = vlm.refine( - structured_prompt=_get_structured_prompt(args), + seed=args.seed, + structured_prompt=FiboVLMUtil.get_structured_prompt(args.prompt_file), editing_instructions=args.instructions, top_p=args.top_p, temperature=args.temperature, max_tokens=args.max_tokens, - seed=args.seed, ) # 2. Parse and save refined JSON prompt - _save_prompt(args, refined_json) + FiboVLMUtil.save_json_prompt(args.output, refined_json) except (PromptFileReadError, ValueError) as exc: print(exc) -def _save_prompt(args, refined_json): - try: - refined_json_parsed = json.loads(refined_json) - except json.JSONDecodeError as e: - raise ValueError(f"VLM did not return valid JSON: {e}") - args.output.parent.mkdir(parents=True, exist_ok=True) - with open(args.output, "w") as f: - json.dump(refined_json_parsed, f, indent=2, ensure_ascii=False) - - -def _get_structured_prompt(args): - if not args.prompt_file.exists(): - raise PromptFileReadError(f"Prompt file does not exist: {args.prompt_file}") - with open(args.prompt_file, "rt") as f: - structured_prompt = f.read().strip() - if not structured_prompt: - raise PromptFileReadError(f"Prompt file is empty: {args.prompt_file}") - try: - json.loads(structured_prompt) - except json.JSONDecodeError as e: - raise PromptFileReadError(f"Prompt file does not contain valid JSON: {e}") - return structured_prompt - - if __name__ == "__main__": main() diff --git a/src/mflux/models/fibo_vlm/fibo_vlm_initializer.py b/src/mflux/models/fibo_vlm/fibo_vlm_initializer.py index 9643307..54f8282 100644 --- a/src/mflux/models/fibo_vlm/fibo_vlm_initializer.py +++ b/src/mflux/models/fibo_vlm/fibo_vlm_initializer.py @@ -1,112 +1,55 @@ -import os -from pathlib import Path - -from transformers import Qwen2Tokenizer - -from mflux.models.fibo.tokenizer.qwen2vl_processor import Qwen2VLProcessor +from mflux.models.common.tokenizer import TokenizerLoader +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_loader import WeightLoader from mflux.models.fibo_vlm.model.qwen3_vl_decoder import Qwen3VLDecoder from mflux.models.fibo_vlm.model.qwen3_vl_vision_model import Qwen3VLVisionModel -from mflux.models.fibo_vlm.weights.fibo_vlm_weight_handler import FIBOVLMWeightHandler -from mflux.utils.download import snapshot_download +from mflux.models.fibo_vlm.weights.fibo_vlm_weight_definition import FIBOVLMWeightDefinition -class FIBOVLMInitializer: +class FiboVLMInitializer: @staticmethod def init( - vlm_model, - model_id: str = "briaai/FIBO-vlm", - local_path: str | None = None, + model, + model_path: str = "briaai/FIBO-vlm", + quantize: int | None = None, ) -> None: - # 1. Load VLM weights - weights = FIBOVLMWeightHandler.load_vlm_regular_weights( - repo_id=model_id, - local_path=local_path, + FiboVLMInitializer._init_config(model, model_path) + weights = FiboVLMInitializer._load_weights(model_path) + FiboVLMInitializer._init_tokenizers(model, model_path) + FiboVLMInitializer._init_models(model) + FiboVLMInitializer._apply_weights(model, weights, quantize) + + @staticmethod + def _init_config(model, model_path: str) -> None: + model.model_path = model_path + + @staticmethod + def _load_weights(model_path: str) -> LoadedWeights: + return WeightLoader.load( + weight_definition=FIBOVLMWeightDefinition, + model_path=model_path, ) - # 2. Initialize processor for tokenization - tokenizer = FIBOVLMInitializer._get_tokenizer(local_path, model_id) - vlm_model.processor = Qwen2VLProcessor(tokenizer=tokenizer) - - # 3. Initialize all models - vlm_model.decoder = Qwen3VLDecoder(visual=Qwen3VLVisionModel()) - - # 4. Apply weights to decoder and visual encoder - vlm_model.decoder.update(weights.decoder, strict=False) - vlm_model.decoder.visual.update(weights.visual, strict=False) - - # Store model ID and local path - vlm_model.model_id = model_id - vlm_model.local_path = local_path + @staticmethod + def _init_tokenizers(model, model_path: str) -> None: + model.tokenizers = TokenizerLoader.load_all( + definitions=FIBOVLMWeightDefinition.get_tokenizers(), + model_path=model_path, + ) @staticmethod - def _get_tokenizer(local_path, model_id): - # Get the root path - if local_path: - root_path = Path(local_path) - else: - root_path = FIBOVLMInitializer._get_model_path(model_id) - - # Try different possible tokenizer paths (like FiboTokenizerHandler does) - tokenizer_path = root_path / "tokenizer" - if not tokenizer_path.exists(): - tokenizer_path = root_path / "text_encoder" - if not tokenizer_path.exists(): - # Tokenizer files are in the root - this is the case for FIBO-vlm - tokenizer_path = root_path - - # Set HF_HUB_OFFLINE to force offline mode - old_offline = os.environ.get("HF_HUB_OFFLINE") - try: - os.environ["HF_HUB_OFFLINE"] = "1" - # Use Qwen2Tokenizer directly instead of AutoTokenizer to avoid config confusion - tokenizer = Qwen2Tokenizer.from_pretrained( - pretrained_model_name_or_path=str(tokenizer_path), - local_files_only=True, - ) - finally: - # Restore original value - if old_offline is None: - os.environ.pop("HF_HUB_OFFLINE", None) - else: - os.environ["HF_HUB_OFFLINE"] = old_offline - - return tokenizer + def _init_models(model) -> None: + model.decoder = Qwen3VLDecoder(visual=Qwen3VLVisionModel()) @staticmethod - def _get_model_path(model_id: str) -> Path: - # Try to use cached path first - try: - root_path = Path( - snapshot_download( - repo_id=model_id, - local_files_only=True, - ) - ) - # Check if tokenizer files actually exist - the model weights might be cached - # but tokenizer files may not have been downloaded yet - if not FIBOVLMInitializer._tokenizer_files_exist(root_path): - # Tokenizer files missing, need to download them - root_path = Path( - snapshot_download( - repo_id=model_id, - local_files_only=False, - ) - ) - except (FileNotFoundError, OSError): - # Model not in cache, allow download if online - root_path = Path( - snapshot_download( - repo_id=model_id, - ) - ) - return root_path - - @staticmethod - def _tokenizer_files_exist(root_path: Path) -> bool: - # Check for vocab.json in possible tokenizer locations - possible_paths = [ - root_path / "tokenizer" / "vocab.json", - root_path / "text_encoder" / "vocab.json", - root_path / "vocab.json", - ] - return any(p.exists() for p in possible_paths) + def _apply_weights(model, weights: LoadedWeights, quantize: int | None) -> None: + model.bits = WeightApplier.apply_and_quantize( + weights=weights, + quantize_arg=quantize, + weight_definition=FIBOVLMWeightDefinition, + models={ + "decoder": model.decoder, + "visual": model.decoder.visual, + }, + ) diff --git a/src/mflux/models/fibo_vlm/model/fibo_vlm.py b/src/mflux/models/fibo_vlm/model/fibo_vlm.py index 13063d0..cc2d3eb 100644 --- a/src/mflux/models/fibo_vlm/model/fibo_vlm.py +++ b/src/mflux/models/fibo_vlm/model/fibo_vlm.py @@ -6,27 +6,31 @@ import mlx.core as mx import numpy as np from PIL import Image -from mflux.models.fibo.tokenizer.qwen2vl_processor import Qwen2VLProcessor -from mflux.models.fibo_vlm.fibo_vlm_initializer import FIBOVLMInitializer +from mflux.models.fibo_vlm.fibo_vlm_initializer import FiboVLMInitializer from mflux.models.fibo_vlm.model.qwen3_vl_decoder import Qwen3VLDecoder from mflux.models.fibo_vlm.model.qwen3_vl_util import Qwen3VLUtil +from mflux.models.fibo_vlm.tokenizer.qwen2vl_processor import Qwen2VLProcessor class FiboVLM: decoder = Qwen3VLDecoder - processor: Qwen2VLProcessor def __init__( self, model_id: str = "briaai/FIBO-vlm", - local_path: str | None = None, + model_path: str | None = None, + quantize: int | None = None, ): - FIBOVLMInitializer.init( - vlm_model=self, - model_id=model_id, - local_path=local_path, + FiboVLMInitializer.init( + model=self, + model_path=model_path if model_path else model_id, + quantize=quantize, ) + @property + def processor(self) -> Qwen2VLProcessor: + return self.tokenizers["fibo_vlm"].processor + def generate( self, prompt: str, @@ -107,8 +111,8 @@ class FiboVLM: messages = FiboVLM._build_messages( task=task, image=image, - refine_image=refine_image, prompt=prompt, + refine_image=refine_image, structured_prompt=FiboVLM._normalize_json(structured_prompt), editing_instructions=editing_instructions, ) @@ -119,17 +123,17 @@ class FiboVLM: image_grid_thw = mx.array(formatted["image_grid_thw"]) if formatted.get("image_grid_thw") is not None else None stop_token_sequences = [self.processor.tokenizer.encode(s, add_special_tokens=False) for s in stop] generated_ids = Qwen3VLUtil.generate_text( + seed=seed, + top_p=top_p, decoder=self.decoder, input_ids=input_ids, attention_mask=attention_mask, pixel_values=pixel_values, image_grid_thw=image_grid_thw, max_new_tokens=max_tokens, - top_p=top_p, temperature=temperature, stop_token_sequences=stop_token_sequences, eos_token_id=self.processor.tokenizer.eos_token_id, - seed=seed, ) generated_tokens = generated_ids[:, input_ids.shape[1] :] generated_text = self.processor.tokenizer.decode(np.array(generated_tokens[0]), skip_special_tokens=True) diff --git a/src/mflux/models/fibo_vlm/model/util.py b/src/mflux/models/fibo_vlm/model/util.py new file mode 100644 index 0000000..e8382c2 --- /dev/null +++ b/src/mflux/models/fibo_vlm/model/util.py @@ -0,0 +1,42 @@ +import json +from pathlib import Path + +from PIL import Image + +from mflux.utils.exceptions import PromptFileReadError + + +class FiboVLMUtil: + @staticmethod + def save_json_prompt(output_path: Path, json_string: str) -> None: + try: + json_parsed = json.loads(json_string) + except json.JSONDecodeError as e: + raise ValueError(f"VLM did not return valid JSON: {e}") + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as f: + json.dump(json_parsed, f, indent=2, ensure_ascii=False) + + @staticmethod + def get_structured_prompt(prompt_file: Path) -> str: + if not prompt_file.exists(): + raise PromptFileReadError(f"Prompt file does not exist: {prompt_file}") + with open(prompt_file, "rt") as f: + structured_prompt = f.read().strip() + if not structured_prompt: + raise PromptFileReadError(f"Prompt file is empty: {prompt_file}") + try: + json.loads(structured_prompt) + except json.JSONDecodeError as e: + raise PromptFileReadError(f"Prompt file does not contain valid JSON: {e}") + return structured_prompt + + @staticmethod + def load_image(image_path: Path) -> Image.Image: + if not image_path.exists(): + raise PromptFileReadError(f"Image file does not exist: {image_path}") + try: + image = Image.open(image_path).convert("RGB") + except (OSError, IOError, ValueError) as e: + raise PromptFileReadError(f"Failed to load image: {e}") + return image diff --git a/src/mflux/models/fibo_vlm/tokenizer/__init__.py b/src/mflux/models/fibo_vlm/tokenizer/__init__.py new file mode 100644 index 0000000..269c166 --- /dev/null +++ b/src/mflux/models/fibo_vlm/tokenizer/__init__.py @@ -0,0 +1 @@ +# FIBO Tokenizer - uses unified tokenizer system from mflux.models.common.tokenizer diff --git a/src/mflux/models/fibo/tokenizer/qwen2vl_image_processor.py b/src/mflux/models/fibo_vlm/tokenizer/qwen2vl_image_processor.py similarity index 100% rename from src/mflux/models/fibo/tokenizer/qwen2vl_image_processor.py rename to src/mflux/models/fibo_vlm/tokenizer/qwen2vl_image_processor.py diff --git a/src/mflux/models/fibo/tokenizer/qwen2vl_processor.py b/src/mflux/models/fibo_vlm/tokenizer/qwen2vl_processor.py similarity index 99% rename from src/mflux/models/fibo/tokenizer/qwen2vl_processor.py rename to src/mflux/models/fibo_vlm/tokenizer/qwen2vl_processor.py index aa03e6a..37a23da 100644 --- a/src/mflux/models/fibo/tokenizer/qwen2vl_processor.py +++ b/src/mflux/models/fibo_vlm/tokenizer/qwen2vl_processor.py @@ -3,7 +3,7 @@ from typing import Optional, Union import numpy as np from PIL import Image -from mflux.models.fibo.tokenizer.qwen2vl_image_processor import Qwen2VLImageProcessor +from mflux.models.fibo_vlm.tokenizer.qwen2vl_image_processor import Qwen2VLImageProcessor class Qwen2VLProcessor: diff --git a/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_definition.py b/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_definition.py new file mode 100644 index 0000000..8ac23c9 --- /dev/null +++ b/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_definition.py @@ -0,0 +1,62 @@ +from typing import List + +from mflux.models.common.tokenizer import VisionLanguageTokenizer +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition, TokenizerDefinition +from mflux.models.fibo_vlm.tokenizer.qwen2vl_processor import Qwen2VLProcessor +from mflux.models.fibo_vlm.weights.fibo_vlm_weight_mapping import FIBOVLMWeightMapping + + +class FIBOVLMWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="decoder", + hf_subdir="", + num_blocks=36, + loading_mode="torch_bfloat16", + weight_prefix_filters=["model.language_model", "lm_head"], + mapping_getter=lambda: FIBOVLMWeightMapping.get_vlm_decoder_mapping(36), + ), + ComponentDefinition( + name="visual", + hf_subdir="", + num_blocks=24, + loading_mode="torch_bfloat16", + weight_prefix_filters=["model.visual"], + mapping_getter=lambda: FIBOVLMWeightMapping.get_vlm_visual_mapping(24), + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [ + TokenizerDefinition( + name="fibo_vlm", + hf_subdir=".", # Root directory + tokenizer_class="Qwen2Tokenizer", + encoder_class=VisionLanguageTokenizer, + processor_class=Qwen2VLProcessor, + max_length=1024, + fallback_subdirs=["tokenizer", "text_encoder"], + download_patterns=["vocab.json", "merges.txt", "tokenizer.json", "tokenizer_config.json"], + ), + ] + + @staticmethod + def get_download_patterns() -> List[str]: + return [ + "*.safetensors", + "*.json", + "vocab.json", + "merges.txt", + "tokenizer.json", + "tokenizer_config.json", + ] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + # Skip lm_head - it maps to vocab size and quantization corrupts dimensions + if "lm_head" in path: + return False + return hasattr(module, "to_quantized") diff --git a/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_handler.py b/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_handler.py deleted file mode 100644 index 66531f8..0000000 --- a/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_handler.py +++ /dev/null @@ -1,74 +0,0 @@ -import mlx.core as mx -import torch -from transformers import Qwen3VLForConditionalGeneration - -from mflux.models.common.weights.mapping.weight_mapper import WeightMapper -from mflux.models.fibo.weights import FIBOWeightHandler -from mflux.models.fibo_vlm.weights.fibo_vlm_weight_mapping import FIBOVLMWeightMapping -from mflux.models.flux.weights.weight_handler import MetaData - - -class FIBOVLMWeightHandler: - @staticmethod - def load_vlm_regular_weights( - repo_id: str = "briaai/FIBO-vlm", - local_path: str | None = None, - ) -> "FIBOWeightHandler": - # Load model - try offline first (use cache), fall back to online if needed - pretrained_path = local_path or repo_id - if local_path: - # If explicit local path, use local_files_only - model = Qwen3VLForConditionalGeneration.from_pretrained( - pretrained_model_name_or_path=pretrained_path, - dtype=torch.bfloat16, - local_files_only=True, - ) - else: - # Try offline first (use cache), fall back to online if not cached - try: - model = Qwen3VLForConditionalGeneration.from_pretrained( - pretrained_model_name_or_path=pretrained_path, - dtype=torch.bfloat16, - local_files_only=True, - ) - except (FileNotFoundError, OSError): - # Model not in cache, allow download if online - model = Qwen3VLForConditionalGeneration.from_pretrained( - pretrained_model_name_or_path=pretrained_path, - dtype=torch.bfloat16, - local_files_only=False, - ) - num_layers = model.config.text_config.num_hidden_layers - depth = model.config.vision_config.depth - state_dict = model.state_dict() - - raw_decoder_weights = { - k: FIBOVLMWeightHandler._to_mlx(v) - for k, v in state_dict.items() - if k.startswith(("model.language_model", "lm_head")) - } - - raw_visual_weights = { - k: FIBOVLMWeightHandler._to_mlx(v) for k, v in state_dict.items() if k.startswith("model.visual") - } - - decoder_weights = WeightMapper.apply_mapping( - hf_weights=raw_decoder_weights, - mapping=FIBOVLMWeightMapping.get_vlm_decoder_mapping(num_layers=num_layers), - num_blocks=num_layers, - ) - visual_weights = WeightMapper.apply_mapping( - hf_weights=raw_visual_weights, - mapping=FIBOVLMWeightMapping.get_vlm_visual_mapping(depth=depth), - num_blocks=depth, - ) - - return FIBOWeightHandler( - decoder=decoder_weights, - visual=visual_weights, - meta_data=MetaData(quantization_level=None, scale=None, is_lora=False, mflux_version=None), - ) - - @staticmethod - def _to_mlx(tensor: torch.Tensor) -> mx.array: - return mx.array((tensor.to(torch.float16) if tensor.dtype == torch.bfloat16 else tensor).detach().cpu().numpy()) diff --git a/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_mapping.py b/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_mapping.py index debf71d..628a199 100644 --- a/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_mapping.py +++ b/src/mflux/models/fibo_vlm/weights/fibo_vlm_weight_mapping.py @@ -1,7 +1,7 @@ from typing import List from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget -from mflux.models.fibo.weights.fibo_weight_mapping import transpose_conv3d_weight +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms class FIBOVLMWeightMapping(WeightMapping): @@ -9,214 +9,210 @@ class FIBOVLMWeightMapping(WeightMapping): def get_vlm_decoder_mapping(num_layers: int = 36) -> List[WeightTarget]: return [ WeightTarget( - mlx_path="embed_tokens.weight", - hf_patterns=["model.language_model.embed_tokens.weight"], + to_pattern="embed_tokens.weight", + from_pattern=["model.language_model.embed_tokens.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.q_proj.weight", - hf_patterns=["model.language_model.layers.{block}.self_attn.q_proj.weight"], + to_pattern="layers.{block}.self_attn.q_proj.weight", + from_pattern=["model.language_model.layers.{block}.self_attn.q_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.q_proj.bias", - hf_patterns=["model.language_model.layers.{block}.self_attn.q_proj.bias"], + to_pattern="layers.{block}.self_attn.q_proj.bias", + from_pattern=["model.language_model.layers.{block}.self_attn.q_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.k_proj.weight", - hf_patterns=["model.language_model.layers.{block}.self_attn.k_proj.weight"], + to_pattern="layers.{block}.self_attn.k_proj.weight", + from_pattern=["model.language_model.layers.{block}.self_attn.k_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.k_proj.bias", - hf_patterns=["model.language_model.layers.{block}.self_attn.k_proj.bias"], + to_pattern="layers.{block}.self_attn.k_proj.bias", + from_pattern=["model.language_model.layers.{block}.self_attn.k_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.v_proj.weight", - hf_patterns=["model.language_model.layers.{block}.self_attn.v_proj.weight"], + to_pattern="layers.{block}.self_attn.v_proj.weight", + from_pattern=["model.language_model.layers.{block}.self_attn.v_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.v_proj.bias", - hf_patterns=["model.language_model.layers.{block}.self_attn.v_proj.bias"], + to_pattern="layers.{block}.self_attn.v_proj.bias", + from_pattern=["model.language_model.layers.{block}.self_attn.v_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.o_proj.weight", - hf_patterns=["model.language_model.layers.{block}.self_attn.o_proj.weight"], + to_pattern="layers.{block}.self_attn.o_proj.weight", + from_pattern=["model.language_model.layers.{block}.self_attn.o_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.o_proj.bias", - hf_patterns=["model.language_model.layers.{block}.self_attn.o_proj.bias"], + to_pattern="layers.{block}.self_attn.o_proj.bias", + from_pattern=["model.language_model.layers.{block}.self_attn.o_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.q_norm.weight", - hf_patterns=["model.language_model.layers.{block}.self_attn.q_norm.weight"], + to_pattern="layers.{block}.self_attn.q_norm.weight", + from_pattern=["model.language_model.layers.{block}.self_attn.q_norm.weight"], ), WeightTarget( - mlx_path="layers.{block}.self_attn.k_norm.weight", - hf_patterns=["model.language_model.layers.{block}.self_attn.k_norm.weight"], + to_pattern="layers.{block}.self_attn.k_norm.weight", + from_pattern=["model.language_model.layers.{block}.self_attn.k_norm.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.gate_proj.weight", - hf_patterns=["model.language_model.layers.{block}.mlp.gate_proj.weight"], + to_pattern="layers.{block}.mlp.gate_proj.weight", + from_pattern=["model.language_model.layers.{block}.mlp.gate_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.gate_proj.bias", - hf_patterns=["model.language_model.layers.{block}.mlp.gate_proj.bias"], + to_pattern="layers.{block}.mlp.gate_proj.bias", + from_pattern=["model.language_model.layers.{block}.mlp.gate_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.mlp.up_proj.weight", - hf_patterns=["model.language_model.layers.{block}.mlp.up_proj.weight"], + to_pattern="layers.{block}.mlp.up_proj.weight", + from_pattern=["model.language_model.layers.{block}.mlp.up_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.up_proj.bias", - hf_patterns=["model.language_model.layers.{block}.mlp.up_proj.bias"], + to_pattern="layers.{block}.mlp.up_proj.bias", + from_pattern=["model.language_model.layers.{block}.mlp.up_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.mlp.down_proj.weight", - hf_patterns=["model.language_model.layers.{block}.mlp.down_proj.weight"], + to_pattern="layers.{block}.mlp.down_proj.weight", + from_pattern=["model.language_model.layers.{block}.mlp.down_proj.weight"], ), WeightTarget( - mlx_path="layers.{block}.mlp.down_proj.bias", - hf_patterns=["model.language_model.layers.{block}.mlp.down_proj.bias"], + to_pattern="layers.{block}.mlp.down_proj.bias", + from_pattern=["model.language_model.layers.{block}.mlp.down_proj.bias"], ), WeightTarget( - mlx_path="layers.{block}.input_layernorm.weight", - hf_patterns=["model.language_model.layers.{block}.input_layernorm.weight"], + to_pattern="layers.{block}.input_layernorm.weight", + from_pattern=["model.language_model.layers.{block}.input_layernorm.weight"], ), WeightTarget( - mlx_path="layers.{block}.post_attention_layernorm.weight", - hf_patterns=["model.language_model.layers.{block}.post_attention_layernorm.weight"], + to_pattern="layers.{block}.post_attention_layernorm.weight", + from_pattern=["model.language_model.layers.{block}.post_attention_layernorm.weight"], ), WeightTarget( - mlx_path="norm.weight", - hf_patterns=["model.language_model.norm.weight"], + to_pattern="norm.weight", + from_pattern=["model.language_model.norm.weight"], ), WeightTarget( - mlx_path="lm_head.weight", - hf_patterns=["lm_head.weight"], + to_pattern="lm_head.weight", + from_pattern=["model.language_model.embed_tokens.weight"], ), ] @staticmethod def get_vlm_visual_mapping(depth: int = 24) -> List[WeightTarget]: return [ - # Patch embedding WeightTarget( - mlx_path="patch_embed.proj.weight", - hf_patterns=["model.visual.patch_embed.proj.weight"], - transform=transpose_conv3d_weight, + to_pattern="patch_embed.proj.weight", + from_pattern=["model.visual.patch_embed.proj.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="patch_embed.proj.bias", - hf_patterns=["model.visual.patch_embed.proj.bias"], - ), - # Position embeddings - WeightTarget( - mlx_path="pos_embed.weight", - hf_patterns=["model.visual.pos_embed.weight"], - ), - # Vision transformer blocks - WeightTarget( - mlx_path="blocks.{block}.norm1.weight", - hf_patterns=["model.visual.blocks.{block}.norm1.weight"], + to_pattern="patch_embed.proj.bias", + from_pattern=["model.visual.patch_embed.proj.bias"], ), WeightTarget( - mlx_path="blocks.{block}.norm1.bias", - hf_patterns=["model.visual.blocks.{block}.norm1.bias"], + to_pattern="pos_embed.weight", + from_pattern=["model.visual.pos_embed.weight"], ), WeightTarget( - mlx_path="blocks.{block}.norm2.weight", - hf_patterns=["model.visual.blocks.{block}.norm2.weight"], + to_pattern="blocks.{block}.norm1.weight", + from_pattern=["model.visual.blocks.{block}.norm1.weight"], ), WeightTarget( - mlx_path="blocks.{block}.norm2.bias", - hf_patterns=["model.visual.blocks.{block}.norm2.bias"], + to_pattern="blocks.{block}.norm1.bias", + from_pattern=["model.visual.blocks.{block}.norm1.bias"], + ), + WeightTarget( + to_pattern="blocks.{block}.norm2.weight", + from_pattern=["model.visual.blocks.{block}.norm2.weight"], + ), + WeightTarget( + to_pattern="blocks.{block}.norm2.bias", + from_pattern=["model.visual.blocks.{block}.norm2.bias"], ), # Attention WeightTarget( - mlx_path="blocks.{block}.attn.qkv.weight", - hf_patterns=["model.visual.blocks.{block}.attn.qkv.weight"], + to_pattern="blocks.{block}.attn.qkv.weight", + from_pattern=["model.visual.blocks.{block}.attn.qkv.weight"], ), WeightTarget( - mlx_path="blocks.{block}.attn.qkv.bias", - hf_patterns=["model.visual.blocks.{block}.attn.qkv.bias"], + to_pattern="blocks.{block}.attn.qkv.bias", + from_pattern=["model.visual.blocks.{block}.attn.qkv.bias"], ), WeightTarget( - mlx_path="blocks.{block}.attn.proj.weight", - hf_patterns=["model.visual.blocks.{block}.attn.proj.weight"], + to_pattern="blocks.{block}.attn.proj.weight", + from_pattern=["model.visual.blocks.{block}.attn.proj.weight"], ), WeightTarget( - mlx_path="blocks.{block}.attn.proj.bias", - hf_patterns=["model.visual.blocks.{block}.attn.proj.bias"], + to_pattern="blocks.{block}.attn.proj.bias", + from_pattern=["model.visual.blocks.{block}.attn.proj.bias"], ), # MLP WeightTarget( - mlx_path="blocks.{block}.mlp.linear_fc1.weight", - hf_patterns=["model.visual.blocks.{block}.mlp.linear_fc1.weight"], + to_pattern="blocks.{block}.mlp.linear_fc1.weight", + from_pattern=["model.visual.blocks.{block}.mlp.linear_fc1.weight"], ), WeightTarget( - mlx_path="blocks.{block}.mlp.linear_fc1.bias", - hf_patterns=["model.visual.blocks.{block}.mlp.linear_fc1.bias"], + to_pattern="blocks.{block}.mlp.linear_fc1.bias", + from_pattern=["model.visual.blocks.{block}.mlp.linear_fc1.bias"], ), WeightTarget( - mlx_path="blocks.{block}.mlp.linear_fc2.weight", - hf_patterns=["model.visual.blocks.{block}.mlp.linear_fc2.weight"], + to_pattern="blocks.{block}.mlp.linear_fc2.weight", + from_pattern=["model.visual.blocks.{block}.mlp.linear_fc2.weight"], ), WeightTarget( - mlx_path="blocks.{block}.mlp.linear_fc2.bias", - hf_patterns=["model.visual.blocks.{block}.mlp.linear_fc2.bias"], + to_pattern="blocks.{block}.mlp.linear_fc2.bias", + from_pattern=["model.visual.blocks.{block}.mlp.linear_fc2.bias"], ), # Final merger WeightTarget( - mlx_path="merger.norm.weight", - hf_patterns=["model.visual.merger.norm.weight"], + to_pattern="merger.norm.weight", + from_pattern=["model.visual.merger.norm.weight"], ), WeightTarget( - mlx_path="merger.norm.bias", - hf_patterns=["model.visual.merger.norm.bias"], + to_pattern="merger.norm.bias", + from_pattern=["model.visual.merger.norm.bias"], ), WeightTarget( - mlx_path="merger.linear_fc1.weight", - hf_patterns=["model.visual.merger.linear_fc1.weight"], + to_pattern="merger.linear_fc1.weight", + from_pattern=["model.visual.merger.linear_fc1.weight"], ), WeightTarget( - mlx_path="merger.linear_fc1.bias", - hf_patterns=["model.visual.merger.linear_fc1.bias"], + to_pattern="merger.linear_fc1.bias", + from_pattern=["model.visual.merger.linear_fc1.bias"], ), WeightTarget( - mlx_path="merger.linear_fc2.weight", - hf_patterns=["model.visual.merger.linear_fc2.weight"], + to_pattern="merger.linear_fc2.weight", + from_pattern=["model.visual.merger.linear_fc2.weight"], ), WeightTarget( - mlx_path="merger.linear_fc2.bias", - hf_patterns=["model.visual.merger.linear_fc2.bias"], - ), - # DeepStack mergers (fixed 3 items, not depth-based) - WeightTarget( - mlx_path="deepstack_merger_list.{block}.norm.weight", - hf_patterns=["model.visual.deepstack_merger_list.{block}.norm.weight"], - max_blocks=3, # Fixed 3 items (0, 1, 2) + to_pattern="merger.linear_fc2.bias", + from_pattern=["model.visual.merger.linear_fc2.bias"], ), WeightTarget( - mlx_path="deepstack_merger_list.{block}.norm.bias", - hf_patterns=["model.visual.deepstack_merger_list.{block}.norm.bias"], + to_pattern="deepstack_merger_list.{block}.norm.weight", + from_pattern=["model.visual.deepstack_merger_list.{block}.norm.weight"], max_blocks=3, ), WeightTarget( - mlx_path="deepstack_merger_list.{block}.linear_fc1.weight", - hf_patterns=["model.visual.deepstack_merger_list.{block}.linear_fc1.weight"], + to_pattern="deepstack_merger_list.{block}.norm.bias", + from_pattern=["model.visual.deepstack_merger_list.{block}.norm.bias"], max_blocks=3, ), WeightTarget( - mlx_path="deepstack_merger_list.{block}.linear_fc1.bias", - hf_patterns=["model.visual.deepstack_merger_list.{block}.linear_fc1.bias"], + to_pattern="deepstack_merger_list.{block}.linear_fc1.weight", + from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc1.weight"], max_blocks=3, ), WeightTarget( - mlx_path="deepstack_merger_list.{block}.linear_fc2.weight", - hf_patterns=["model.visual.deepstack_merger_list.{block}.linear_fc2.weight"], + to_pattern="deepstack_merger_list.{block}.linear_fc1.bias", + from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc1.bias"], max_blocks=3, ), WeightTarget( - mlx_path="deepstack_merger_list.{block}.linear_fc2.bias", - hf_patterns=["model.visual.deepstack_merger_list.{block}.linear_fc2.bias"], + to_pattern="deepstack_merger_list.{block}.linear_fc2.weight", + from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc2.weight"], + max_blocks=3, + ), + WeightTarget( + to_pattern="deepstack_merger_list.{block}.linear_fc2.bias", + from_pattern=["model.visual.deepstack_merger_list.{block}.linear_fc2.bias"], max_blocks=3, ), ] diff --git a/src/mflux/models/flux/cli/__init__.py b/src/mflux/models/flux/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/concept.py b/src/mflux/models/flux/cli/flux_concept.py similarity index 71% rename from src/mflux/concept.py rename to src/mflux/models/flux/cli/flux_concept.py index 55c4863..1dbeba6 100644 --- a/src/mflux/concept.py +++ b/src/mflux/models/flux/cli/flux_concept.py @@ -1,11 +1,11 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.concept_attention.flux_concept import Flux1Concept -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -28,31 +28,33 @@ def main(): flux = Flux1Concept( model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, concept=args.concept, + guidance=args.guidance, + image_path=args.image_path, + num_inference_steps=args.steps, + image_strength=args.image_strength, heatmap_timesteps=args.heatmap_timesteps, heatmap_layer_indices=args.heatmap_layer_indices, - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - image_strength=args.image_strength, - ), ) # 4. Save the image and heatmap image.save_with_heatmap(path=args.output.format(seed=seed), export_json_metadata=args.metadata) diff --git a/src/mflux/concept_from_image.py b/src/mflux/models/flux/cli/flux_concept_from_image.py similarity index 73% rename from src/mflux/concept_from_image.py rename to src/mflux/models/flux/cli/flux_concept_from_image.py index c45ec1f..2bbcac4 100644 --- a/src/mflux/concept_from_image.py +++ b/src/mflux/models/flux/cli/flux_concept_from_image.py @@ -1,11 +1,11 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.concept_attention.flux_concept_from_image import Flux1ConceptFromImage -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -28,32 +28,32 @@ def main(): flux = Flux1ConceptFromImage( model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, concept=args.concept, + guidance=args.guidance, + num_inference_steps=args.steps, image_path=str(args.input_image_path), heatmap_timesteps=args.heatmap_timesteps, heatmap_layer_indices=args.heatmap_layer_indices, - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - image_strength=args.image_strength, - ), ) # 4. Save the image and heatmap image.save_with_heatmap(path=args.output.format(seed=seed), export_json_metadata=args.metadata) diff --git a/src/mflux/generate.py b/src/mflux/models/flux/cli/flux_generate.py similarity index 53% rename from src/mflux/generate.py rename to src/mflux/models/flux/cli/flux_generate.py index 820a10a..35349cc 100644 --- a/src/mflux/generate.py +++ b/src/mflux/models/flux/cli/flux_generate.py @@ -1,11 +1,12 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.txt2img.flux import Flux1 -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.utils.dimension_resolver import DimensionResolver from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -14,7 +15,7 @@ def main(): parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) parser.add_lora_arguments() - parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_generator_arguments(supports_metadata_config=True, supports_dimension_scale_factor=True) parser.add_image_to_image_arguments(required=False) parser.add_output_arguments() args = parser.parse_args() @@ -24,33 +25,42 @@ def main(): args.guidance = ui_defaults.GUIDANCE_SCALE # 1. Load the model - model = Flux1( + flux = Flux1( model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=model) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: + # Resolve dimensions (supports ScaleFactor like "2x" when --image-path is provided) + width, height = DimensionResolver.resolve( + height=args.height, + width=args.width, + reference_image_path=args.image_path, + ) + for seed in args.seed: # 3. Generate an image for each seed value - image = model.generate_image( + image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - negative_prompt=PromptUtils.get_effective_negative_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - image_strength=args.image_strength, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + width=width, + height=height, + guidance=args.guidance, + scheduler=args.scheduler, + image_path=args.image_path, + num_inference_steps=args.steps, + image_strength=args.image_strength, + negative_prompt=PromptUtil.read_negative_prompt(args), ) # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) diff --git a/src/mflux/generate_controlnet.py b/src/mflux/models/flux/cli/flux_generate_controlnet.py similarity index 69% rename from src/mflux/generate_controlnet.py rename to src/mflux/models/flux/cli/flux_generate_controlnet.py index 20148b5..87a7782 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/models/flux/cli/flux_generate_controlnet.py @@ -1,11 +1,11 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -27,29 +27,32 @@ def main(): flux = Flux1Controlnet( model_config=_get_controlnet_model_config(args.model), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux, enable_canny_saver=True) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + enable_canny_saver=True, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + scheduler=args.scheduler, + num_inference_steps=args.steps, + controlnet_strength=args.controlnet_strength, controlnet_image_path=args.controlnet_image_path, - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - controlnet_strength=args.controlnet_strength, - scheduler=args.scheduler, - ), ) # 4. Save the image diff --git a/src/mflux/generate_depth.py b/src/mflux/models/flux/cli/flux_generate_depth.py similarity index 65% rename from src/mflux/generate_depth.py rename to src/mflux/models/flux/cli/flux_generate_depth.py index bef1ac6..fae2467 100644 --- a/src/mflux/generate_depth.py +++ b/src/mflux/models/flux/cli/flux_generate_depth.py @@ -1,10 +1,10 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.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 PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -25,29 +25,32 @@ def main(): # 1. Load the model flux = Flux1Depth( quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux, enable_depth_saver=True) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + enable_depth_saver=True, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - depth_image_path=args.depth_image_path, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + scheduler=args.scheduler, + image_path=args.image_path, + num_inference_steps=args.steps, + depth_image_path=args.depth_image_path, ) # 4. Save the image diff --git a/src/mflux/generate_fill.py b/src/mflux/models/flux/cli/flux_generate_fill.py similarity index 66% rename from src/mflux/generate_fill.py rename to src/mflux/models/flux/cli/flux_generate_fill.py index 0e27ea9..74789f1 100644 --- a/src/mflux/generate_fill.py +++ b/src/mflux/models/flux/cli/flux_generate_fill.py @@ -1,10 +1,10 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.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 PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -25,29 +25,31 @@ def main(): # 1. Load the model flux = Flux1Fill( quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - masked_image_path=args.masked_image_path, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + scheduler=args.scheduler, + image_path=args.image_path, + num_inference_steps=args.steps, + masked_image_path=args.masked_image_path, ) # 4. Save the image diff --git a/src/mflux/generate_in_context_catvton.py b/src/mflux/models/flux/cli/flux_generate_in_context_catvton.py similarity index 76% rename from src/mflux/generate_in_context_catvton.py rename to src/mflux/models/flux/cli/flux_generate_in_context_catvton.py index b80fd7b..7928fdf 100644 --- a/src/mflux/generate_in_context_catvton.py +++ b/src/mflux/models/flux/cli/flux_generate_in_context_catvton.py @@ -1,13 +1,13 @@ from pathlib import Path from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -38,31 +38,32 @@ def main(): flux = Flux1InContextFill( model_config=ModelConfig.dev_fill_catvton(), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + scheduler=args.scheduler, + num_inference_steps=args.steps, 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, - scheduler=args.scheduler, - ), + masked_image_path=args.person_mask, ) # 4. Save the image(s) diff --git a/src/mflux/generate_in_context_dev.py b/src/mflux/models/flux/cli/flux_generate_in_context_dev.py similarity index 61% rename from src/mflux/generate_in_context_dev.py rename to src/mflux/models/flux/cli/flux_generate_in_context_dev.py index 42dcca6..da5a1ad 100644 --- a/src/mflux/generate_in_context_dev.py +++ b/src/mflux/models/flux/cli/flux_generate_in_context_dev.py @@ -1,14 +1,14 @@ from pathlib import Path from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.in_context.flux_in_context_dev import Flux1InContextDev -from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.models.flux.variants.in_context.utils.in_context_loras import get_lora_path from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -31,33 +31,44 @@ def main(): if args.vae_tiling: args.vae_tiling_split = "vertical" + # Build lora_paths: style LoRA (if specified) + user-provided LoRAs + lora_paths = [] + lora_scales = [] + if args.lora_style: + lora_paths.append(get_lora_path(args.lora_style)) + lora_scales.append(1.0) + if args.lora_paths: + lora_paths.extend(args.lora_paths) + lora_scales.extend(args.lora_scales or [1.0] * len(args.lora_paths)) + # 1. Load the model flux = Flux1InContextDev( model_config=ModelConfig.dev(), quantize=args.quantize, - lora_names=[get_lora_filename(args.lora_style)] if args.lora_style else None, - lora_repo_id=LORA_REPO_ID if args.lora_style else None, - lora_paths=args.lora_paths, - lora_scales=args.lora_scales, + model_path=args.model_path, + lora_paths=lora_paths or None, + lora_scales=lora_scales or None, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + scheduler=args.scheduler, + image_path=args.image_path, + num_inference_steps=args.steps, ) # 4. Save the image output_path = Path(args.output.format(seed=seed)) diff --git a/src/mflux/generate_in_context_edit.py b/src/mflux/models/flux/cli/flux_generate_in_context_edit.py similarity index 50% rename from src/mflux/generate_in_context_edit.py rename to src/mflux/models/flux/cli/flux_generate_in_context_edit.py index cd1daa4..d3b85be 100644 --- a/src/mflux/generate_in_context_edit.py +++ b/src/mflux/models/flux/cli/flux_generate_in_context_edit.py @@ -1,15 +1,13 @@ from pathlib import Path -from PIL import Image - from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill -from mflux.models.flux.variants.in_context.utils.in_context_loras import prepare_ic_edit_loras -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.models.flux.variants.in_context.utils.in_context_fill_util import FluxInContextFillUtil +from mflux.models.flux.variants.in_context.utils.in_context_loras import IC_EDIT_LORA_SCALE, get_ic_edit_lora_path from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException @@ -34,35 +32,44 @@ def main(): args.vae_tiling_split = "vertical" # Auto-resize to optimal width for IC-Edit - width, height = _resize_for_ic_edit_optimal_width(args) + width, height = FluxInContextFillUtil.resize_for_ic_edit_optimal_width(args) + + # Build lora_paths: IC-Edit LoRA (required) + user-provided LoRAs + lora_paths = [get_ic_edit_lora_path()] + lora_scales = [IC_EDIT_LORA_SCALE] + if args.lora_paths: + lora_paths.extend(args.lora_paths) + lora_scales.extend(args.lora_scales or [1.0] * len(args.lora_paths)) # 1. Load the model with IC-Edit LoRA flux = Flux1InContextFill( 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, + model_path=args.model_path, + lora_paths=lora_paths, + lora_scales=lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) 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), + prompt=FluxInContextFillUtil.get_effective_ic_edit_prompt(args), + width=width, + height=height, + guidance=args.guidance, + scheduler=args.scheduler, + num_inference_steps=args.steps, left_image_path=args.reference_image, right_image_path=None, - config=Config( - num_inference_steps=args.steps, - height=height, - width=width, - guidance=args.guidance, - scheduler=args.scheduler, - ), ) # 4. Save the image(s) @@ -78,28 +85,5 @@ def main(): 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 PromptUtils.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() diff --git a/src/mflux/generate_kontext.py b/src/mflux/models/flux/cli/flux_generate_kontext.py similarity index 60% rename from src/mflux/generate_kontext.py rename to src/mflux/models/flux/cli/flux_generate_kontext.py index 6f89a29..794fd51 100644 --- a/src/mflux/generate_kontext.py +++ b/src/mflux/models/flux/cli/flux_generate_kontext.py @@ -1,12 +1,13 @@ from pathlib import Path from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.utils.dimension_resolver import DimensionResolver from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -15,7 +16,7 @@ def main(): parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) parser.add_lora_arguments() - parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_generator_arguments(supports_metadata_config=True, supports_dimension_scale_factor=True) parser.add_image_to_image_arguments(required=True) parser.add_output_arguments() args = parser.parse_args() @@ -27,28 +28,37 @@ def main(): # 1. Load the model flux = Flux1Kontext( quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: + # Resolve dimensions (supports ScaleFactor like "2x" based on --image-path) + width, height = DimensionResolver.resolve( + width=args.width, + height=args.height, + reference_image_path=args.image_path, + ) + for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + width=width, + height=height, + guidance=args.guidance, + scheduler=args.scheduler, + image_path=args.image_path, + num_inference_steps=args.steps, ) # 4. Save the image diff --git a/src/mflux/generate_redux.py b/src/mflux/models/flux/cli/flux_generate_redux.py similarity index 50% rename from src/mflux/generate_redux.py rename to src/mflux/models/flux/cli/flux_generate_redux.py index e198259..04621cd 100644 --- a/src/mflux/generate_redux.py +++ b/src/mflux/models/flux/cli/flux_generate_redux.py @@ -1,13 +1,12 @@ -from pathlib import Path - from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.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 PromptUtils +from mflux.models.flux.variants.redux.redux_util import ReduxUtil from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -26,7 +25,7 @@ def main(): args.guidance = ui_defaults.GUIDANCE_SCALE # Validate and normalize redux image strengths - redux_image_strengths = _validate_redux_image_strengths( + redux_image_strengths = ReduxUtil.validate_redux_image_strengths( redux_image_paths=args.redux_image_paths, redux_image_strengths=args.redux_image_strengths, ) @@ -35,29 +34,31 @@ def main(): flux = Flux1Redux( model_config=ModelConfig.dev_redux(), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) try: for seed in args.seed: # 3. Generate an image for each seed value image = flux.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - redux_image_paths=args.redux_image_paths, - redux_image_strengths=redux_image_strengths, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + scheduler=args.scheduler, + num_inference_steps=args.steps, + redux_image_paths=args.redux_image_paths, + redux_image_strengths=redux_image_strengths, ) # 4. Save the image @@ -69,25 +70,5 @@ def main(): print(memory_saver.memory_stats()) -def _validate_redux_image_strengths( - redux_image_paths: list[Path], - redux_image_strengths: list[float] | None, -) -> list[float] | None: - if not redux_image_strengths or len(redux_image_strengths) == 0: - return redux_image_strengths - - # If strengths provided but not enough for all images, fill with default (1.0) - if len(redux_image_strengths) < len(redux_image_paths): - redux_image_strengths.extend([1.0] * (len(redux_image_paths) - len(redux_image_strengths))) - - # If too many strengths provided, raise error - elif len(redux_image_strengths) > len(redux_image_paths): - raise ValueError( - f"Too many strengths provided ({len(redux_image_strengths)}), expted at most {len(redux_image_paths)}." - ) - - return redux_image_strengths - - if __name__ == "__main__": main() diff --git a/src/mflux/models/flux/cli/flux_upscale.py b/src/mflux/models/flux/cli/flux_upscale.py new file mode 100644 index 0000000..87fd1a6 --- /dev/null +++ b/src/mflux/models/flux/cli/flux_upscale.py @@ -0,0 +1,68 @@ +from mflux.callbacks.callback_manager import CallbackManager +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.config import ModelConfig +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator +from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet +from mflux.utils.dimension_resolver import DimensionResolver +from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil + + +def main(): + # 0. Parse command line arguments + parser = CommandLineParser(description="Upscale an image.") + 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, supports_dimension_scale_factor=True) + parser.add_controlnet_arguments(require_image=True) + parser.add_output_arguments() + args = parser.parse_args() + + # 1. Load the model + flux = Flux1Controlnet( + model_config=ModelConfig.dev_controlnet_upscaler(), + quantize=args.quantize, + model_path=args.model_path, + lora_paths=args.lora_paths, + lora_scales=args.lora_scales, + ) + + # 2. Register the optional callbacks + memory_saver = CallbackManager.register_callbacks( + args=args, + model=flux, + latent_creator=FluxLatentCreator, + ) + + try: + # Resolve dimensions using the controlnet image as reference + width, height = DimensionResolver.resolve( + height=args.height, + width=args.width, + reference_image_path=args.controlnet_image_path, + ) + + for seed in args.seed: + # 3. Generate an upscaled image for each seed value + image = flux.generate_image( + seed=seed, + prompt=PromptUtil.read_prompt(args), + width=width, + height=height, + num_inference_steps=args.steps, + controlnet_strength=args.controlnet_strength, + controlnet_image_path=args.controlnet_image_path, + ) + + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) + finally: + if memory_saver: + print(memory_saver.memory_stats()) + + +if __name__ == "__main__": + main() diff --git a/src/mflux/models/flux/flux_initializer.py b/src/mflux/models/flux/flux_initializer.py index 4db5cd3..36e480a 100644 --- a/src/mflux/models/flux/flux_initializer.py +++ b/src/mflux/models/flux/flux_initializer.py @@ -1,236 +1,217 @@ -from mflux.config.model_config import ModelConfig -from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.models.common.config import ModelConfig +from mflux.models.common.lora.mapping.lora_loader import LoRALoader +from mflux.models.common.tokenizer import TokenizerLoader +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_loader import WeightLoader +from mflux.models.depth_pro.model.depth_pro import DepthPro from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.model.redux_encoder.redux_encoder import ReduxEncoder from mflux.models.flux.model.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer -from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP -from mflux.models.flux.tokenizer.t5_tokenizer import TokenizerT5 -from mflux.models.flux.tokenizer.tokenizer_handler import TokenizerHandler from mflux.models.flux.variants.controlnet.transformer_controlnet import TransformerControlnet -from mflux.models.flux.variants.controlnet.weight_handler_controlnet import WeightHandlerControlnet -from mflux.models.flux.variants.redux.weight_handler_redux import WeightHandlerRedux -from mflux.models.common.lora.download.lora_huggingface_downloader import LoRAHuggingFaceDownloader -from mflux.models.flux.weights.weight_handler import WeightHandler -from mflux.models.flux.weights.weight_handler_lora import WeightHandlerLoRA -from mflux.models.flux.weights.weight_util import WeightUtil +from mflux.models.flux.weights.flux_lora_mapping import FluxLoRAMapping +from mflux.models.flux.weights.flux_weight_definition import ( + FluxControlnetWeightDefinition, + FluxReduxWeightDefinition, + FluxWeightDefinition, +) class FluxInitializer: @staticmethod def init( - flux_model, + model, model_config: ModelConfig, quantize: int | None, - local_path: str | None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, custom_transformer=None, ) -> None: - # 0. Set paths, configs, and prompt_cache for later - lora_paths = lora_paths or [] - flux_model.prompt_cache = {} - flux_model.model_config = model_config - - # 1. Load the regular weights - weights = WeightHandler.load_regular_weights( - repo_id=model_config.model_name, - local_path=local_path, - transformer_repo_id=model_config.custom_transformer_model, - ) - - # 2. Initialize tokenizers - tokenizers = TokenizerHandler( - repo_id=model_config.model_name, - max_t5_length=model_config.max_sequence_length, - local_path=local_path, - ) - flux_model.t5_tokenizer = TokenizerT5( - tokenizer=tokenizers.t5, - max_length=model_config.max_sequence_length, - ) - flux_model.clip_tokenizer = TokenizerCLIP( - tokenizer=tokenizers.clip, - ) - - # 3. Initialize all models - flux_model.vae = VAE() - flux_model.t5_text_encoder = T5Encoder() - flux_model.clip_text_encoder = CLIPEncoder() - if custom_transformer is not None: - flux_model.transformer = custom_transformer - else: - flux_model.transformer = Transformer( - model_config=model_config, - num_transformer_blocks=weights.num_transformer_blocks(), - num_single_transformer_blocks=weights.num_single_transformer_blocks(), - ) - - # 4. Apply weights and quantize the models - flux_model.bits = WeightUtil.set_weights_and_quantize( - quantize_arg=quantize, - weights=weights, - vae=flux_model.vae, - transformer=flux_model.transformer, - t5_text_encoder=flux_model.t5_text_encoder, - clip_text_encoder=flux_model.clip_text_encoder, - ) - - # 5. Set LoRA weights - hf_lora_paths = LoRAHuggingFaceDownloader.download_loras( - lora_names=lora_names, - repo_id=lora_repo_id, - model_name="Flux", - ) - flux_model.lora_paths = lora_paths + hf_lora_paths - flux_model.lora_scales = (lora_scales or []) + [1.0] * len(hf_lora_paths) - lora_weights = WeightHandlerLoRA.load_lora_weights( - transformer=flux_model.transformer, - lora_files=flux_model.lora_paths, - lora_scales=flux_model.lora_scales, - ) - WeightHandlerLoRA.set_lora_weights( - transformer=flux_model.transformer, - loras=lora_weights, - ) + path = model_path if model_path else model_config.model_name + FluxInitializer._init_config(model, model_config) + weights = FluxInitializer._load_weights(path) + FluxInitializer._init_tokenizers(model, path, model_config) + FluxInitializer._init_models(model, model_config, weights, custom_transformer) + FluxInitializer._apply_weights(model, weights, quantize) + FluxInitializer._apply_lora(model, lora_paths, lora_scales) @staticmethod def init_depth( - flux_model, + model, model_config: ModelConfig, quantize: int | None, - local_path: str | None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, - ): - # 1. Start with the same init as regular Flux + ) -> None: FluxInitializer.init( - flux_model=flux_model, + model=model, model_config=model_config, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, ) - - # 2. Initialize the DepthPro model - flux_model.depth_pro = DepthPro() + model.depth_pro = DepthPro() @staticmethod def init_redux( - flux_model, + model, + model_config: ModelConfig, quantize: int | None, - local_path: str | None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, - ): - # 1. Start with the same init as regular Flux dev + ) -> None: FluxInitializer.init( - flux_model=flux_model, - model_config=ModelConfig.dev(), + model=model, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, + model_config=model_config, ) - # 2. Initialize the redux specific addons - redux_weights = WeightHandlerRedux.load_weights() - flux_model.image_embedder = ReduxEncoder() - flux_model.image_encoder = SiglipVisionTransformer() - WeightUtil.set_redux_weights_and_quantize( - quantize_arg=quantize, + redux_weights = WeightLoader.load( + weight_definition=FluxReduxWeightDefinition, + model_path=ModelConfig.dev_redux().model_name, + ) + model.image_embedder = ReduxEncoder() + model.image_encoder = SiglipVisionTransformer() + WeightApplier.apply_and_quantize( weights=redux_weights, - redux_encoder=flux_model.image_embedder, - siglip_vision_transformer=flux_model.image_encoder, + quantize_arg=quantize, + weight_definition=FluxReduxWeightDefinition, + models={ + "siglip": model.image_encoder, + "redux_encoder": model.image_embedder, + }, ) @staticmethod def init_controlnet( - flux_model, + model, model_config: ModelConfig, quantize: int | None, - local_path: str | None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, ) -> None: - # 1. Start with the same init as regular Flux FluxInitializer.init( - flux_model=flux_model, + model=model, model_config=model_config, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, ) - # 2. Apply ControlNet-specific initialization - weights_controlnet = WeightHandlerControlnet.load_controlnet_transformer( - controlnet_model=model_config.controlnet_model + controlnet_component = FluxControlnetWeightDefinition.get_controlnet_component() + controlnet_weights = WeightLoader.load_single( + component=controlnet_component, + repo_id=model_config.controlnet_model, ) - flux_model.transformer_controlnet = TransformerControlnet( + model.transformer_controlnet = TransformerControlnet( model_config=model_config, - num_transformer_blocks=weights_controlnet.num_transformer_blocks(), - num_single_transformer_blocks=weights_controlnet.num_single_transformer_blocks(), + num_transformer_blocks=controlnet_weights.num_transformer_blocks(), + num_single_transformer_blocks=controlnet_weights.num_single_transformer_blocks(), ) - WeightUtil.set_controlnet_weights_and_quantize( + WeightApplier.apply_and_quantize_single( + weights=controlnet_weights, + model=model.transformer_controlnet, + component=controlnet_component, quantize_arg=quantize, - weights=weights_controlnet, - transformer_controlnet=flux_model.transformer_controlnet, + quantization_predicate=FluxWeightDefinition.quantization_predicate, ) @staticmethod def init_concept( - flux_model, + model, model_config: ModelConfig, quantize: int | None, - local_path: str | None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, - ): - # Import here to avoid circular dependency + ) -> None: from mflux.models.flux.variants.concept_attention.transformer_concept import TransformerConcept - # 1. Load weights first to get flux_transformer dimensions - weights = WeightHandler.load_regular_weights( - repo_id=model_config.model_name, - local_path=local_path, - ) - - # 2. Create custom TransformerConcept + path = model_path if model_path else model_config.model_name + FluxInitializer._init_config(model, model_config) + weights = FluxInitializer._load_weights(path) + FluxInitializer._init_tokenizers(model, path, model_config) custom_transformer = TransformerConcept( model_config=model_config, num_transformer_blocks=weights.num_transformer_blocks(), num_single_transformer_blocks=weights.num_single_transformer_blocks(), ) + FluxInitializer._init_models(model, model_config, weights, custom_transformer) + FluxInitializer._apply_weights(model, weights, quantize) + FluxInitializer._apply_lora(model, lora_paths, lora_scales) - # 3. Use the improved FluxInitializer with custom flux_transformer - FluxInitializer.init( - flux_model=flux_model, - model_config=model_config, - quantize=quantize, - local_path=local_path, + @staticmethod + def _init_config(model, model_config: ModelConfig) -> None: + model.prompt_cache = {} + model.model_config = model_config + model.callbacks = CallbackRegistry() + + @staticmethod + def _load_weights(model_path: str) -> LoadedWeights: + return WeightLoader.load( + weight_definition=FluxWeightDefinition, + model_path=model_path, + ) + + @staticmethod + def _init_tokenizers(model, model_path: str, model_config: ModelConfig) -> None: + model.tokenizers = TokenizerLoader.load_all( + definitions=FluxWeightDefinition.get_tokenizers(), + model_path=model_path, + max_length_overrides={"t5": model_config.max_sequence_length}, + ) + + @staticmethod + def _init_models( + model, + model_config: ModelConfig, + weights: LoadedWeights, + custom_transformer=None, + ) -> None: + model.vae = VAE() + model.t5_text_encoder = T5Encoder() + model.clip_text_encoder = CLIPEncoder() + if custom_transformer is not None: + model.transformer = custom_transformer + else: + model.transformer = Transformer( + model_config=model_config, + num_transformer_blocks=weights.num_transformer_blocks(), + num_single_transformer_blocks=weights.num_single_transformer_blocks(), + ) + + @staticmethod + def _apply_weights(model, weights: LoadedWeights, quantize: int | None) -> None: + model.bits = WeightApplier.apply_and_quantize( + weights=weights, + quantize_arg=quantize, + weight_definition=FluxWeightDefinition, + models={ + "vae": model.vae, + "transformer": model.transformer, + "t5_encoder": model.t5_text_encoder, + "clip_encoder": model.clip_text_encoder, + }, + ) + + @staticmethod + def _apply_lora(model, lora_paths: list[str] | None, lora_scales: list[float] | None) -> None: + model.lora_paths, model.lora_scales = LoRALoader.load_and_apply_lora( + lora_mapping=FluxLoRAMapping.get_mapping(), + transformer=model.transformer, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, - custom_transformer=custom_transformer, ) diff --git a/src/mflux/models/flux/latent_creator/flux_latent_creator.py b/src/mflux/models/flux/latent_creator/flux_latent_creator.py index 89d393e..2b41462 100644 --- a/src/mflux/models/flux/latent_creator/flux_latent_creator.py +++ b/src/mflux/models/flux/latent_creator/flux_latent_creator.py @@ -10,7 +10,13 @@ class FluxLatentCreator: ) @staticmethod - def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: - latents = mx.reshape(latents, (1, 16, height // 16, 2, width // 16, 2)) + def pack_latents(latents: mx.array, height: int, width: int, num_channels_latents: int = 16) -> mx.array: + latents = mx.reshape(latents, (1, num_channels_latents, height // 16, 2, width // 16, 2)) latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5)) - return mx.reshape(latents, (1, (width // 16) * (height // 16), 64)) + return mx.reshape(latents, (1, (width // 16) * (height // 16), num_channels_latents * 4)) + + @staticmethod + def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: + latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2)) + latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5)) + return mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2)) diff --git a/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.py b/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.py index cc6c553..5f5782d 100644 --- a/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.py +++ b/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_embeddings.py @@ -1,13 +1,11 @@ import mlx.core as mx from mlx import nn -from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP - class CLIPEmbeddings(nn.Module): def __init__(self, dims: int): super().__init__() - self.position_embedding = nn.Embedding(num_embeddings=TokenizerCLIP.MAX_TOKEN_LENGTH, dims=dims) + self.position_embedding = nn.Embedding(num_embeddings=77, dims=dims) self.token_embedding = nn.Embedding(num_embeddings=49408, dims=dims) def __call__(self, tokens: mx.array) -> mx.array: diff --git a/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_sdpa_attention.py b/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_sdpa_attention.py index cc560ea..7127387 100644 --- a/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_sdpa_attention.py +++ b/src/mflux/models/flux/model/flux_text_encoder/clip_encoder/clip_sdpa_attention.py @@ -2,7 +2,7 @@ import mlx.core as mx from mlx import nn from mlx.core.fast import scaled_dot_product_attention -from mflux.config.config import Config +from mflux.models.common.config import ModelConfig class CLIPSdpaAttention(nn.Module): @@ -18,7 +18,7 @@ class CLIPSdpaAttention(nn.Module): self.out_proj = nn.Linear(input_dims=768, output_dims=768) def __call__(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array: - causal_attention_mask = causal_attention_mask.astype(Config.precision) + causal_attention_mask = causal_attention_mask.astype(ModelConfig.precision) query = self.q_proj(hidden_states) key = self.k_proj(hidden_states) diff --git a/src/mflux/models/flux/model/flux_text_encoder/prompt_encoder.py b/src/mflux/models/flux/model/flux_text_encoder/prompt_encoder.py index 5b702d6..a119cd6 100644 --- a/src/mflux/models/flux/model/flux_text_encoder/prompt_encoder.py +++ b/src/mflux/models/flux/model/flux_text_encoder/prompt_encoder.py @@ -1,9 +1,8 @@ import mlx.core as mx +from mflux.models.common.tokenizer import Tokenizer from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder -from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP -from mflux.models.flux.tokenizer.t5_tokenizer import TokenizerT5 class PromptEncoder: @@ -11,8 +10,8 @@ class PromptEncoder: def encode_prompt( prompt: str, prompt_cache: dict[str, tuple[mx.array, mx.array]], - t5_tokenizer: TokenizerT5, - clip_tokenizer: TokenizerCLIP, + t5_tokenizer: Tokenizer, + clip_tokenizer: Tokenizer, t5_text_encoder: T5Encoder, clip_text_encoder: CLIPEncoder, ) -> tuple[mx.array, mx.array]: @@ -21,10 +20,10 @@ class PromptEncoder: return prompt_cache[prompt] # 1. Encode the prompt - t5_tokens = t5_tokenizer.tokenize(prompt) - clip_tokens = clip_tokenizer.tokenize(prompt) - prompt_embeds = t5_text_encoder(t5_tokens) - pooled_prompt_embeds = clip_text_encoder(clip_tokens) + t5_output = t5_tokenizer.tokenize(prompt) + clip_output = clip_tokenizer.tokenize(prompt) + prompt_embeds = t5_text_encoder(t5_output.input_ids) + pooled_prompt_embeds = clip_text_encoder(clip_output.input_ids) # 2. Cache the encoded prompt prompt_cache[prompt] = (prompt_embeds, pooled_prompt_embeds) diff --git a/src/mflux/models/flux/model/flux_transformer/ada_layer_norm_continuous.py b/src/mflux/models/flux/model/flux_transformer/ada_layer_norm_continuous.py index 801e54c..3599851 100644 --- a/src/mflux/models/flux/model/flux_transformer/ada_layer_norm_continuous.py +++ b/src/mflux/models/flux/model/flux_transformer/ada_layer_norm_continuous.py @@ -1,7 +1,7 @@ import mlx.core as mx from mlx import nn -from mflux.config.config import Config +from mflux.models.common.config import ModelConfig class AdaLayerNormContinuous(nn.Module): @@ -12,7 +12,7 @@ class AdaLayerNormContinuous(nn.Module): self.norm = nn.LayerNorm(dims=embedding_dim, eps=1e-6, affine=False) def __call__(self, x: mx.array, text_embeddings: mx.array) -> mx.array: - text_embeddings = self.linear(nn.silu(text_embeddings).astype(Config.precision)) + text_embeddings = self.linear(nn.silu(text_embeddings).astype(ModelConfig.precision)) chunk_size = self.embedding_dim scale = text_embeddings[:, 0 * chunk_size : 1 * chunk_size] shift = text_embeddings[:, 1 * chunk_size : 2 * chunk_size] diff --git a/src/mflux/models/flux/model/flux_transformer/time_text_embed.py b/src/mflux/models/flux/model/flux_transformer/time_text_embed.py index a979fd6..87d008d 100644 --- a/src/mflux/models/flux/model/flux_transformer/time_text_embed.py +++ b/src/mflux/models/flux/model/flux_transformer/time_text_embed.py @@ -3,8 +3,7 @@ import math import mlx.core as mx from mlx import nn -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.model.flux_transformer.guidance_embedder import GuidanceEmbedder from mflux.models.flux.model.flux_transformer.text_embedder import TextEmbedder from mflux.models.flux.model.flux_transformer.timestep_embedder import TimestepEmbedder @@ -29,7 +28,7 @@ class TimeTextEmbed(nn.Module): time_steps_emb += self.guidance_embedder(self._time_proj(guidance)) pooled_projections = self.text_embedder(pooled_projection) conditioning = time_steps_emb + pooled_projections - return conditioning.astype(Config.precision) + return conditioning.astype(ModelConfig.precision) @staticmethod def _time_proj(time_steps: mx.array) -> mx.array: diff --git a/src/mflux/models/flux/model/flux_transformer/transformer.py b/src/mflux/models/flux/model/flux_transformer/transformer.py index e8228ba..e72eb38 100644 --- a/src/mflux/models/flux/model/flux_transformer/transformer.py +++ b/src/mflux/models/flux/model/flux_transformer/transformer.py @@ -3,8 +3,8 @@ import math import mlx.core as mx from mlx import nn -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.model.flux_transformer.ada_layer_norm_continuous import AdaLayerNormContinuous from mflux.models.flux.model.flux_transformer.embed_nd import EmbedND from mflux.models.flux.model.flux_transformer.joint_transformer_block import JointTransformerBlock @@ -32,7 +32,7 @@ class Transformer(nn.Module): def __call__( self, t: int, - config: RuntimeConfig, + config: Config, hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, @@ -130,7 +130,7 @@ class Transformer(nn.Module): def compute_rotary_embeddings( prompt_embeds: mx.array, pos_embed: EmbedND, - config: RuntimeConfig, + config: Config, kontext_image_ids: mx.array | None = None, ) -> mx.array: txt_ids = Transformer._prepare_text_ids(seq_len=prompt_embeds.shape[1]) @@ -148,7 +148,7 @@ class Transformer(nn.Module): t: int, pooled_prompt_embeds: mx.array, time_text_embed: TimeTextEmbed, - config: RuntimeConfig, + config: Config, ) -> mx.array: time_step = config.scheduler.sigmas[t] * config.num_train_steps time_step = mx.broadcast_to(time_step, (1,)).astype(config.precision) diff --git a/src/mflux/models/flux/model/flux_vae/common/attention.py b/src/mflux/models/flux/model/flux_vae/common/attention.py index 1ad1568..374610a 100644 --- a/src/mflux/models/flux/model/flux_vae/common/attention.py +++ b/src/mflux/models/flux/model/flux_vae/common/attention.py @@ -2,7 +2,7 @@ import mlx.core as mx from mlx import nn from mlx.core.fast import scaled_dot_product_attention -from mflux.config.config import Config +from mflux.models.common.config import ModelConfig class Attention(nn.Module): @@ -19,7 +19,7 @@ class Attention(nn.Module): B, H, W, C = input_array.shape - y = self.group_norm(input_array.astype(mx.float32)).astype(Config.precision) + y = self.group_norm(input_array.astype(mx.float32)).astype(ModelConfig.precision) queries = self.to_q(y).reshape(B, H * W, 1, C) keys = self.to_k(y).reshape(B, H * W, 1, C) diff --git a/src/mflux/models/flux/model/flux_vae/common/resnet_block_2d.py b/src/mflux/models/flux/model/flux_vae/common/resnet_block_2d.py index fe1e65a..3e1932b 100644 --- a/src/mflux/models/flux/model/flux_vae/common/resnet_block_2d.py +++ b/src/mflux/models/flux/model/flux_vae/common/resnet_block_2d.py @@ -1,7 +1,7 @@ import mlx.core as mx from mlx import nn -from mflux.config.config import Config +from mflux.models.common.config import ModelConfig class ResnetBlock2D(nn.Module): @@ -60,10 +60,10 @@ class ResnetBlock2D(nn.Module): def __call__(self, input_array: mx.array) -> mx.array: input_array = mx.transpose(input_array, (0, 2, 3, 1)) - hidden_states = self.norm1(input_array.astype(mx.float32)).astype(Config.precision) + hidden_states = self.norm1(input_array.astype(mx.float32)).astype(ModelConfig.precision) hidden_states = nn.silu(hidden_states) hidden_states = self.conv1(hidden_states) - hidden_states = self.norm2(hidden_states.astype(mx.float32)).astype(Config.precision) + hidden_states = self.norm2(hidden_states.astype(mx.float32)).astype(ModelConfig.precision) hidden_states = nn.silu(hidden_states) hidden_states = self.conv2(hidden_states) if self.conv_shortcut is not None: diff --git a/src/mflux/models/flux/model/flux_vae/decoder/conv_norm_out.py b/src/mflux/models/flux/model/flux_vae/decoder/conv_norm_out.py index 05e5fe8..7937e75 100644 --- a/src/mflux/models/flux/model/flux_vae/decoder/conv_norm_out.py +++ b/src/mflux/models/flux/model/flux_vae/decoder/conv_norm_out.py @@ -1,7 +1,7 @@ import mlx.core as mx import mlx.nn as nn -from mflux.config.config import Config +from mflux.models.common.config import ModelConfig class ConvNormOut(nn.Module): @@ -17,5 +17,5 @@ class ConvNormOut(nn.Module): def __call__(self, input_array: mx.array) -> mx.array: input_array = mx.transpose(input_array, (0, 2, 3, 1)) - hidden_states = self.norm(input_array.astype(mx.float32)).astype(Config.precision) + hidden_states = self.norm(input_array.astype(mx.float32)).astype(ModelConfig.precision) return mx.transpose(hidden_states, (0, 3, 1, 2)) diff --git a/src/mflux/models/flux/tokenizer/clip_tokenizer.py b/src/mflux/models/flux/tokenizer/clip_tokenizer.py deleted file mode 100644 index 2508241..0000000 --- a/src/mflux/models/flux/tokenizer/clip_tokenizer.py +++ /dev/null @@ -1,20 +0,0 @@ -import mlx.core as mx -from transformers import CLIPTokenizer - - -class TokenizerCLIP: - MAX_TOKEN_LENGTH = 77 - - def __init__(self, tokenizer: CLIPTokenizer): - self.tokenizer = tokenizer - - def tokenize(self, prompt: str) -> mx.array: - return self.tokenizer( - [prompt], - padding="max_length", - max_length=TokenizerCLIP.MAX_TOKEN_LENGTH, - truncation=True, - return_length=False, - return_overflowing_tokens=False, - return_tensors="mlx", - ).input_ids diff --git a/src/mflux/models/flux/tokenizer/t5_tokenizer.py b/src/mflux/models/flux/tokenizer/t5_tokenizer.py deleted file mode 100644 index b59c1f0..0000000 --- a/src/mflux/models/flux/tokenizer/t5_tokenizer.py +++ /dev/null @@ -1,19 +0,0 @@ -import mlx.core as mx -from transformers import T5Tokenizer - - -class TokenizerT5: - def __init__(self, tokenizer: T5Tokenizer, max_length: int = 256): - self.tokenizer = tokenizer - self.max_length = max_length - - def tokenize(self, prompt: str) -> mx.array: - return self.tokenizer( - [prompt], - padding="max_length", - max_length=self.max_length, - truncation=True, - return_length=False, - return_overflowing_tokens=False, - return_tensors="mlx", - ).input_ids diff --git a/src/mflux/models/flux/tokenizer/tokenizer_handler.py b/src/mflux/models/flux/tokenizer/tokenizer_handler.py deleted file mode 100644 index e144025..0000000 --- a/src/mflux/models/flux/tokenizer/tokenizer_handler.py +++ /dev/null @@ -1,36 +0,0 @@ -from pathlib import Path - -import transformers - -from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP -from mflux.utils.download import snapshot_download - - -class TokenizerHandler: - def __init__( - self, - repo_id: str, - max_t5_length: int = 256, - local_path: str | None = None, - ): - root_path = Path(local_path) if local_path else TokenizerHandler._download_or_get_cached_tokenizers(repo_id) - - self.clip = transformers.CLIPTokenizer.from_pretrained( - pretrained_model_name_or_path=root_path / "tokenizer", - local_files_only=True, - max_length=TokenizerCLIP.MAX_TOKEN_LENGTH, - ) - self.t5 = transformers.T5Tokenizer.from_pretrained( - pretrained_model_name_or_path=root_path / "tokenizer_2", - local_files_only=True, - max_length=max_t5_length, - ) - - @staticmethod - def _download_or_get_cached_tokenizers(repo_id: str) -> Path: - return Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=["tokenizer/**", "tokenizer_2/**"], - ) - ) diff --git a/src/mflux/models/flux/variants/concept_attention/flux_concept.py b/src/mflux/models/flux/variants/concept_attention/flux_concept.py index 98d107a..7b90034 100644 --- a/src/mflux/models/flux/variants/concept_attention/flux_concept.py +++ b/src/mflux/models/flux/variants/concept_attention/flux_concept.py @@ -1,11 +1,10 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator @@ -16,7 +15,6 @@ from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.concept_attention.attention_data import GenerationAttentionData from mflux.models.flux.variants.concept_attention.concept_util import ConceptUtil from mflux.models.flux.variants.concept_attention.transformer_concept import TransformerConcept -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -30,20 +28,20 @@ class Flux1Concept(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.schnell(), ): super().__init__() FluxInitializer.init_concept( - flux_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( @@ -51,19 +49,33 @@ class Flux1Concept(nn.Module): seed: int, prompt: str, concept: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", heatmap_layer_indices: list[int] | None = None, heatmap_timesteps: list[int] | None = None, ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = LatentCreator.create_for_txt2img_or_img2img( seed=seed, - height=config.height, width=config.width, + height=config.height, img2img=Img2Img( vae=self.vae, latent_creator=FluxLatentCreator, @@ -77,8 +89,8 @@ class Flux1Concept(nn.Module): 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) @@ -87,27 +99,23 @@ class Flux1Concept(nn.Module): prompt_embeds_concept, pooled_prompt_embeds_concept = PromptEncoder.encode_prompt( prompt=concept, prompt_cache=self.prompt_cache, - t5_tokenizer=self.t5_tokenizer, - clip_tokenizer=self.clip_tokenizer, + t5_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) attention_data = GenerationAttentionData() - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Predict the noise + # 5.t Predict the noise noise, attention = self.transformer( t=t, config=config, @@ -119,46 +127,25 @@ class Flux1Concept(nn.Module): ) attention_data.append(attention) - # 5.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 6.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 7.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 8. Call subscribers after loop + ctx.after_loop(latents) - # 6. Generate concept attention heatmap + # 9. Generate concept attention heatmap concept_heatmap = ConceptUtil.create_heatmap( concept=concept, attention_data=attention_data, @@ -168,10 +155,9 @@ class Flux1Concept(nn.Module): timesteps=heatmap_timesteps or list(range(config.num_inference_steps)), ) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) - return ImageUtil.to_image( decoded_latents=decoded, config=config, @@ -182,6 +168,6 @@ class Flux1Concept(nn.Module): lora_scales=self.lora_scales, image_path=config.image_path, image_strength=config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], concept_heatmap=concept_heatmap, ) diff --git a/src/mflux/models/flux/variants/concept_attention/flux_concept_from_image.py b/src/mflux/models/flux/variants/concept_attention/flux_concept_from_image.py index 36a4035..b30764f 100644 --- a/src/mflux/models/flux/variants/concept_attention/flux_concept_from_image.py +++ b/src/mflux/models/flux/variants/concept_attention/flux_concept_from_image.py @@ -1,11 +1,10 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.latent_creator.latent_creator import LatentCreator from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator @@ -16,7 +15,6 @@ from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.concept_attention.attention_data import GenerationAttentionData from mflux.models.flux.variants.concept_attention.concept_util import ConceptUtil from mflux.models.flux.variants.concept_attention.transformer_concept import TransformerConcept -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -30,20 +28,20 @@ class Flux1ConceptFromImage(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.schnell(), ): super().__init__() FluxInitializer.init_concept( - flux_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( @@ -51,33 +49,46 @@ class Flux1ConceptFromImage(nn.Module): seed: int, prompt: str, concept: str, - image_path: str, - config: Config, + image_path: Path | str, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_strength: float | None = None, + scheduler: str = "linear", heatmap_layer_indices: list[int] | None = None, heatmap_timesteps: list[int] | None = None, ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents from the reference image encoded_image = LatentCreator.encode_image( vae=self.vae, - image_path=image_path, - height=config.height, width=config.width, + height=config.height, + image_path=image_path, ) # Create static noise for blending at each timestep static_noise = FluxLatentCreator.create_noise( seed=seed, - height=config.height, width=config.width, + height=config.height, ) # Start with an appropriately noised version of the encoded image latents = LatentCreator.add_noise_by_interpolation( - clean=ArrayUtil.pack_latents(latents=encoded_image, height=config.height, width=config.width), + clean=FluxLatentCreator.pack_latents(latents=encoded_image, height=config.height, width=config.width), noise=static_noise, sigma=config.scheduler.sigmas[config.init_time_step], ) @@ -86,8 +97,8 @@ class Flux1ConceptFromImage(nn.Module): 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) @@ -96,27 +107,23 @@ class Flux1ConceptFromImage(nn.Module): prompt_embeds_concept, pooled_prompt_embeds_concept = PromptEncoder.encode_prompt( prompt=concept, prompt_cache=self.prompt_cache, - t5_tokenizer=self.t5_tokenizer, - clip_tokenizer=self.clip_tokenizer, + t5_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) attention_data = GenerationAttentionData() - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Predict the noise (we don't use the noise, only the attention) + # 5.t Predict the noise (we don't use the noise, only the attention) _, attention = self.transformer( t=t, config=config, @@ -128,22 +135,17 @@ class Flux1ConceptFromImage(nn.Module): ) attention_data.append(attention) - # 5.t Follow reverse diffusion trajectory + # 6.t Follow reverse diffusion trajectory latents = LatentCreator.add_noise_by_interpolation( - clean=ArrayUtil.pack_latents(latents=encoded_image, height=config.height, width=config.width), + clean=FluxLatentCreator.pack_latents( + latents=encoded_image, height=config.height, width=config.width + ), noise=static_noise, sigma=config.scheduler.sigmas[t + 1], ) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 7.t Call subscribers in-loop + ctx.in_loop(t, latents) # Evaluate attention data to force MLX computation for progress tracking mx.eval( @@ -152,25 +154,15 @@ class Flux1ConceptFromImage(nn.Module): ) except KeyboardInterrupt: # noqa: PERF203 - Callbacks.interruption( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 8. Call subscribers after loop + ctx.after_loop(latents) - # 6. Generate concept attention heatmap + # 9. Generate concept attention heatmap concept_heatmap = ConceptUtil.create_heatmap( concept=concept, attention_data=attention_data, @@ -180,8 +172,8 @@ class Flux1ConceptFromImage(nn.Module): timesteps=heatmap_timesteps or list(range(config.num_inference_steps)), ) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( @@ -194,6 +186,6 @@ class Flux1ConceptFromImage(nn.Module): lora_scales=self.lora_scales, image_path=image_path, image_strength=config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], concept_heatmap=concept_heatmap, ) diff --git a/src/mflux/models/flux/variants/concept_attention/transformer_concept.py b/src/mflux/models/flux/variants/concept_attention/transformer_concept.py index 6e452c1..08734a5 100644 --- a/src/mflux/models/flux/variants/concept_attention/transformer_concept.py +++ b/src/mflux/models/flux/variants/concept_attention/transformer_concept.py @@ -1,8 +1,8 @@ import mlx.core as mx from mlx import nn -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.model.flux_transformer.ada_layer_norm_continuous import AdaLayerNormContinuous from mflux.models.flux.model.flux_transformer.embed_nd import EmbedND from mflux.models.flux.model.flux_transformer.single_transformer_block import SingleTransformerBlock @@ -32,7 +32,7 @@ class TransformerConcept(nn.Module): def __call__( self, t: int, - config: RuntimeConfig, + config: Config, hidden_states: mx.array, prompt_embeds: mx.array, prompt_embeds_concept: mx.array, diff --git a/src/mflux/models/flux/variants/controlnet/controlnet_util.py b/src/mflux/models/flux/variants/controlnet/controlnet_util.py index 911d436..6dabd60 100644 --- a/src/mflux/models/flux/variants/controlnet/controlnet_util.py +++ b/src/mflux/models/flux/variants/controlnet/controlnet_util.py @@ -5,8 +5,8 @@ import mlx.core as mx import numpy as np import PIL.Image +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_vae.vae import VAE -from mflux.utils.array_util import ArrayUtil from mflux.utils.image_util import StrOrBytesPath log = logging.getLogger(__name__) @@ -31,7 +31,7 @@ class ControlnetUtil: controlnet_cond = vae.encode(controlnet_cond) if is_canny: controlnet_cond = (controlnet_cond / vae.scaling_factor) + vae.shift_factor - controlnet_cond = ArrayUtil.pack_latents(latents=controlnet_cond, height=height, width=width) + controlnet_cond = FluxLatentCreator.pack_latents(latents=controlnet_cond, height=height, width=width) return controlnet_cond, control_image @staticmethod diff --git a/src/mflux/models/flux/variants/controlnet/flux_controlnet.py b/src/mflux/models/flux/variants/controlnet/flux_controlnet.py index 2061a18..8599897 100644 --- a/src/mflux/models/flux/variants/controlnet/flux_controlnet.py +++ b/src/mflux/models/flux/variants/controlnet/flux_controlnet.py @@ -1,12 +1,9 @@ import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig -from mflux.models.common.weights.model_saver import ModelSaver +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.common.weights.saving.model_saver import ModelSaver from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -16,7 +13,7 @@ from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.controlnet.controlnet_util import ControlnetUtil from mflux.models.flux.variants.controlnet.transformer_controlnet import TransformerControlnet -from mflux.utils.array_util import ArrayUtil +from mflux.models.flux.weights.flux_weight_definition import FluxControlnetWeightDefinition from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil, StrOrBytesPath @@ -31,21 +28,21 @@ class Flux1Controlnet(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, controlnet_path: str | None = None, + model_config: ModelConfig = ModelConfig.dev_controlnet_canny(), ): super().__init__() FluxInitializer.init_controlnet( - flux_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( @@ -53,17 +50,29 @@ class Flux1Controlnet(nn.Module): seed: int, prompt: str, controlnet_image_path: StrOrBytesPath, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + controlnet_strength: float = 1.0, + scheduler: str = "linear", ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + controlnet_strength=controlnet_strength, + ) # 1. Encode the controlnet reference image controlnet_condition, canny_image = ControlnetUtil.encode_image( vae=self.vae, - height=config.height, width=config.width, + height=config.height, controlnet_image_path=controlnet_image_path, is_canny=self.model_config.is_canny(), ) @@ -71,35 +80,30 @@ class Flux1Controlnet(nn.Module): # 2. Create the initial latents latents = FluxLatentCreator.create_noise( seed=seed, - height=config.height, width=config.width, + height=config.height, ) # 3. 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - canny_image=canny_image, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents, canny_image=canny_image) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Compute controlnet samples + # 5.t Compute controlnet samples controlnet_block_samples, controlnet_single_block_samples = self.transformer_controlnet( t=t, config=config, @@ -109,7 +113,7 @@ class Flux1Controlnet(nn.Module): controlnet_condition=controlnet_condition, ) - # 5.t Predict the noise + # 6.t Predict the noise noise = self.transformer( t=t, config=config, @@ -120,47 +124,26 @@ class Flux1Controlnet(nn.Module): controlnet_single_block_samples=controlnet_single_block_samples, ) - # 6.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 7.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 8.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 9. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -171,7 +154,7 @@ class Flux1Controlnet(nn.Module): lora_paths=self.lora_paths, lora_scales=self.lora_scales, controlnet_image_path=controlnet_image_path, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) def save_model(self, base_path: str) -> None: @@ -179,15 +162,5 @@ class Flux1Controlnet(nn.Module): model=self, bits=self.bits, base_path=base_path, - tokenizers=[ - ("clip_tokenizer.tokenizer", "tokenizer"), - ("t5_tokenizer.tokenizer", "tokenizer_2"), - ], - components=[ - ("vae", "vae"), - ("transformer", "transformer"), - ("clip_text_encoder", "text_encoder"), - ("t5_text_encoder", "text_encoder_2"), - ("transformer_controlnet", "transformer_controlnet"), - ], + weight_definition=FluxControlnetWeightDefinition, ) diff --git a/src/mflux/models/flux/variants/controlnet/transformer_controlnet.py b/src/mflux/models/flux/variants/controlnet/transformer_controlnet.py index 39da5a0..5f5b23d 100644 --- a/src/mflux/models/flux/variants/controlnet/transformer_controlnet.py +++ b/src/mflux/models/flux/variants/controlnet/transformer_controlnet.py @@ -1,8 +1,8 @@ import mlx.core as mx from mlx import nn -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.model.flux_transformer.embed_nd import EmbedND from mflux.models.flux.model.flux_transformer.joint_transformer_block import JointTransformerBlock from mflux.models.flux.model.flux_transformer.single_transformer_block import SingleTransformerBlock @@ -32,7 +32,7 @@ class TransformerControlnet(nn.Module): def __call__( self, t: int, - config: RuntimeConfig, + config: Config, hidden_states: mx.array, prompt_embeds: mx.array, pooled_prompt_embeds: mx.array, @@ -80,7 +80,7 @@ class TransformerControlnet(nn.Module): def _apply_single_transformer_block( self, idx: int, - config: RuntimeConfig, + config: Config, block: SingleTransformerBlock, hidden_states: mx.array, encoder_hidden_states: mx.array, @@ -98,7 +98,7 @@ class TransformerControlnet(nn.Module): # 2. Apply controlnet block states = hidden_states[:, encoder_hidden_states.shape[1] :] controlnet_sample = self.controlnet_single_blocks[idx](states) - scaled_controlnet_sample = controlnet_sample * config.config.controlnet_strength + scaled_controlnet_sample = controlnet_sample * config.controlnet_strength controlnet_single_block_samples.append(scaled_controlnet_sample) return hidden_states @@ -106,7 +106,7 @@ class TransformerControlnet(nn.Module): def _apply_joint_transformer_block( self, idx: int, - config: RuntimeConfig, + config: Config, block: JointTransformerBlock, hidden_states: mx.array, encoder_hidden_states: mx.array, @@ -124,7 +124,7 @@ class TransformerControlnet(nn.Module): # 2. Apply controlnet block controlnet_sample = self.controlnet_blocks[idx](hidden_states) - scaled_controlnet_example = controlnet_sample * config.config.controlnet_strength + scaled_controlnet_example = controlnet_sample * config.controlnet_strength controlnet_block_samples.append(scaled_controlnet_example) return encoder_hidden_states, hidden_states diff --git a/src/mflux/models/flux/variants/controlnet/weight_handler_controlnet.py b/src/mflux/models/flux/variants/controlnet/weight_handler_controlnet.py deleted file mode 100644 index 59a3bc3..0000000 --- a/src/mflux/models/flux/variants/controlnet/weight_handler_controlnet.py +++ /dev/null @@ -1,68 +0,0 @@ -import json -from pathlib import Path - -import mlx.core as mx -from mlx.utils import tree_unflatten - -from mflux.models.flux.weights.weight_handler import MetaData -from mflux.models.flux.weights.weight_util import WeightUtil -from mflux.utils.download import snapshot_download - - -class WeightHandlerControlnet: - def __init__(self, meta_data: MetaData, config: dict, controlnet_transformer: dict | None = None): - self.meta_data = meta_data - self.controlnet_transformer = controlnet_transformer - self.config = config - - @staticmethod - def load_controlnet_transformer(controlnet_model: str) -> "WeightHandlerControlnet": - controlnet_path = Path(snapshot_download(repo_id=controlnet_model, allow_patterns=["*.safetensors", "config.json"])) # fmt:off - file = next(controlnet_path.glob("diffusion_pytorch_model.safetensors")) - quantization_level = mx.load(str(file), return_metadata=True)[1].get("quantization_level") - weights = list(mx.load(str(file)).items()) - config = json.load(open(controlnet_path / "config.json")) - - if quantization_level is not None: - return WeightHandlerControlnet( - config=config, - controlnet_transformer=tree_unflatten(weights), - meta_data=MetaData(quantization_level=quantization_level), - ) - - weights = [WeightUtil.reshape_weights(k, v) for k, v in weights] - weights = WeightUtil.flatten(weights) - weights = tree_unflatten(weights) - - # Quantized weights (i.e. ones exported from this project) don't need any post-processing. - if quantization_level is not None: - return WeightHandlerControlnet( - config=config, - controlnet_transformer=weights, - meta_data=MetaData(quantization_level=quantization_level), - ) - - # Reshape and process the huggingface weights - if "transformer_blocks" in weights: - for block in weights["transformer_blocks"]: - 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], - } - - return WeightHandlerControlnet( - config=config, - controlnet_transformer=weights, - meta_data=MetaData(quantization_level=quantization_level), - ) - - def num_transformer_blocks(self) -> int: - return self.config["num_layers"] - - def num_single_transformer_blocks(self) -> int: - return self.config["num_single_layers"] diff --git a/src/mflux/models/flux/variants/depth/depth_util.py b/src/mflux/models/flux/variants/depth/depth_util.py index dae93c1..76a42bc 100644 --- a/src/mflux/models/flux/variants/depth/depth_util.py +++ b/src/mflux/models/flux/variants/depth/depth_util.py @@ -5,10 +5,10 @@ from pathlib import Path import mlx.core as mx import PIL.Image -from mflux.config.runtime_config import RuntimeConfig -from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.models.common.config.config import Config +from mflux.models.depth_pro.model.depth_pro import DepthPro +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_vae.vae import VAE -from mflux.utils.array_util import ArrayUtil from mflux.utils.image_util import ImageUtil logger = logging.getLogger(__name__) @@ -19,7 +19,7 @@ class DepthUtil: def encode_depth_map( vae: VAE, depth_pro: DepthPro, - config: RuntimeConfig, + config: Config, image_path: str | Path | None = None, depth_image_path: str | Path | None = None, ) -> tuple[mx.array, PIL.Image.Image]: @@ -38,7 +38,7 @@ class DepthUtil: ) depth_map_array = ImageUtil.to_array(scaled_depth_map) encoded_depth = vae.encode(depth_map_array) - depth_latents = ArrayUtil.pack_latents(latents=encoded_depth, height=config.height, width=config.width) + depth_latents = FluxLatentCreator.pack_latents(latents=encoded_depth, height=config.height, width=config.width) return depth_latents, depth_image diff --git a/src/mflux/models/flux/variants/depth/flux_depth.py b/src/mflux/models/flux/variants/depth/flux_depth.py index a415058..8a9558b 100644 --- a/src/mflux/models/flux/variants/depth/flux_depth.py +++ b/src/mflux/models/flux/variants/depth/flux_depth.py @@ -1,12 +1,11 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig -from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.depth_pro.model.depth_pro import DepthPro from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -15,7 +14,6 @@ from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5En from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.depth.depth_util import DepthUtil -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -31,29 +29,46 @@ class Flux1Depth(nn.Module): def __init__( self, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.dev_depth(), ): super().__init__() FluxInitializer.init_depth( - flux_model=self, - model_config=ModelConfig.dev_depth(), + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + depth_image_path: Path | str | None = None, + scheduler: str = "linear", ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + depth_image_path=depth_image_path, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = FluxLatentCreator.create_noise( @@ -66,8 +81,8 @@ class Flux1Depth(nn.Module): 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) @@ -81,24 +96,19 @@ class Flux1Depth(nn.Module): depth_image_path=config.depth_image_path, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - depth_image=depth_image, - ) # fmt: off + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents, depth_image=depth_image) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Concatenate the updated latents with the (static) depth map + # 5.t Concatenate the updated latents with the (static) depth map hidden_states = mx.concatenate([latents, static_depth_map], axis=-1) - # 5.t Predict the noise + # 6.t Predict the noise noise = self.transformer( t=t, config=config, @@ -107,47 +117,26 @@ class Flux1Depth(nn.Module): pooled_prompt_embeds=pooled_prompt_embeds, ) - # 6.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 7.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) # fmt: off + # 8.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) # fmt: off + # 9. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -159,5 +148,5 @@ class Flux1Depth(nn.Module): lora_scales=self.lora_scales, image_path=config.image_path, depth_image_path=config.depth_image_path, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) diff --git a/src/mflux/models/flux/variants/dreambooth/dataset/dataset.py b/src/mflux/models/flux/variants/dreambooth/dataset/dataset.py index 3d6a66f..21424f9 100644 --- a/src/mflux/models/flux/variants/dreambooth/dataset/dataset.py +++ b/src/mflux/models/flux/variants/dreambooth/dataset/dataset.py @@ -5,11 +5,11 @@ import PIL.Image from mlx import nn from tqdm import tqdm +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.variants.dreambooth.dataset.batch import Example from mflux.models.flux.variants.dreambooth.dataset.dreambooth_preprocessing import DreamBoothPreProcessing from mflux.models.flux.variants.dreambooth.state.training_spec import ExampleSpec from mflux.models.flux.variants.txt2img.flux import Flux1 -from mflux.utils.array_util import ArrayUtil from mflux.utils.image_util import ImageUtil @@ -50,8 +50,10 @@ class Dataset: encoded_image = Dataset._encode_image(flux.vae, entry.image, width=width, height=height) # Encode the prompt - prompt_embeds = flux.t5_text_encoder(flux.t5_tokenizer.tokenize(entry.prompt)) - pooled_prompt_embeds = flux.clip_text_encoder(flux.clip_tokenizer.tokenize(entry.prompt)) + t5_output = flux.tokenizers["t5"].tokenize(entry.prompt) + clip_output = flux.tokenizers["clip"].tokenize(entry.prompt) + prompt_embeds = flux.t5_text_encoder(t5_output.input_ids) + pooled_prompt_embeds = flux.clip_text_encoder(clip_output.input_ids) # Create the example object example = Example( @@ -65,9 +67,7 @@ class Dataset: examples.append(example) # Evaluate to enable progress tracking - mx.eval(encoded_image) - mx.eval(prompt_embeds) - mx.eval(pooled_prompt_embeds) + mx.eval(encoded_image, prompt_embeds, pooled_prompt_embeds) return examples @@ -76,5 +76,5 @@ class Dataset: image = PIL.Image.open(image_path.resolve()).convert("RGB") scaled_user_image = ImageUtil.scale_to_dimensions(image, target_width=width, target_height=height) encoded = vae.encode(ImageUtil.to_array(scaled_user_image)) - latents = ArrayUtil.pack_latents(encoded, width=width, height=height) + latents = FluxLatentCreator.pack_latents(encoded, width=width, height=height) return latents diff --git a/src/mflux/models/flux/variants/dreambooth/dataset/iterator.py b/src/mflux/models/flux/variants/dreambooth/dataset/iterator.py index c605ba3..0af5d4f 100644 --- a/src/mflux/models/flux/variants/dreambooth/dataset/iterator.py +++ b/src/mflux/models/flux/variants/dreambooth/dataset/iterator.py @@ -102,7 +102,6 @@ class Iterator: return Batch(examples=examples, rng=self.rng) def to_dict(self) -> dict[str, Any]: - """Returns the current state as a dictionary.""" return { "position": self._position, "epoch": self._epoch, @@ -116,7 +115,6 @@ class Iterator: @classmethod def from_dict(cls, state_dict: dict[str, Any], dataset: Dataset) -> "Iterator": - """Creates an iterator from a state dictionary.""" # Convert the RNG state elements to tuples rng_state = state_dict["rng_state"] rng_state = (rng_state[0], tuple(rng_state[1]), rng_state[2]) @@ -134,7 +132,6 @@ class Iterator: ) def to_json(self, filepath: str) -> None: - """Save iterator state to a JSON file.""" with open(filepath, "w") as f: json.dump(self.to_dict(), f, indent=4) diff --git a/src/mflux/models/flux/variants/dreambooth/dreambooth.py b/src/mflux/models/flux/variants/dreambooth/dreambooth.py index 7ba51a3..3cb0413 100644 --- a/src/mflux/models/flux/variants/dreambooth/dreambooth.py +++ b/src/mflux/models/flux/variants/dreambooth/dreambooth.py @@ -1,34 +1,33 @@ from mlx import nn from tqdm import tqdm -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear from mflux.models.flux.variants.dreambooth.optimization.dreambooth_loss import DreamBoothLoss from mflux.models.flux.variants.dreambooth.state.training_spec import TrainingSpec from mflux.models.flux.variants.dreambooth.state.training_state import TrainingState from mflux.models.flux.variants.dreambooth.statistics.plotter import Plotter from mflux.models.flux.variants.txt2img.flux import Flux1 -from mflux.models.flux.weights.weight_handler_lora import WeightHandlerLoRA class DreamBooth: @staticmethod def train( flux: Flux1, - runtime_config: RuntimeConfig, + config: Config, training_spec: TrainingSpec, training_state: TrainingState, ): - # Freeze the model and assign the LoRA layers to the model + # Freeze the model (LoRA layers are already applied to transformer in from_spec) flux.freeze() - WeightHandlerLoRA.set_lora_layers( - transformer_module=flux.transformer, - lora_layers=training_state.lora_layers, - ) + + # Unfreeze LoRA layers so they can be trained + DreamBooth._unfreeze_lora_layers(flux.transformer) # Define loss computation as a function of a batch 'b' train_step_function = nn.value_and_grad( model=flux, - fn=lambda b: DreamBoothLoss.compute_loss(flux, runtime_config, b), + fn=lambda b: DreamBoothLoss.compute_loss(flux, config, b), ) # Setup progress bar @@ -48,7 +47,7 @@ class DreamBooth: # Plot loss progress periodically if training_state.should_plot_loss(training_spec): validation_batch = training_state.iterator.get_validation_batch() - validation_loss = DreamBoothLoss.compute_loss(flux, runtime_config, validation_batch) + validation_loss = DreamBoothLoss.compute_loss(flux, config, validation_batch) training_state.statistics.append_values(step=training_state.iterator.num_iterations, loss=validation_loss) # fmt: off Plotter.update_loss_plot(training_spec=training_spec, training_state=training_state) del validation_loss @@ -57,7 +56,6 @@ class DreamBooth: if training_state.should_generate_image(training_spec): image = flux.generate_image( seed=training_spec.seed, - config=runtime_config.config, prompt=training_spec.instrumentation.validation_prompt, ) image.save(path=training_state.get_current_validation_image_path(training_spec)) @@ -66,7 +64,13 @@ class DreamBooth: # Save checkpoint periodically if training_state.should_save(training_spec): - training_state.save(training_spec) + training_state.save(flux, training_spec) # Save the final state - training_state.save(training_spec) + training_state.save(flux, training_spec) + + @staticmethod + def _unfreeze_lora_layers(module: nn.Module) -> None: + for name, child in module.named_modules(): + if isinstance(child, LoRALinear): + child.unfreeze(keys=["lora_A", "lora_B"], strict=False) diff --git a/src/mflux/models/flux/variants/dreambooth/dreambooth_initializer.py b/src/mflux/models/flux/variants/dreambooth/dreambooth_initializer.py index d8c1e5f..5bc471f 100644 --- a/src/mflux/models/flux/variants/dreambooth/dreambooth_initializer.py +++ b/src/mflux/models/flux/variants/dreambooth/dreambooth_initializer.py @@ -1,8 +1,7 @@ import mlx.core.random as random -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.variants.dreambooth.dataset.dataset import Dataset from mflux.models.flux.variants.dreambooth.dataset.iterator import Iterator from mflux.models.flux.variants.dreambooth.lora_layers.lora_layers import LoRALayers @@ -18,7 +17,7 @@ class DreamBoothInitializer: def initialize( config_path: str | None, checkpoint_path: str | None, - ) -> tuple[Flux1, RuntimeConfig, TrainingSpec, TrainingState]: + ) -> tuple[Flux1, Config, TrainingSpec, TrainingState]: # The training specification describing the details of the training process. It is resolved # differently depending on if training starts from scratch or resumes from checkpoint. training_spec = TrainingSpec.resolve( @@ -35,21 +34,19 @@ class DreamBoothInitializer: model_config=model_config, quantize=training_spec.quantize, ) - runtime_config = RuntimeConfig( + config = Config( model_config=model_config, - config=Config( - num_inference_steps=training_spec.steps, - width=training_spec.width, - height=training_spec.height, - guidance=training_spec.guidance, - ), + num_inference_steps=training_spec.steps, + width=training_spec.width, + height=training_spec.height, + guidance=training_spec.guidance, ) # Create the optimizer optimizer = Optimizer.from_spec(training_spec) - # Create the LoRA layers by matching them against the corresponding Flux layers - lora_layers = LoRALayers.from_spec(flux=flux, training_spec=training_spec) + # Create and apply the LoRA layers directly to the transformer + LoRALayers.from_spec(flux=flux, training_spec=training_spec) # Prepare the fine-tuning dataset and create the iterator dataset = Dataset.prepare_dataset( @@ -69,9 +66,8 @@ class DreamBoothInitializer: # The training state consisting of everything that moves during training training_state = TrainingState( optimizer=optimizer, - lora_layers=lora_layers, iterator=iterator, statistics=statistics, ) - return flux, runtime_config, training_spec, training_state + return flux, config, training_spec, training_state diff --git a/src/mflux/models/flux/variants/dreambooth/lora_layers/lora_layers.py b/src/mflux/models/flux/variants/dreambooth/lora_layers/lora_layers.py index 6628f93..e9b8f3d 100644 --- a/src/mflux/models/flux/variants/dreambooth/lora_layers/lora_layers.py +++ b/src/mflux/models/flux/variants/dreambooth/lora_layers/lora_layers.py @@ -1,12 +1,11 @@ from pathlib import Path from typing import TYPE_CHECKING -import mlx import mlx.core as mx -from mlx import nn from mlx.utils import tree_flatten from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear +from mflux.models.common.lora.mapping.lora_loader import LoRALoader from mflux.models.flux.model.flux_transformer.joint_transformer_block import JointTransformerBlock from mflux.models.flux.model.flux_transformer.single_transformer_block import SingleTransformerBlock from mflux.models.flux.variants.dreambooth.state.training_spec import ( @@ -15,7 +14,7 @@ from mflux.models.flux.variants.dreambooth.state.training_spec import ( TransformerBlocks, ) from mflux.models.flux.variants.dreambooth.state.zip_util import ZipUtil -from mflux.models.flux.weights.weight_handler import MetaData, WeightHandler +from mflux.models.flux.weights.flux_lora_mapping import FluxLoRAMapping from mflux.utils.version_util import VersionUtil if TYPE_CHECKING: @@ -23,59 +22,46 @@ if TYPE_CHECKING: class LoRALayers: - def __init__(self, weights: "WeightHandler"): - self.layers = weights - @staticmethod - def from_spec(flux: "Flux1", training_spec: TrainingSpec) -> "LoRALayers": + def from_spec(flux: "Flux1", training_spec: TrainingSpec) -> None: if training_spec.lora_layers.state_path is not None: - # Load from state if present in the spec - from mflux.models.flux.weights.weight_handler_lora import WeightHandlerLoRA - - weights = ZipUtil.unzip( + # Load from checkpoint using the unified LoRALoader + ZipUtil.unzip( zip_path=training_spec.checkpoint_path, filename=training_spec.lora_layers.state_path, - loader=lambda x: WeightHandlerLoRA.load_lora_weights( - transformer=flux.transformer, lora_files=[x], lora_scales=[1.0] + loader=lambda lora_file: LoRALoader.load_and_apply_lora( + lora_mapping=FluxLoRAMapping.get_mapping(), + transformer=flux.transformer, + lora_paths=[lora_file], + lora_scales=[1.0], ), ) - return LoRALayers(weights=weights[0]) else: - # Construct the LoRA weights from the spec - transformer_lora_layers = {} - single_transformer_lora_layers = {} - + # Construct fresh LoRA layers and apply directly to transformer if training_spec.lora_layers.transformer_blocks: - transformer_lora_layers = LoRALayers._construct_layers( + LoRALayers._construct_and_apply_layers( + transformer=flux.transformer, blocks=flux.transformer.transformer_blocks, block_spec=training_spec.lora_layers.transformer_blocks, - block_prefix="transformer.transformer_blocks", + block_attr="transformer_blocks", ) if training_spec.lora_layers.single_transformer_blocks: - single_transformer_lora_layers = LoRALayers._construct_layers( + LoRALayers._construct_and_apply_layers( + transformer=flux.transformer, blocks=flux.transformer.single_transformer_blocks, block_spec=training_spec.lora_layers.single_transformer_blocks, - block_prefix="transformer.single_transformer_blocks", + block_attr="single_transformer_blocks", ) - lora_layers = {**transformer_lora_layers, **single_transformer_lora_layers} - - weights = WeightHandler( - meta_data=MetaData(mflux_version=VersionUtil.get_mflux_version()), - transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))["transformer"], - ) - - return LoRALayers(weights=weights) - @staticmethod - def _construct_layers( - block_spec: TransformerBlocks | SingleTransformerBlocks, + def _construct_and_apply_layers( + transformer, blocks: list[JointTransformerBlock] | list[SingleTransformerBlock], - block_prefix: str, - ) -> dict: + block_spec: TransformerBlocks | SingleTransformerBlocks, + block_attr: str, + ) -> None: block_indices = block_spec.block_range.get_blocks() - lora_layers = {} for idx in block_indices: if idx >= len(blocks): raise IndexError(f"Index {idx} over range") @@ -89,192 +75,40 @@ class LoRALayers: linear=original_layer[0] if is_list else original_layer, r=block_spec.lora_rank, ) - layer_path = f"{block_prefix}.{idx}.{layer_type}" - lora_layers[layer_path] = [lora_layer] if is_list else lora_layer - - return lora_layers + # Apply the LoRA layer directly to the transformer + LoRALayers._set_layer_at_path( + transformer=transformer, + block_attr=block_attr, + block_idx=idx, + layer_type=layer_type, + lora_layer=lora_layer, + is_list=is_list, + ) @staticmethod - def transformer_dict_from_template(weights: dict, transformer: nn.Module, scale: float) -> dict: - lora_layers = {} - for key in weights.keys(): - if key.endswith(".lora_A"): - base_path = key[: -len(".lora_A")] - parts = base_path.split(".") + def _set_layer_at_path( + transformer, + block_attr: str, + block_idx: int, + layer_type: str, + lora_layer: LoRALinear, + is_list: bool, + ) -> None: + block = getattr(transformer, block_attr)[block_idx] + parts = layer_type.split(".") - if parts[1] == "transformer_blocks": - LoRALayers._handle_transformer_blocks( - weights=weights, - scale=scale, - transformer=transformer, - lora_layers=lora_layers, - base_path=base_path, - ) + # Navigate to parent + current = block + for part in parts[:-1]: + current = getattr(current, part) - elif parts[1] == "single_transformer_blocks": - LoRALayers._handle_single_transformer_blocks( - weights=weights, - scale=scale, - transformer=transformer, - lora_layers=lora_layers, - 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 - get the attribute and check if it's actually a list - original_layer = getattr(module, attr_name) - - # Some layers (like attn.to_out) are legitimately lists in the transformer architecture - # In such cases, we need to extract the first element for LoRA creation - if isinstance(original_layer, list): - # For list layers, we use the first element for LoRA creation - # but we need to return the list itself for proper storage - return original_layer, ".".join(parts) - - 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] - - # Handle the case where original_layer is a list (like attn.to_out) - # In such cases, use the first element for LoRA creation (same as construction code) - is_list = isinstance(original_layer, list) - layer_for_lora = original_layer[0] if is_list else original_layer - - lora_layer = LoRALinear.from_linear(linear=layer_for_lora, 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 - ): - parts = base_path.split(".") - block_idx = int(parts[2]) - module_name = parts[3] - attr_name = parts[4] - block = transformer.transformer_blocks[block_idx] - module = getattr(block, module_name) - - # 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 and store LoRA layer - lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale) - - # Store as list if original layer was a list (matching construction logic) - is_list = isinstance(original_layer, list) - lora_layers[storage_path] = [lora_layer] if is_list else lora_layer - - @staticmethod - def _handle_single_transformer_blocks( - weights: dict, scale: float, transformer: nn.Module, lora_layers: dict, base_path: str - ): - 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) + # Set the final attribute + final_attr = parts[-1] + if is_list: + getattr(current, final_attr)[0] = lora_layer else: - attr_name = parts[4] - block = transformer.single_transformer_blocks[block_idx] - module = getattr(block, module_name) - original_layer = getattr(module, attr_name) - - # Create and store LoRA layer - lora_layer = LoRALayers._create_lora_layer(weights, base_path, original_layer, scale) - - # Store as list if original layer was a list (matching construction logic) - is_list = isinstance(original_layer, list) - lora_layers[base_path] = [lora_layer] if is_list else lora_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) - - # Store as list if original layer was a list (matching construction logic) - is_list = isinstance(original_layer, list) - lora_layers[base_path] = [lora_layer] if is_list else lora_layer - - @staticmethod - def set_transformer_block(transformer_block, dictionary: dict): - for key, val in dictionary.items(): - if key == "attn": - LoRALayers._set_attribute(transformer_block, key, val, "to_q") - LoRALayers._set_attribute(transformer_block, key, val, "to_k") - LoRALayers._set_attribute(transformer_block, key, val, "to_v") - LoRALayers._set_attribute(transformer_block, key, val, "to_out") - LoRALayers._set_attribute(transformer_block, key, val, "add_q_proj") - LoRALayers._set_attribute(transformer_block, key, val, "add_k_proj") - LoRALayers._set_attribute(transformer_block, key, val, "add_v_proj") - LoRALayers._set_attribute(transformer_block, key, val, "to_add_out") - elif key == "ff" or key == "ff_context": - LoRALayers._set_attribute(transformer_block, key, val, "linear1") - LoRALayers._set_attribute(transformer_block, key, val, "linear2") - elif key == "norm1" or key == "norm1_context": - LoRALayers._set_attribute(transformer_block, key, val, "linear") - else: - raise Exception("Could not set LoRA weights") - - @staticmethod - def set_single_transformer_block(single_transformer_block, dictionary: dict): - for key, val in dictionary.items(): - if key == "attn": - LoRALayers._set_attribute(single_transformer_block, key, val, "to_q") - LoRALayers._set_attribute(single_transformer_block, key, val, "to_k") - LoRALayers._set_attribute(single_transformer_block, key, val, "to_v") - elif key == "norm": - LoRALayers._set_attribute(single_transformer_block, key, val, "linear") - elif key == "proj_mlp" or key == "proj_out": - single_transformer_block[key] = val - else: - raise Exception("Could not set LoRA weights") - - @staticmethod - def _set_attribute(block, key: str, val: dict, name: str): - if block[key].get(name, False) and val.get(name, False): - if name == "to_out" and isinstance(block[key][name], list): - if len(block[key][name]) > 0: - lora_layer = val[name][0] if isinstance(val[name], list) else val[name] - block[key][name][0] = lora_layer - else: - block[key][name] = val[name] + setattr(current, final_attr, lora_layer) @staticmethod def _get_nested_attr(obj, attr_path): @@ -283,19 +117,19 @@ class LoRALayers: obj = getattr(obj, attr) return obj - def save(self, path: Path, training_spec: TrainingSpec) -> None: + @staticmethod + def save(transformer, path: Path, training_spec: TrainingSpec) -> None: weights = {} - for entry in tree_flatten(self.layers.transformer): + for entry in tree_flatten(transformer): name = entry[0] weight = entry[1] if name.endswith(".lora_A") or name.endswith(".lora_B"): weights[name] = weight - weights = {key: mx.transpose(val) for key, val in weights.items()} - weights = {"transformer": weights} + weights = {f"transformer.{key}": mx.transpose(val) for key, val in weights.items()} mx.save_safetensors( str(path), - dict(tree_flatten(weights)), + weights, metadata={ "mflux_version": VersionUtil.get_mflux_version(), "transformer_blocks": str(training_spec.lora_layers.transformer_blocks), diff --git a/src/mflux/models/flux/variants/dreambooth/optimization/dreambooth_loss.py b/src/mflux/models/flux/variants/dreambooth/optimization/dreambooth_loss.py index b2b49a4..8631a5b 100644 --- a/src/mflux/models/flux/variants/dreambooth/optimization/dreambooth_loss.py +++ b/src/mflux/models/flux/variants/dreambooth/optimization/dreambooth_loss.py @@ -2,8 +2,8 @@ import random import mlx.core as mx -from mflux.config.config import Config -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.latent_creator.latent_creator import LatentCreator from mflux.models.flux.variants.dreambooth.dataset.batch import Batch, Example from mflux.models.flux.variants.txt2img.flux import Flux1 @@ -11,7 +11,7 @@ from mflux.models.flux.variants.txt2img.flux import Flux1 class DreamBoothLoss: @staticmethod - def compute_loss(flux: Flux1, config: RuntimeConfig, batch: Batch) -> mx.float16: + def compute_loss(flux: Flux1, config: Config, batch: Batch) -> mx.float16: losses = [ DreamBoothLoss._single_example_loss(flux, config, example, batch.rng) for example in batch.examples @@ -19,7 +19,7 @@ class DreamBoothLoss: return mx.mean(mx.array(losses)) @staticmethod - def _single_example_loss(flux: Flux1, config: RuntimeConfig, example: Example, rng: random.Random) -> mx.float16: + def _single_example_loss(flux: Flux1, config: Config, example: Example, rng: random.Random) -> mx.float16: # Must be a better way to handle the randomness than this, but we already # save/restore the random state via the iterator so this is a continent shortcut. time_seed = rng.randint(0, 2**32 - 1) @@ -41,7 +41,7 @@ class DreamBoothLoss: # Generate pure noise pure_noise = mx.random.normal( shape=clean_image.shape, - dtype=Config.precision, + dtype=ModelConfig.precision, key=mx.random.key(noise_seed), ) diff --git a/src/mflux/models/flux/variants/dreambooth/state/training_state.py b/src/mflux/models/flux/variants/dreambooth/state/training_state.py index 824a298..b2f2660 100644 --- a/src/mflux/models/flux/variants/dreambooth/state/training_state.py +++ b/src/mflux/models/flux/variants/dreambooth/state/training_state.py @@ -3,6 +3,7 @@ import json import tempfile import zipfile from pathlib import Path +from typing import TYPE_CHECKING from mflux.models.flux.variants.dreambooth.dataset.iterator import Iterator from mflux.models.flux.variants.dreambooth.lora_layers.lora_layers import LoRALayers @@ -11,6 +12,9 @@ from mflux.models.flux.variants.dreambooth.state.training_spec import TrainingSp from mflux.models.flux.variants.dreambooth.state.zip_util import ZipUtil from mflux.models.flux.variants.dreambooth.statistics.statistics import Statistics +if TYPE_CHECKING: + from mflux.models.flux.variants.txt2img.flux import Flux1 + DREAMBOOTH_PATH_CHECKPOINTS = "_checkpoints" DREAMBOOTH_FILE_NAME_CHECKPOINT = "checkpoint" @@ -31,16 +35,14 @@ class TrainingState: def __init__( self, iterator: Iterator, - lora_layers: LoRALayers, optimizer: Optimizer, statistics: Statistics, ): self.iterator = iterator self.optimizer = optimizer - self.lora_layers = lora_layers self.statistics = statistics - def save(self, training_spec: TrainingSpec) -> None: + def save(self, flux: "Flux1", training_spec: TrainingSpec) -> None: # Create a temporary directory to store files before zipping with tempfile.TemporaryDirectory() as temp_dir: # Save files to temporary directory @@ -54,7 +56,7 @@ class TrainingState: # Save individual files to temporary directory self.optimizer.save(optimizer_path) - self.lora_layers.save(lora_path, training_spec) + LoRALayers.save(flux.transformer, lora_path, training_spec) self.iterator.save(iterator_path) self.statistics.save(loss_path) self._save_train_config(config_path, training_spec) # For completeness, we also save the (initial) config diff --git a/src/mflux/models/flux/variants/fill/flux_fill.py b/src/mflux/models/flux/variants/fill/flux_fill.py index 261c8b4..05b4bb2 100644 --- a/src/mflux/models/flux/variants/fill/flux_fill.py +++ b/src/mflux/models/flux/variants/fill/flux_fill.py @@ -1,11 +1,10 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -14,7 +13,6 @@ from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5En from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.fill.mask_util import MaskUtil -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -29,43 +27,60 @@ class Flux1Fill(nn.Module): def __init__( self, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.dev_fill(), ): super().__init__() FluxInitializer.init( - flux_model=self, - model_config=ModelConfig.dev_fill(), + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + image_path: Path | str, + masked_image_path: Path | str, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_strength: float | None = None, + scheduler: str = "linear", ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + masked_image_path=masked_image_path, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = FluxLatentCreator.create_noise( seed=seed, - height=config.height, width=config.width, + height=config.height, ) # 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) @@ -73,29 +88,25 @@ class Flux1Fill(nn.Module): # 3. Create the static masked latents static_masked_latents = MaskUtil.create_masked_latents( vae=self.vae, - height=config.height, width=config.width, + height=config.height, img_path=config.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, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Concatenate the updated latents with the static masked latents + # 5.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 + # 6.t Predict the noise noise = self.transformer( t=t, config=config, @@ -104,47 +115,26 @@ class Flux1Fill(nn.Module): pooled_prompt_embeds=pooled_prompt_embeds, ) - # 6.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 7.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 8.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 9. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -157,5 +147,5 @@ class Flux1Fill(nn.Module): image_path=config.image_path, image_strength=config.image_strength, masked_image_path=config.masked_image_path, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) diff --git a/src/mflux/models/flux/variants/fill/mask_util.py b/src/mflux/models/flux/variants/fill/mask_util.py index 90da69e..b872598 100644 --- a/src/mflux/models/flux/variants/fill/mask_util.py +++ b/src/mflux/models/flux/variants/fill/mask_util.py @@ -2,8 +2,8 @@ from pathlib import Path import mlx.core as mx +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_vae.vae import VAE -from mflux.utils.array_util import ArrayUtil from mflux.utils.image_util import ImageUtil @@ -39,11 +39,11 @@ class MaskUtil: # 3. Create and pack the masked image masked_image = image * (1 - the_mask) masked_image = vae.encode(masked_image) - masked_image = ArrayUtil.pack_latents(latents=masked_image, height=height, width=width) + masked_image = FluxLatentCreator.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 = ArrayUtil.pack_latents(latents=mask, height=height, width=width, num_channels_latents=64) + mask = FluxLatentCreator.pack_latents(latents=mask, height=height, width=width, num_channels_latents=64) # 5. Concat the masked_image and the mask masked_image_latents = mx.concatenate([masked_image, mask], axis=-1) diff --git a/src/mflux/models/flux/variants/in_context/flux_in_context_dev.py b/src/mflux/models/flux/variants/in_context/flux_in_context_dev.py index 03a7f27..3c3fe28 100644 --- a/src/mflux/models/flux/variants/in_context/flux_in_context_dev.py +++ b/src/mflux/models/flux/variants/in_context/flux_in_context_dev.py @@ -1,19 +1,18 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.latent_creator.latent_creator import LatentCreator from mflux.models.flux.flux_initializer import FluxInitializer +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -27,42 +26,52 @@ class Flux1InContextDev(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, + model_config: ModelConfig = ModelConfig.dev(), ): super().__init__() FluxInitializer.init( - flux_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Encode the reference image encoded_image = LatentCreator.encode_image( vae=self.vae, - image_path=config.image_path, - height=config.height, width=config.width, + height=config.height, + image_path=config.image_path, ) # 2. Create the initial latents and keep the initial static noise for later blending @@ -73,26 +82,22 @@ class Flux1InContextDev(nn.Module): 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Predict the noise + # 5.t Predict the noise noise = self.transformer( t=t, config=config, @@ -101,14 +106,10 @@ class Flux1InContextDev(nn.Module): pooled_prompt_embeds=pooled_prompt_embeds, ) - # 5.t Take one denoise step and update latents - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 6.t Take one denoise step and update latents + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # 6.t Override the left-hand side of latents by linearly interpolating between latents and static noise + # 7.t Override the left-hand side of latents by linearly interpolating between latents and static noise latents = Flux1InContextDev._update_latents( t=t, config=config, @@ -118,40 +119,23 @@ class Flux1InContextDev(nn.Module): sigmas=config.scheduler.sigmas, ) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 8.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 9. Call subscribers after loop + ctx.after_loop(latents) - # 6. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -163,11 +147,11 @@ class Flux1InContextDev(nn.Module): lora_scales=self.lora_scales, image_path=config.image_path, image_strength=config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) @staticmethod - def _create_in_context_latents(seed: int, config: RuntimeConfig): + def _create_in_context_latents(seed: int, config: Config): # 1. Double the width for side-by-side generation config.width = 2 * config.width @@ -177,21 +161,23 @@ class Flux1InContextDev(nn.Module): # 3. Create noise with appropriate dimensions static_noise = mx.random.normal(shape=[1, 16, latent_height, latent_width], key=mx.random.key(seed)) - latents = ArrayUtil.pack_latents(latents=static_noise, height=config.height, width=config.width) + latents = FluxLatentCreator.pack_latents(latents=static_noise, height=config.height, width=config.width) return latents @staticmethod def _update_latents( t: int, - config: RuntimeConfig, + config: Config, latents: mx.array, encoded_image: mx.array, static_noise: mx.array, sigmas: mx.array, ) -> mx.array: # 1. Unpack the latents - unpacked = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) - unpacked_static_noise = ArrayUtil.unpack_latents(latents=static_noise, height=config.height, width=config.width) + unpacked = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) + unpacked_static_noise = FluxLatentCreator.unpack_latents( + latents=static_noise, height=config.height, width=config.width + ) # 2. Calculate latent_width from the config (original width is half of current width) latent_width = (config.width // 2) // 8 @@ -204,4 +190,4 @@ class Flux1InContextDev(nn.Module): ) # 4. Repack the latents - return ArrayUtil.pack_latents(latents=unpacked, height=config.height, width=config.width) + return FluxLatentCreator.pack_latents(latents=unpacked, height=config.height, width=config.width) diff --git a/src/mflux/models/flux/variants/in_context/flux_in_context_fill.py b/src/mflux/models/flux/variants/in_context/flux_in_context_fill.py index 1655fb0..18dc333 100644 --- a/src/mflux/models/flux/variants/in_context/flux_in_context_fill.py +++ b/src/mflux/models/flux/variants/in_context/flux_in_context_fill.py @@ -1,11 +1,10 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -14,7 +13,6 @@ from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5En from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.in_context.utils.in_context_mask_util import InContextMaskUtil -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -28,38 +26,51 @@ class Flux1InContextFill(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.dev(), ): super().__init__() FluxInitializer.init( - flux_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, left_image_path: str, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, right_image_path: str | None = None, + masked_image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", ) -> 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 + original_width = width + doubled_width = original_width * 2 - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + height=height, + width=doubled_width, # Use doubled width for in-context + guidance=guidance, + scheduler=scheduler, + image_strength=image_strength, + model_config=self.model_config, + masked_image_path=masked_image_path, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = FluxLatentCreator.create_noise( @@ -72,8 +83,8 @@ class Flux1InContextFill(nn.Module): 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) @@ -82,30 +93,26 @@ class Flux1InContextFill(nn.Module): static_masked_latents = InContextMaskUtil.create_masked_latents( vae=self.vae, height=config.height, - width=config.width, + width=doubled_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, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Concatenate the updated latents with the static masked latents + # 5.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 + # 6.t Predict the noise noise = self.transformer( t=t, config=config, @@ -114,47 +121,26 @@ class Flux1InContextFill(nn.Module): pooled_prompt_embeds=pooled_prompt_embeds, ) - # 6.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 7.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 8.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 9. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the full image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 10. Decode the latent array and return the full image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -167,5 +153,5 @@ class Flux1InContextFill(nn.Module): image_path=right_image_path, image_strength=config.image_strength, masked_image_path=config.masked_image_path, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) diff --git a/src/mflux/models/flux/variants/in_context/utils/in_context_fill_util.py b/src/mflux/models/flux/variants/in_context/utils/in_context_fill_util.py new file mode 100644 index 0000000..0a09fea --- /dev/null +++ b/src/mflux/models/flux/variants/in_context/utils/in_context_fill_util.py @@ -0,0 +1,28 @@ +from PIL import Image + +from mflux.utils.prompt_util import PromptUtil + + +class FluxInContextFillUtil: + @staticmethod + 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 PromptUtil.read_prompt(args) + + @staticmethod + 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 diff --git a/src/mflux/models/flux/variants/in_context/utils/in_context_loras.py b/src/mflux/models/flux/variants/in_context/utils/in_context_loras.py index f82351b..19b8bc5 100644 --- a/src/mflux/models/flux/variants/in_context/utils/in_context_loras.py +++ b/src/mflux/models/flux/variants/in_context/utils/in_context_loras.py @@ -1,6 +1,3 @@ -"""Configuration for In-Context LoRAs from Hugging Face.""" -from mflux.models.common.lora.download.lora_huggingface_downloader import LoRAHuggingFaceDownloader - # Default Hugging Face repository for In-Context LoRAs LORA_REPO_ID = "ali-vilab/In-Context-LoRA" @@ -24,31 +21,12 @@ IC_EDIT_LORA_FILENAME = "pytorch_lora_weights.safetensors" IC_EDIT_LORA_SCALE = 1.0 -def get_lora_filename(simple_name: str) -> str: +def get_lora_path(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] + return f"{LORA_REPO_ID}:{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 = LoRAHuggingFaceDownloader.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 +def get_ic_edit_lora_path() -> str: + return f"{IC_EDIT_LORA_REPO_ID}:{IC_EDIT_LORA_FILENAME}" diff --git a/src/mflux/models/flux/variants/in_context/utils/in_context_mask_util.py b/src/mflux/models/flux/variants/in_context/utils/in_context_mask_util.py index 2484d84..3fb6266 100644 --- a/src/mflux/models/flux/variants/in_context/utils/in_context_mask_util.py +++ b/src/mflux/models/flux/variants/in_context/utils/in_context_mask_util.py @@ -1,8 +1,8 @@ import mlx.core as mx +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.variants.fill.mask_util import MaskUtil -from mflux.utils.array_util import ArrayUtil from mflux.utils.image_util import ImageUtil @@ -67,7 +67,7 @@ class InContextMaskUtil: # Encode and process exactly like MaskUtil.create_masked_latents encoded_image = vae.encode(masked_concatenated_image) - encoded_image = ArrayUtil.pack_latents( + encoded_image = FluxLatentCreator.pack_latents( latents=encoded_image, height=height, width=width, @@ -78,7 +78,7 @@ class InContextMaskUtil: height=height, width=width, ) - processed_mask = ArrayUtil.pack_latents( + processed_mask = FluxLatentCreator.pack_latents( latents=processed_mask, height=height, width=width, diff --git a/src/mflux/models/flux/variants/kontext/flux_kontext.py b/src/mflux/models/flux/variants/kontext/flux_kontext.py index 9827e12..56db271 100644 --- a/src/mflux/models/flux/variants/kontext/flux_kontext.py +++ b/src/mflux/models/flux/variants/kontext/flux_kontext.py @@ -1,11 +1,10 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -13,8 +12,7 @@ from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncod from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE -from mflux.models.flux.variants.kontext.utils.kontext_util import KontextUtil -from mflux.utils.array_util import ArrayUtil +from mflux.models.flux.variants.kontext.kontext_util import KontextUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -29,43 +27,58 @@ class Flux1Kontext(nn.Module): def __init__( self, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.dev_kontext(), ): super().__init__() FluxInitializer.init( - flux_model=self, - model_config=ModelConfig.dev_kontext(), + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = FluxLatentCreator.create_noise( seed=seed, - height=config.height, width=config.width, + height=config.height, ) # 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) @@ -73,28 +86,24 @@ class Flux1Kontext(nn.Module): # 3. Create the static image conditioning latents and IDs static_image_latents, kontext_image_ids = KontextUtil.create_image_conditioning_latents( vae=self.vae, - height=config.height, width=config.width, + height=config.height, image_path=config.image_path, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 4.t Concatenate the updated latents with the static image latents + # 5.t Concatenate the updated latents with the static image latents hidden_states = mx.concatenate([latents, static_image_latents], axis=1) - # 5.t Predict the noise + # 6.t Predict the noise noise = self.transformer( t=t, config=config, @@ -104,50 +113,29 @@ class Flux1Kontext(nn.Module): kontext_image_ids=kontext_image_ids, ) - # 6.t Extract only the noise for the generation latents (first part) + # 7.t Extract only the noise for the generation latents (first part) noise = noise[:, : latents.shape[1]] - # 7.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 8.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 9.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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, - ) + # 10. Call subscribers after loop + ctx.after_loop(latents) - # 8. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 11. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -158,5 +146,5 @@ class Flux1Kontext(nn.Module): lora_paths=self.lora_paths, lora_scales=self.lora_scales, image_path=config.image_path, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) diff --git a/src/mflux/models/flux/variants/kontext/utils/kontext_util.py b/src/mflux/models/flux/variants/kontext/kontext_util.py similarity index 94% rename from src/mflux/models/flux/variants/kontext/utils/kontext_util.py rename to src/mflux/models/flux/variants/kontext/kontext_util.py index fd67393..f87b930 100644 --- a/src/mflux/models/flux/variants/kontext/utils/kontext_util.py +++ b/src/mflux/models/flux/variants/kontext/kontext_util.py @@ -1,7 +1,7 @@ import mlx.core as mx from mflux.models.common.latent_creator.latent_creator import LatentCreator -from mflux.utils.array_util import ArrayUtil +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator class KontextUtil: @@ -21,7 +21,7 @@ class KontextUtil: ) # Pack image latents for conditioning - image_latents = ArrayUtil.pack_latents( + image_latents = FluxLatentCreator.pack_latents( latents=input_image, height=height, width=width, diff --git a/src/mflux/models/flux/variants/redux/flux_redux.py b/src/mflux/models/flux/variants/redux/flux_redux.py index d590a4f..9042fd3 100644 --- a/src/mflux/models/flux/variants/redux/flux_redux.py +++ b/src/mflux/models/flux/variants/redux/flux_redux.py @@ -2,12 +2,10 @@ from pathlib import Path import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.common.tokenizer import Tokenizer from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -17,10 +15,7 @@ from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE from mflux.models.flux.model.redux_encoder.redux_encoder import ReduxEncoder from mflux.models.flux.model.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer -from mflux.models.flux.tokenizer.clip_tokenizer import TokenizerCLIP -from mflux.models.flux.tokenizer.t5_tokenizer import TokenizerT5 from mflux.models.flux.variants.redux.redux_util import ReduxUtil -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -36,136 +31,128 @@ class Flux1Redux(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.dev(), ): super().__init__() FluxInitializer.init_redux( - flux_model=self, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + redux_image_paths: list[Path | str], + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + redux_image_strengths: list[float] | None = None, + image_strength: float | None = None, + scheduler: str = "linear", ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - runtime_config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(runtime_config.init_time_step, runtime_config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_strength=image_strength, + model_config=self.model_config, + redux_image_paths=redux_image_paths, + num_inference_steps=num_inference_steps, + redux_image_strengths=redux_image_strengths, + ) # 1. Create the initial latents latents = FluxLatentCreator.create_noise( seed=seed, - height=runtime_config.height, - width=runtime_config.width, + width=config.width, + height=config.height, ) # 2. Get prompt embeddings by fusing the prompt and image embeddings prompt_embeds, pooled_prompt_embeds = Flux1Redux._get_prompt_embeddings( prompt=prompt, prompt_cache=self.prompt_cache, - t5_tokenizer=self.t5_tokenizer, - clip_tokenizer=self.clip_tokenizer, + t5_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, - image_paths=runtime_config.redux_image_paths, + image_paths=config.redux_image_paths, image_encoder=self.image_encoder, image_embedder=self.image_embedder, - image_strengths=runtime_config.redux_image_strengths, - ) # fmt: off + image_strengths=config.redux_image_strengths, + ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - ) # fmt: off + # 3. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler - latents = runtime_config.scheduler.scale_model_input(latents, t) + latents = config.scheduler.scale_model_input(latents, t) - # 3.t Predict the noise + # 4.t Predict the noise noise = self.transformer( t=t, - config=runtime_config, + config=config, hidden_states=latents, prompt_embeds=prompt_embeds, pooled_prompt_embeds=pooled_prompt_embeds, ) - # 4.t Take one denoise step - latents = runtime_config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 5.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, - ) # fmt: off + # 6.t Call subscribers in-loop + ctx.in_loop(t, latents) # (Optional) Evaluate to enable progress tracking mx.eval(latents) except KeyboardInterrupt: # noqa: PERF203 - Callbacks.interruption( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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=runtime_config, - ) # fmt: off + # 7. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width) + # 8. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, - config=runtime_config, + config=config, seed=seed, prompt=prompt, quantization=self.bits, lora_paths=self.lora_paths, lora_scales=self.lora_scales, - redux_image_paths=runtime_config.redux_image_paths, - redux_image_strengths=runtime_config.redux_image_strengths, - image_strength=runtime_config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + redux_image_paths=config.redux_image_paths, + redux_image_strengths=config.redux_image_strengths, + image_strength=config.image_strength, + generation_time=config.time_steps.format_dict["elapsed"], ) @staticmethod def _get_prompt_embeddings( prompt: str, prompt_cache: dict[str, tuple[mx.array, mx.array]], - t5_tokenizer: TokenizerT5, - clip_tokenizer: TokenizerCLIP, + t5_tokenizer: Tokenizer, + clip_tokenizer: Tokenizer, t5_text_encoder: T5Encoder, clip_text_encoder: CLIPEncoder, image_paths: list[str] | list[Path], @@ -189,7 +176,7 @@ class Flux1Redux(nn.Module): image_encoder=image_encoder, image_embedder=image_embedder, image_strengths=image_strengths, - ) # fmt:off + ) # 3. Join text embeddings with all image embeddings prompt_embeds = mx.concatenate([prompt_embeds_txt] + image_embeds, axis=1) diff --git a/src/mflux/models/flux/variants/redux/redux_util.py b/src/mflux/models/flux/variants/redux/redux_util.py index cd12353..2803c3f 100644 --- a/src/mflux/models/flux/variants/redux/redux_util.py +++ b/src/mflux/models/flux/variants/redux/redux_util.py @@ -8,6 +8,26 @@ from mflux.utils.image_util import ImageUtil class ReduxUtil: + @staticmethod + def validate_redux_image_strengths( + redux_image_paths: list[Path], + redux_image_strengths: list[float] | None, + ) -> list[float] | None: + if not redux_image_strengths or len(redux_image_strengths) == 0: + return redux_image_strengths + + # If strengths provided but not enough for all images, fill with default (1.0) + if len(redux_image_strengths) < len(redux_image_paths): + redux_image_strengths.extend([1.0] * (len(redux_image_paths) - len(redux_image_strengths))) + + # If too many strengths provided, raise error + elif len(redux_image_strengths) > len(redux_image_paths): + raise ValueError( + f"Too many strengths provided ({len(redux_image_strengths)}), expted at most {len(redux_image_paths)}." + ) + + return redux_image_strengths + @staticmethod def embed_images( image_paths: list[str] | list[Path], diff --git a/src/mflux/models/flux/variants/redux/weight_handler_redux.py b/src/mflux/models/flux/variants/redux/weight_handler_redux.py deleted file mode 100644 index aee798b..0000000 --- a/src/mflux/models/flux/variants/redux/weight_handler_redux.py +++ /dev/null @@ -1,35 +0,0 @@ -from pathlib import Path - -from mflux.config.model_config import ModelConfig -from mflux.models.flux.weights.weight_handler import MetaData, WeightHandler -from mflux.utils.download import snapshot_download - - -class WeightHandlerRedux: - def __init__(self, siglip: dict, redux_encoder: dict, meta_data: MetaData): - self.siglip = siglip - self.redux_encoder = redux_encoder - self.meta_data = meta_data - - @staticmethod - def load_weights() -> "WeightHandlerRedux": - root_path = Path(snapshot_download(repo_id=ModelConfig.dev_redux().model_name, allow_patterns=["*.safetensors", "config.json"])) # fmt:off - - siglip, _, _ = WeightHandlerRedux._load_siglip_weights(root_path=root_path) - redux_encoder, _, _ = WeightHandlerRedux._load_redux_encoder_weights(root_path=root_path) - - return WeightHandlerRedux( - siglip=siglip, - redux_encoder=redux_encoder, - meta_data=MetaData(quantization_level=None) - ) # fmt:off - - @staticmethod - def _load_siglip_weights(root_path: Path) -> tuple[dict, int, str | None]: - weights, _, _ = WeightHandler.get_weights("image_encoder", root_path) - return weights, _, _ - - @staticmethod - def _load_redux_encoder_weights(root_path: Path) -> tuple[dict, int, str | None]: - weights, _, _ = WeightHandler.get_weights("image_embedder", root_path) - return weights, _, _ diff --git a/src/mflux/models/flux/variants/txt2img/flux.py b/src/mflux/models/flux/variants/txt2img/flux.py index adb1986..0a20027 100644 --- a/src/mflux/models/flux/variants/txt2img/flux.py +++ b/src/mflux/models/flux/variants/txt2img/flux.py @@ -1,13 +1,12 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator -from mflux.models.common.weights.model_saver import ModelSaver +from mflux.models.common.weights.saving.model_saver import ModelSaver from mflux.models.flux.flux_initializer import FluxInitializer from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -15,7 +14,7 @@ from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncod from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder from mflux.models.flux.model.flux_transformer.transformer import Transformer from mflux.models.flux.model.flux_vae.vae import VAE -from mflux.utils.array_util import ArrayUtil +from mflux.models.flux.weights.flux_weight_definition import FluxWeightDefinition from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -29,38 +28,52 @@ class Flux1(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.schnell(), ): super().__init__() FluxInitializer.init( - flux_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", negative_prompt: str | None = None, ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = LatentCreator.create_for_txt2img_or_img2img( seed=seed, - height=config.height, width=config.width, + height=config.height, img2img=Img2Img( vae=self.vae, latent_creator=FluxLatentCreator, @@ -74,26 +87,22 @@ class Flux1(nn.Module): 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_tokenizer=self.tokenizers["t5"], + clip_tokenizer=self.tokenizers["clip"], t5_text_encoder=self.t5_text_encoder, clip_text_encoder=self.clip_text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=config, - ) + # 3. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler latents = config.scheduler.scale_model_input(latents, t) - # 3.t Predict the noise + # 4.t Predict the noise noise = self.transformer( t=t, config=config, @@ -102,47 +111,26 @@ class Flux1(nn.Module): pooled_prompt_embeds=pooled_prompt_embeds, ) - # 4.t Take one denoise step - latents = config.scheduler.step( - model_output=noise, - timestep=t, - sample=latents, - ) + # 5.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=config, - time_steps=time_steps, - ) + # 6.t Call subscribers in-loop + ctx.in_loop(t, latents) # (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, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width) + # 8. Decode the latent array and return the image + latents = FluxLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, @@ -154,7 +142,7 @@ class Flux1(nn.Module): lora_scales=self.lora_scales, image_path=config.image_path, image_strength=config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + generation_time=config.time_steps.format_dict["elapsed"], ) @staticmethod @@ -169,16 +157,7 @@ class Flux1(nn.Module): model=self, bits=self.bits, base_path=base_path, - tokenizers=[ - ("clip_tokenizer.tokenizer", "tokenizer"), - ("t5_tokenizer.tokenizer", "tokenizer_2"), - ], - components=[ - ("vae", "vae"), - ("transformer", "transformer"), - ("clip_text_encoder", "text_encoder"), - ("t5_text_encoder", "text_encoder_2"), - ], + weight_definition=FluxWeightDefinition, ) def freeze(self, **kwargs): diff --git a/src/mflux/models/flux/weights/__init__.py b/src/mflux/models/flux/weights/__init__.py index e69de29..8001de8 100644 --- a/src/mflux/models/flux/weights/__init__.py +++ b/src/mflux/models/flux/weights/__init__.py @@ -0,0 +1,4 @@ +from mflux.models.flux.weights.flux_weight_definition import FluxWeightDefinition +from mflux.models.flux.weights.flux_weight_mapping import FluxWeightMapping + +__all__ = ["FluxWeightDefinition", "FluxWeightMapping"] diff --git a/src/mflux/models/flux/weights/flux_lora_mapping.py b/src/mflux/models/flux/weights/flux_lora_mapping.py new file mode 100644 index 0000000..60821d7 --- /dev/null +++ b/src/mflux/models/flux/weights/flux_lora_mapping.py @@ -0,0 +1,576 @@ +from mflux.models.common.lora.mapping.lora_mapping import LoRAMapping, LoRATarget +from mflux.models.common.lora.mapping.lora_transforms import LoraTransforms + + +class FluxLoRAMapping(LoRAMapping): + @staticmethod + def get_mapping() -> list[LoRATarget]: + targets = [] + + targets.extend(FluxLoRAMapping._get_standard_transformer_block_targets()) + targets.extend(FluxLoRAMapping._get_standard_single_transformer_block_targets()) + + targets.extend(FluxLoRAMapping._get_bfl_transformer_block_targets()) + targets.extend(FluxLoRAMapping._get_bfl_single_transformer_block_targets()) + + return targets + + @staticmethod + def _get_standard_transformer_block_targets() -> list[LoRATarget]: + return [ + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_q", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_q.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.to_q.lora_B", + "transformer_blocks.{block}.attn.to_q.lora_up.weight", + "transformer_blocks.{block}.attn.to_q.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_q.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.to_q.lora_A", + "transformer_blocks.{block}.attn.to_q.lora_down.weight", + "transformer_blocks.{block}.attn.to_q.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_q.alpha", + "transformer_blocks.{block}.attn.to_q.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_k", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_k.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.to_k.lora_B", + "transformer_blocks.{block}.attn.to_k.lora_up.weight", + "transformer_blocks.{block}.attn.to_k.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_k.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.to_k.lora_A", + "transformer_blocks.{block}.attn.to_k.lora_down.weight", + "transformer_blocks.{block}.attn.to_k.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_k.alpha", + "transformer_blocks.{block}.attn.to_k.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_v", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_v.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.to_v.lora_B", + "transformer_blocks.{block}.attn.to_v.lora_up.weight", + "transformer_blocks.{block}.attn.to_v.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_v.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.to_v.lora_A", + "transformer_blocks.{block}.attn.to_v.lora_down.weight", + "transformer_blocks.{block}.attn.to_v.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_v.alpha", + "transformer_blocks.{block}.attn.to_v.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_out.0", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_out.0.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.to_out.0.lora_B", + "transformer_blocks.{block}.attn.to_out.0.lora_up.weight", + "transformer_blocks.{block}.attn.to_out.0.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_out.0.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.to_out.0.lora_A", + "transformer_blocks.{block}.attn.to_out.0.lora_down.weight", + "transformer_blocks.{block}.attn.to_out.0.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_out.0.alpha", + "transformer_blocks.{block}.attn.to_out.0.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.add_q_proj", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_q_proj.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.add_q_proj.lora_B", + "transformer_blocks.{block}.attn.add_q_proj.lora_up.weight", + "transformer_blocks.{block}.attn.add_q_proj.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_q_proj.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.add_q_proj.lora_A", + "transformer_blocks.{block}.attn.add_q_proj.lora_down.weight", + "transformer_blocks.{block}.attn.add_q_proj.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_q_proj.alpha", + "transformer_blocks.{block}.attn.add_q_proj.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.add_k_proj", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_k_proj.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.add_k_proj.lora_B", + "transformer_blocks.{block}.attn.add_k_proj.lora_up.weight", + "transformer_blocks.{block}.attn.add_k_proj.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_k_proj.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.add_k_proj.lora_A", + "transformer_blocks.{block}.attn.add_k_proj.lora_down.weight", + "transformer_blocks.{block}.attn.add_k_proj.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_k_proj.alpha", + "transformer_blocks.{block}.attn.add_k_proj.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.add_v_proj", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_v_proj.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.add_v_proj.lora_B", + "transformer_blocks.{block}.attn.add_v_proj.lora_up.weight", + "transformer_blocks.{block}.attn.add_v_proj.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_v_proj.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.add_v_proj.lora_A", + "transformer_blocks.{block}.attn.add_v_proj.lora_down.weight", + "transformer_blocks.{block}.attn.add_v_proj.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.add_v_proj.alpha", + "transformer_blocks.{block}.attn.add_v_proj.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_add_out", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_add_out.lora_B.weight", + "transformer.transformer_blocks.{block}.attn.to_add_out.lora_B", + "transformer_blocks.{block}.attn.to_add_out.lora_up.weight", + "transformer_blocks.{block}.attn.to_add_out.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_add_out.lora_A.weight", + "transformer.transformer_blocks.{block}.attn.to_add_out.lora_A", + "transformer_blocks.{block}.attn.to_add_out.lora_down.weight", + "transformer_blocks.{block}.attn.to_add_out.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.attn.to_add_out.alpha", + "transformer_blocks.{block}.attn.to_add_out.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff.linear1", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.ff.net.0.proj.lora_B.weight", + "transformer.transformer_blocks.{block}.ff.net.0.proj.lora_B", + "transformer.transformer_blocks.{block}.ff.linear1.lora_B.weight", + "transformer.transformer_blocks.{block}.ff.linear1.lora_B", + "transformer_blocks.{block}.ff.net.0.proj.lora_up.weight", + "transformer_blocks.{block}.ff.linear1.lora_up.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.ff.net.0.proj.lora_A.weight", + "transformer.transformer_blocks.{block}.ff.net.0.proj.lora_A", + "transformer.transformer_blocks.{block}.ff.linear1.lora_A.weight", + "transformer.transformer_blocks.{block}.ff.linear1.lora_A", + "transformer_blocks.{block}.ff.net.0.proj.lora_down.weight", + "transformer_blocks.{block}.ff.linear1.lora_down.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.ff.net.0.proj.alpha", + "transformer.transformer_blocks.{block}.ff.linear1.alpha", + "transformer_blocks.{block}.ff.net.0.proj.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff.linear2", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.ff.net.2.lora_B.weight", + "transformer.transformer_blocks.{block}.ff.net.2.lora_B", + "transformer.transformer_blocks.{block}.ff.linear2.lora_B.weight", + "transformer.transformer_blocks.{block}.ff.linear2.lora_B", + "transformer_blocks.{block}.ff.net.2.lora_up.weight", + "transformer_blocks.{block}.ff.linear2.lora_up.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.ff.net.2.lora_A.weight", + "transformer.transformer_blocks.{block}.ff.net.2.lora_A", + "transformer.transformer_blocks.{block}.ff.linear2.lora_A.weight", + "transformer.transformer_blocks.{block}.ff.linear2.lora_A", + "transformer_blocks.{block}.ff.net.2.lora_down.weight", + "transformer_blocks.{block}.ff.linear2.lora_down.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.ff.net.2.alpha", + "transformer.transformer_blocks.{block}.ff.linear2.alpha", + "transformer_blocks.{block}.ff.net.2.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff_context.linear1", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.ff_context.net.0.proj.lora_B.weight", + "transformer.transformer_blocks.{block}.ff_context.net.0.proj.lora_B", + "transformer.transformer_blocks.{block}.ff_context.linear1.lora_B.weight", + "transformer.transformer_blocks.{block}.ff_context.linear1.lora_B", + "transformer_blocks.{block}.ff_context.net.0.proj.lora_up.weight", + "transformer_blocks.{block}.ff_context.linear1.lora_up.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.ff_context.net.0.proj.lora_A.weight", + "transformer.transformer_blocks.{block}.ff_context.net.0.proj.lora_A", + "transformer.transformer_blocks.{block}.ff_context.linear1.lora_A.weight", + "transformer.transformer_blocks.{block}.ff_context.linear1.lora_A", + "transformer_blocks.{block}.ff_context.net.0.proj.lora_down.weight", + "transformer_blocks.{block}.ff_context.linear1.lora_down.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.ff_context.net.0.proj.alpha", + "transformer.transformer_blocks.{block}.ff_context.linear1.alpha", + "transformer_blocks.{block}.ff_context.net.0.proj.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff_context.linear2", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.ff_context.net.2.lora_B.weight", + "transformer.transformer_blocks.{block}.ff_context.net.2.lora_B", + "transformer.transformer_blocks.{block}.ff_context.linear2.lora_B.weight", + "transformer.transformer_blocks.{block}.ff_context.linear2.lora_B", + "transformer_blocks.{block}.ff_context.net.2.lora_up.weight", + "transformer_blocks.{block}.ff_context.linear2.lora_up.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.ff_context.net.2.lora_A.weight", + "transformer.transformer_blocks.{block}.ff_context.net.2.lora_A", + "transformer.transformer_blocks.{block}.ff_context.linear2.lora_A.weight", + "transformer.transformer_blocks.{block}.ff_context.linear2.lora_A", + "transformer_blocks.{block}.ff_context.net.2.lora_down.weight", + "transformer_blocks.{block}.ff_context.linear2.lora_down.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.ff_context.net.2.alpha", + "transformer.transformer_blocks.{block}.ff_context.linear2.alpha", + "transformer_blocks.{block}.ff_context.net.2.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.norm1.linear", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.norm1.linear.lora_B.weight", + "transformer.transformer_blocks.{block}.norm1.linear.lora_B", + "transformer_blocks.{block}.norm1.linear.lora_up.weight", + "transformer_blocks.{block}.norm1.linear.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.norm1.linear.lora_A.weight", + "transformer.transformer_blocks.{block}.norm1.linear.lora_A", + "transformer_blocks.{block}.norm1.linear.lora_down.weight", + "transformer_blocks.{block}.norm1.linear.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.norm1.linear.alpha", + "transformer_blocks.{block}.norm1.linear.alpha", + ], + ), + LoRATarget( + model_path="transformer_blocks.{block}.norm1_context.linear", + possible_up_patterns=[ + "transformer.transformer_blocks.{block}.norm1_context.linear.lora_B.weight", + "transformer.transformer_blocks.{block}.norm1_context.linear.lora_B", + "transformer_blocks.{block}.norm1_context.linear.lora_up.weight", + "transformer_blocks.{block}.norm1_context.linear.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.transformer_blocks.{block}.norm1_context.linear.lora_A.weight", + "transformer.transformer_blocks.{block}.norm1_context.linear.lora_A", + "transformer_blocks.{block}.norm1_context.linear.lora_down.weight", + "transformer_blocks.{block}.norm1_context.linear.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.transformer_blocks.{block}.norm1_context.linear.alpha", + "transformer_blocks.{block}.norm1_context.linear.alpha", + ], + ), + ] + + @staticmethod + def _get_standard_single_transformer_block_targets() -> list[LoRATarget]: + return [ + LoRATarget( + model_path="single_transformer_blocks.{block}.attn.to_q", + possible_up_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_q.lora_B.weight", + "transformer.single_transformer_blocks.{block}.attn.to_q.lora_B", + "single_transformer_blocks.{block}.attn.to_q.lora_up.weight", + "single_transformer_blocks.{block}.attn.to_q.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_q.lora_A.weight", + "transformer.single_transformer_blocks.{block}.attn.to_q.lora_A", + "single_transformer_blocks.{block}.attn.to_q.lora_down.weight", + "single_transformer_blocks.{block}.attn.to_q.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_q.alpha", + "single_transformer_blocks.{block}.attn.to_q.alpha", + ], + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.attn.to_k", + possible_up_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_k.lora_B.weight", + "transformer.single_transformer_blocks.{block}.attn.to_k.lora_B", + "single_transformer_blocks.{block}.attn.to_k.lora_up.weight", + "single_transformer_blocks.{block}.attn.to_k.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_k.lora_A.weight", + "transformer.single_transformer_blocks.{block}.attn.to_k.lora_A", + "single_transformer_blocks.{block}.attn.to_k.lora_down.weight", + "single_transformer_blocks.{block}.attn.to_k.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_k.alpha", + "single_transformer_blocks.{block}.attn.to_k.alpha", + ], + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.attn.to_v", + possible_up_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_v.lora_B.weight", + "transformer.single_transformer_blocks.{block}.attn.to_v.lora_B", + "single_transformer_blocks.{block}.attn.to_v.lora_up.weight", + "single_transformer_blocks.{block}.attn.to_v.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_v.lora_A.weight", + "transformer.single_transformer_blocks.{block}.attn.to_v.lora_A", + "single_transformer_blocks.{block}.attn.to_v.lora_down.weight", + "single_transformer_blocks.{block}.attn.to_v.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.single_transformer_blocks.{block}.attn.to_v.alpha", + "single_transformer_blocks.{block}.attn.to_v.alpha", + ], + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.proj_mlp", + possible_up_patterns=[ + "transformer.single_transformer_blocks.{block}.proj_mlp.lora_B.weight", + "transformer.single_transformer_blocks.{block}.proj_mlp.lora_B", + "single_transformer_blocks.{block}.proj_mlp.lora_up.weight", + "single_transformer_blocks.{block}.proj_mlp.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.single_transformer_blocks.{block}.proj_mlp.lora_A.weight", + "transformer.single_transformer_blocks.{block}.proj_mlp.lora_A", + "single_transformer_blocks.{block}.proj_mlp.lora_down.weight", + "single_transformer_blocks.{block}.proj_mlp.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.single_transformer_blocks.{block}.proj_mlp.alpha", + "single_transformer_blocks.{block}.proj_mlp.alpha", + ], + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.proj_out", + possible_up_patterns=[ + "transformer.single_transformer_blocks.{block}.proj_out.lora_B.weight", + "transformer.single_transformer_blocks.{block}.proj_out.lora_B", + "single_transformer_blocks.{block}.proj_out.lora_up.weight", + "single_transformer_blocks.{block}.proj_out.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.single_transformer_blocks.{block}.proj_out.lora_A.weight", + "transformer.single_transformer_blocks.{block}.proj_out.lora_A", + "single_transformer_blocks.{block}.proj_out.lora_down.weight", + "single_transformer_blocks.{block}.proj_out.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.single_transformer_blocks.{block}.proj_out.alpha", + "single_transformer_blocks.{block}.proj_out.alpha", + ], + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.norm.linear", + possible_up_patterns=[ + "transformer.single_transformer_blocks.{block}.norm.linear.lora_B.weight", + "transformer.single_transformer_blocks.{block}.norm.linear.lora_B", + "single_transformer_blocks.{block}.norm.linear.lora_up.weight", + "single_transformer_blocks.{block}.norm.linear.lora_B.weight", + ], + possible_down_patterns=[ + "transformer.single_transformer_blocks.{block}.norm.linear.lora_A.weight", + "transformer.single_transformer_blocks.{block}.norm.linear.lora_A", + "single_transformer_blocks.{block}.norm.linear.lora_down.weight", + "single_transformer_blocks.{block}.norm.linear.lora_A.weight", + ], + possible_alpha_patterns=[ + "transformer.single_transformer_blocks.{block}.norm.linear.alpha", + "single_transformer_blocks.{block}.norm.linear.alpha", + ], + ), + ] + + @staticmethod + def _get_bfl_transformer_block_targets() -> list[LoRATarget]: + return [ + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_q", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.alpha"], + up_transform=LoraTransforms.split_q_up, + down_transform=LoraTransforms.split_q_down, + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_k", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.alpha"], + up_transform=LoraTransforms.split_k_up, + down_transform=LoraTransforms.split_k_down, + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_v", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_attn_qkv.alpha"], + up_transform=LoraTransforms.split_v_up, + down_transform=LoraTransforms.split_v_down, + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_out.0", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_attn_proj.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_attn_proj.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_attn_proj.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.add_q_proj", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.alpha"], + up_transform=LoraTransforms.split_q_up, + down_transform=LoraTransforms.split_q_down, + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.add_k_proj", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.alpha"], + up_transform=LoraTransforms.split_k_up, + down_transform=LoraTransforms.split_k_down, + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.add_v_proj", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_attn_qkv.alpha"], + up_transform=LoraTransforms.split_v_up, + down_transform=LoraTransforms.split_v_down, + ), + LoRATarget( + model_path="transformer_blocks.{block}.attn.to_add_out", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_attn_proj.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_attn_proj.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_attn_proj.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff.linear1", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_mlp_0.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_mlp_0.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_mlp_0.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff.linear2", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_mlp_2.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_mlp_2.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_mlp_2.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff_context.linear1", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_mlp_0.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_mlp_0.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_mlp_0.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.ff_context.linear2", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_mlp_2.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_mlp_2.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_mlp_2.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.norm1.linear", + possible_up_patterns=["lora_unet_double_blocks_{block}_img_mod_lin.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_img_mod_lin.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_img_mod_lin.alpha"], + ), + LoRATarget( + model_path="transformer_blocks.{block}.norm1_context.linear", + possible_up_patterns=["lora_unet_double_blocks_{block}_txt_mod_lin.lora_up.weight"], + possible_down_patterns=["lora_unet_double_blocks_{block}_txt_mod_lin.lora_down.weight"], + possible_alpha_patterns=["lora_unet_double_blocks_{block}_txt_mod_lin.alpha"], + ), + ] + + @staticmethod + def _get_bfl_single_transformer_block_targets() -> list[LoRATarget]: + return [ + LoRATarget( + model_path="single_transformer_blocks.{block}.attn.to_q", + possible_up_patterns=["lora_unet_single_blocks_{block}_linear1.lora_up.weight"], + possible_down_patterns=["lora_unet_single_blocks_{block}_linear1.lora_down.weight"], + possible_alpha_patterns=["lora_unet_single_blocks_{block}_linear1.alpha"], + up_transform=LoraTransforms.split_single_q_up, + down_transform=LoraTransforms.split_single_q_down, + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.attn.to_k", + possible_up_patterns=["lora_unet_single_blocks_{block}_linear1.lora_up.weight"], + possible_down_patterns=["lora_unet_single_blocks_{block}_linear1.lora_down.weight"], + possible_alpha_patterns=["lora_unet_single_blocks_{block}_linear1.alpha"], + up_transform=LoraTransforms.split_single_k_up, + down_transform=LoraTransforms.split_single_k_down, + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.attn.to_v", + possible_up_patterns=["lora_unet_single_blocks_{block}_linear1.lora_up.weight"], + possible_down_patterns=["lora_unet_single_blocks_{block}_linear1.lora_down.weight"], + possible_alpha_patterns=["lora_unet_single_blocks_{block}_linear1.alpha"], + up_transform=LoraTransforms.split_single_v_up, + down_transform=LoraTransforms.split_single_v_down, + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.proj_mlp", + possible_up_patterns=["lora_unet_single_blocks_{block}_linear1.lora_up.weight"], + possible_down_patterns=["lora_unet_single_blocks_{block}_linear1.lora_down.weight"], + possible_alpha_patterns=["lora_unet_single_blocks_{block}_linear1.alpha"], + up_transform=LoraTransforms.split_single_mlp_up, + down_transform=LoraTransforms.split_single_mlp_down, + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.proj_out", + possible_up_patterns=["lora_unet_single_blocks_{block}_linear2.lora_up.weight"], + possible_down_patterns=["lora_unet_single_blocks_{block}_linear2.lora_down.weight"], + possible_alpha_patterns=["lora_unet_single_blocks_{block}_linear2.alpha"], + ), + LoRATarget( + model_path="single_transformer_blocks.{block}.norm.linear", + possible_up_patterns=["lora_unet_single_blocks_{block}_modulation_lin.lora_up.weight"], + possible_down_patterns=["lora_unet_single_blocks_{block}_modulation_lin.lora_down.weight"], + possible_alpha_patterns=["lora_unet_single_blocks_{block}_modulation_lin.alpha"], + ), + ] diff --git a/src/mflux/models/flux/weights/flux_weight_definition.py b/src/mflux/models/flux/weights/flux_weight_definition.py new file mode 100644 index 0000000..01a8ab2 --- /dev/null +++ b/src/mflux/models/flux/weights/flux_weight_definition.py @@ -0,0 +1,163 @@ +from typing import List + +import mlx.nn as nn + +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.common.tokenizer import LanguageTokenizer +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition, TokenizerDefinition +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms +from mflux.models.flux.weights.flux_weight_mapping import FluxWeightMapping + + +class FluxWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="vae", + hf_subdir="vae", + precision=ModelConfig.precision, + mapping_getter=FluxWeightMapping.get_vae_mapping, + ), + ComponentDefinition( + name="transformer", + hf_subdir="transformer", + precision=ModelConfig.precision, + mapping_getter=FluxWeightMapping.get_transformer_mapping, + ), + ComponentDefinition( + name="t5_encoder", + hf_subdir="text_encoder_2", + model_attr="t5_text_encoder", + num_blocks=24, + precision=ModelConfig.precision, + mapping_getter=FluxWeightMapping.get_t5_encoder_mapping, + ), + ComponentDefinition( + name="clip_encoder", + hf_subdir="text_encoder", + model_attr="clip_text_encoder", + precision=ModelConfig.precision, + mapping_getter=FluxWeightMapping.get_clip_encoder_mapping, + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [ + TokenizerDefinition( + name="clip", + hf_subdir="tokenizer", + tokenizer_class="CLIPTokenizer", + encoder_class=LanguageTokenizer, + max_length=77, + download_patterns=["tokenizer/**"], + ), + TokenizerDefinition( + name="t5", + hf_subdir="tokenizer_2", + tokenizer_class="T5Tokenizer", + encoder_class=LanguageTokenizer, + max_length=256, # Will be overridden by model_config.max_sequence_length + download_patterns=["tokenizer_2/**"], + ), + ] + + @staticmethod + def get_download_patterns() -> List[str]: + return [ + "text_encoder/*.safetensors", + "text_encoder/*.json", + "text_encoder_2/*.safetensors", + "text_encoder_2/*.json", + "transformer/*.safetensors", + "transformer/*.json", + "vae/*.safetensors", + "vae/*.json", + ] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + return hasattr(module, "to_quantized") + + +class FluxControlnetWeightDefinition: + @staticmethod + def get_controlnet_component() -> ComponentDefinition: + return ComponentDefinition( + name="transformer_controlnet", + hf_subdir="", + loading_mode="single", + precision=ModelConfig.precision, + mapping_getter=FluxWeightMapping.get_controlnet_transformer_mapping, + ) + + @staticmethod + def get_components() -> List[ComponentDefinition]: + return FluxWeightDefinition.get_components() + [ + ComponentDefinition( + name="transformer_controlnet", + hf_subdir="transformer_controlnet", + precision=ModelConfig.precision, + mapping_getter=FluxWeightMapping.get_transformer_mapping, + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return FluxWeightDefinition.get_tokenizers() + + @staticmethod + def get_download_patterns() -> List[str]: + return FluxWeightDefinition.get_download_patterns() + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + return FluxWeightDefinition.quantization_predicate(path, module) + + +class FluxReduxWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="siglip", + hf_subdir="image_encoder", + mapping_getter=None, + model_attr="image_encoder", + precision=ModelConfig.precision, + bulk_transform=WeightTransforms.transpose_conv2d_weight, + weight_subkey="vision_model", + ), + ComponentDefinition( + name="redux_encoder", + hf_subdir="image_embedder", + mapping_getter=None, + model_attr="image_embedder", + precision=ModelConfig.precision, + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [] + + @staticmethod + def get_download_patterns() -> List[str]: + return [ + "image_encoder/*.safetensors", + "image_encoder/config.json", + "image_embedder/*.safetensors", + "image_embedder/config.json", + ] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + if isinstance(module, nn.Conv2d): + return False + if hasattr(module, "weight") and hasattr(module.weight, "shape"): + if module.weight.shape == (1152, 4304): + return False + if module.weight.shape[-1] % 64 != 0: + return False + return hasattr(module, "to_quantized") diff --git a/src/mflux/models/flux/weights/flux_weight_mapping.py b/src/mflux/models/flux/weights/flux_weight_mapping.py new file mode 100644 index 0000000..7d2b559 --- /dev/null +++ b/src/mflux/models/flux/weights/flux_weight_mapping.py @@ -0,0 +1,817 @@ +from typing import List + +from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms + + +class FluxWeightMapping(WeightMapping): + @staticmethod + def get_transformer_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="x_embedder.weight", + from_pattern=["x_embedder.weight"], + ), + WeightTarget( + to_pattern="x_embedder.bias", + from_pattern=["x_embedder.bias"], + ), + WeightTarget( + to_pattern="context_embedder.weight", + from_pattern=["context_embedder.weight"], + ), + WeightTarget( + to_pattern="context_embedder.bias", + from_pattern=["context_embedder.bias"], + ), + WeightTarget( + to_pattern="proj_out.weight", + from_pattern=["proj_out.weight"], + ), + WeightTarget( + to_pattern="proj_out.bias", + from_pattern=["proj_out.bias"], + ), + # TimeTextEmbed + WeightTarget( + to_pattern="time_text_embed.timestep_embedder.linear_1.weight", + from_pattern=["time_text_embed.timestep_embedder.linear_1.weight"], + ), + WeightTarget( + to_pattern="time_text_embed.timestep_embedder.linear_1.bias", + from_pattern=["time_text_embed.timestep_embedder.linear_1.bias"], + ), + WeightTarget( + to_pattern="time_text_embed.timestep_embedder.linear_2.weight", + from_pattern=["time_text_embed.timestep_embedder.linear_2.weight"], + ), + WeightTarget( + to_pattern="time_text_embed.timestep_embedder.linear_2.bias", + from_pattern=["time_text_embed.timestep_embedder.linear_2.bias"], + ), + WeightTarget( + to_pattern="time_text_embed.text_embedder.linear_1.weight", + from_pattern=["time_text_embed.text_embedder.linear_1.weight"], + ), + WeightTarget( + to_pattern="time_text_embed.text_embedder.linear_1.bias", + from_pattern=["time_text_embed.text_embedder.linear_1.bias"], + ), + WeightTarget( + to_pattern="time_text_embed.text_embedder.linear_2.weight", + from_pattern=["time_text_embed.text_embedder.linear_2.weight"], + ), + WeightTarget( + to_pattern="time_text_embed.text_embedder.linear_2.bias", + from_pattern=["time_text_embed.text_embedder.linear_2.bias"], + ), + WeightTarget( + to_pattern="time_text_embed.guidance_embedder.linear_1.weight", + from_pattern=["time_text_embed.guidance_embedder.linear_1.weight"], + required=False, + ), + WeightTarget( + to_pattern="time_text_embed.guidance_embedder.linear_1.bias", + from_pattern=["time_text_embed.guidance_embedder.linear_1.bias"], + required=False, + ), + WeightTarget( + to_pattern="time_text_embed.guidance_embedder.linear_2.weight", + from_pattern=["time_text_embed.guidance_embedder.linear_2.weight"], + required=False, + ), + WeightTarget( + to_pattern="time_text_embed.guidance_embedder.linear_2.bias", + from_pattern=["time_text_embed.guidance_embedder.linear_2.bias"], + required=False, + ), + WeightTarget( + to_pattern="norm_out.linear.weight", + from_pattern=["norm_out.linear.weight"], + ), + WeightTarget( + to_pattern="norm_out.linear.bias", + from_pattern=["norm_out.linear.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.norm1.linear.weight", + from_pattern=["transformer_blocks.{block}.norm1.linear.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.norm1.linear.bias", + from_pattern=["transformer_blocks.{block}.norm1.linear.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.norm1_context.linear.weight", + from_pattern=["transformer_blocks.{block}.norm1_context.linear.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.norm1_context.linear.bias", + from_pattern=["transformer_blocks.{block}.norm1_context.linear.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_q.weight", + from_pattern=["transformer_blocks.{block}.attn.to_q.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_q.bias", + from_pattern=["transformer_blocks.{block}.attn.to_q.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_k.weight", + from_pattern=["transformer_blocks.{block}.attn.to_k.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_k.bias", + from_pattern=["transformer_blocks.{block}.attn.to_k.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_v.weight", + from_pattern=["transformer_blocks.{block}.attn.to_v.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_v.bias", + from_pattern=["transformer_blocks.{block}.attn.to_v.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_out.0.weight", + from_pattern=["transformer_blocks.{block}.attn.to_out.0.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_out.0.bias", + from_pattern=["transformer_blocks.{block}.attn.to_out.0.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.add_q_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_q_proj.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.add_q_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_q_proj.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.add_k_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_k_proj.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.add_k_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_k_proj.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.add_v_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_v_proj.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.add_v_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_v_proj.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_add_out.weight", + from_pattern=["transformer_blocks.{block}.attn.to_add_out.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.to_add_out.bias", + from_pattern=["transformer_blocks.{block}.attn.to_add_out.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.norm_q.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_q.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.norm_k.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_k.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.norm_added_q.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_added_q.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.attn.norm_added_k.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_added_k.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff.linear1.weight", + from_pattern=["transformer_blocks.{block}.ff.net.0.proj.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff.linear1.bias", + from_pattern=["transformer_blocks.{block}.ff.net.0.proj.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff.linear2.weight", + from_pattern=["transformer_blocks.{block}.ff.net.2.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff.linear2.bias", + from_pattern=["transformer_blocks.{block}.ff.net.2.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff_context.linear1.weight", + from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff_context.linear1.bias", + from_pattern=["transformer_blocks.{block}.ff_context.net.0.proj.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff_context.linear2.weight", + from_pattern=["transformer_blocks.{block}.ff_context.net.2.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.ff_context.linear2.bias", + from_pattern=["transformer_blocks.{block}.ff_context.net.2.bias"], + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.norm.linear.weight", + from_pattern=["single_transformer_blocks.{block}.norm.linear.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.norm.linear.bias", + from_pattern=["single_transformer_blocks.{block}.norm.linear.bias"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.to_q.weight", + from_pattern=["single_transformer_blocks.{block}.attn.to_q.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.to_q.bias", + from_pattern=["single_transformer_blocks.{block}.attn.to_q.bias"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.to_k.weight", + from_pattern=["single_transformer_blocks.{block}.attn.to_k.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.to_k.bias", + from_pattern=["single_transformer_blocks.{block}.attn.to_k.bias"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.to_v.weight", + from_pattern=["single_transformer_blocks.{block}.attn.to_v.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.to_v.bias", + from_pattern=["single_transformer_blocks.{block}.attn.to_v.bias"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.norm_q.weight", + from_pattern=["single_transformer_blocks.{block}.attn.norm_q.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.attn.norm_k.weight", + from_pattern=["single_transformer_blocks.{block}.attn.norm_k.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.proj_mlp.weight", + from_pattern=["single_transformer_blocks.{block}.proj_mlp.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.proj_mlp.bias", + from_pattern=["single_transformer_blocks.{block}.proj_mlp.bias"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.proj_out.weight", + from_pattern=["single_transformer_blocks.{block}.proj_out.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="single_transformer_blocks.{block}.proj_out.bias", + from_pattern=["single_transformer_blocks.{block}.proj_out.bias"], + max_blocks=38, + ), + ] + + @staticmethod + def get_controlnet_transformer_mapping() -> List[WeightTarget]: + return FluxWeightMapping.get_transformer_mapping() + [ + WeightTarget( + to_pattern="controlnet_x_embedder.weight", + from_pattern=["controlnet_x_embedder.weight"], + ), + WeightTarget( + to_pattern="controlnet_x_embedder.bias", + from_pattern=["controlnet_x_embedder.bias"], + ), + WeightTarget( + to_pattern="controlnet_blocks.{block}.weight", + from_pattern=["controlnet_blocks.{block}.weight"], + max_blocks=19, + ), + WeightTarget( + to_pattern="controlnet_blocks.{block}.bias", + from_pattern=["controlnet_blocks.{block}.bias"], + max_blocks=19, + ), + WeightTarget( + to_pattern="controlnet_single_blocks.{block}.weight", + from_pattern=["controlnet_single_blocks.{block}.weight"], + max_blocks=38, + ), + WeightTarget( + to_pattern="controlnet_single_blocks.{block}.bias", + from_pattern=["controlnet_single_blocks.{block}.bias"], + max_blocks=38, + ), + ] + + @staticmethod + def get_vae_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="decoder.conv_in.conv2d.weight", + from_pattern=["decoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.conv_in.conv2d.bias", + from_pattern=["decoder.conv_in.bias"], + ), + WeightTarget( + to_pattern="decoder.conv_out.conv2d.weight", + from_pattern=["decoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.conv_out.conv2d.bias", + from_pattern=["decoder.conv_out.bias"], + ), + WeightTarget( + to_pattern="decoder.conv_norm_out.norm.weight", + from_pattern=["decoder.conv_norm_out.weight"], + ), + WeightTarget( + to_pattern="decoder.conv_norm_out.norm.bias", + from_pattern=["decoder.conv_norm_out.bias"], + ), + WeightTarget( + to_pattern="encoder.conv_in.conv2d.weight", + from_pattern=["encoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.conv_in.conv2d.bias", + from_pattern=["encoder.conv_in.bias"], + ), + WeightTarget( + to_pattern="encoder.conv_out.conv2d.weight", + from_pattern=["encoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.conv_out.conv2d.bias", + from_pattern=["encoder.conv_out.bias"], + ), + WeightTarget( + to_pattern="encoder.conv_norm_out.norm.weight", + from_pattern=["encoder.conv_norm_out.weight"], + ), + WeightTarget( + to_pattern="encoder.conv_norm_out.norm.bias", + from_pattern=["encoder.conv_norm_out.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm1.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm1.bias", + from_pattern=["decoder.mid_block.resnets.{i}.norm1.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv1.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv1.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm2.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm2.bias", + from_pattern=["decoder.mid_block.resnets.{i}.norm2.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv2.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv2.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm1.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm1.bias", + from_pattern=["encoder.mid_block.resnets.{i}.norm1.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv1.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv1.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm2.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm2.bias", + from_pattern=["encoder.mid_block.resnets.{i}.norm2.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv2.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv2.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.group_norm.weight", + from_pattern=["decoder.mid_block.attentions.0.group_norm.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.group_norm.bias", + from_pattern=["decoder.mid_block.attentions.0.group_norm.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_q.weight", + from_pattern=["decoder.mid_block.attentions.0.to_q.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_q.bias", + from_pattern=["decoder.mid_block.attentions.0.to_q.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_k.weight", + from_pattern=["decoder.mid_block.attentions.0.to_k.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_k.bias", + from_pattern=["decoder.mid_block.attentions.0.to_k.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_v.weight", + from_pattern=["decoder.mid_block.attentions.0.to_v.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_v.bias", + from_pattern=["decoder.mid_block.attentions.0.to_v.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_out.0.weight", + from_pattern=["decoder.mid_block.attentions.0.to_out.0.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_out.0.bias", + from_pattern=["decoder.mid_block.attentions.0.to_out.0.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.group_norm.weight", + from_pattern=["encoder.mid_block.attentions.0.group_norm.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.group_norm.bias", + from_pattern=["encoder.mid_block.attentions.0.group_norm.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_q.weight", + from_pattern=["encoder.mid_block.attentions.0.to_q.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_q.bias", + from_pattern=["encoder.mid_block.attentions.0.to_q.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_k.weight", + from_pattern=["encoder.mid_block.attentions.0.to_k.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_k.bias", + from_pattern=["encoder.mid_block.attentions.0.to_k.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_v.weight", + from_pattern=["encoder.mid_block.attentions.0.to_v.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_v.bias", + from_pattern=["encoder.mid_block.attentions.0.to_v.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_out.0.weight", + from_pattern=["encoder.mid_block.attentions.0.to_out.0.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_out.0.bias", + from_pattern=["encoder.mid_block.attentions.0.to_out.0.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.weight"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.weight"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"], + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsamplers.0.conv.weight", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsamplers.0.conv.bias", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.conv.bias"], + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.weight"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.weight"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"], + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.downsamplers.0.conv.weight", + from_pattern=["encoder.down_blocks.{block}.downsamplers.0.conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.downsamplers.0.conv.bias", + from_pattern=["encoder.down_blocks.{block}.downsamplers.0.conv.bias"], + required=False, + ), + WeightTarget( + to_pattern="quant_conv.weight", + from_pattern=["quant_conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="quant_conv.bias", + from_pattern=["quant_conv.bias"], + ), + WeightTarget( + to_pattern="post_quant_conv.weight", + from_pattern=["post_quant_conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="post_quant_conv.bias", + from_pattern=["post_quant_conv.bias"], + ), + ] + + @staticmethod + def get_t5_encoder_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="shared.weight", + from_pattern=["shared.weight"], + ), + WeightTarget( + to_pattern="final_layer_norm.weight", + from_pattern=["encoder.final_layer_norm.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.attention.layer_norm.weight", + from_pattern=["encoder.block.{block}.layer.0.layer_norm.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.attention.SelfAttention.q.weight", + from_pattern=["encoder.block.{block}.layer.0.SelfAttention.q.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.attention.SelfAttention.k.weight", + from_pattern=["encoder.block.{block}.layer.0.SelfAttention.k.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.attention.SelfAttention.v.weight", + from_pattern=["encoder.block.{block}.layer.0.SelfAttention.v.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.attention.SelfAttention.o.weight", + from_pattern=["encoder.block.{block}.layer.0.SelfAttention.o.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.attention.SelfAttention.relative_attention_bias.weight", + from_pattern=["encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight"], + max_blocks=24, + ), + WeightTarget( + to_pattern="t5_blocks.{block}.ff.layer_norm.weight", + from_pattern=["encoder.block.{block}.layer.1.layer_norm.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.ff.DenseReluDense.wi_0.weight", + from_pattern=["encoder.block.{block}.layer.1.DenseReluDense.wi_0.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.ff.DenseReluDense.wi_1.weight", + from_pattern=["encoder.block.{block}.layer.1.DenseReluDense.wi_1.weight"], + ), + WeightTarget( + to_pattern="t5_blocks.{block}.ff.DenseReluDense.wo.weight", + from_pattern=["encoder.block.{block}.layer.1.DenseReluDense.wo.weight"], + ), + ] + + @staticmethod + def get_clip_encoder_mapping() -> List[WeightTarget]: + return [ + WeightTarget( + to_pattern="text_model.embeddings.token_embedding.weight", + from_pattern=["text_model.embeddings.token_embedding.weight"], + ), + WeightTarget( + to_pattern="text_model.embeddings.position_embedding.weight", + from_pattern=["text_model.embeddings.position_embedding.weight"], + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.q_proj.weight", + from_pattern=["text_model.encoder.layers.{block}.self_attn.q_proj.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.q_proj.bias", + from_pattern=["text_model.encoder.layers.{block}.self_attn.q_proj.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.k_proj.weight", + from_pattern=["text_model.encoder.layers.{block}.self_attn.k_proj.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.k_proj.bias", + from_pattern=["text_model.encoder.layers.{block}.self_attn.k_proj.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.v_proj.weight", + from_pattern=["text_model.encoder.layers.{block}.self_attn.v_proj.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.v_proj.bias", + from_pattern=["text_model.encoder.layers.{block}.self_attn.v_proj.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.out_proj.weight", + from_pattern=["text_model.encoder.layers.{block}.self_attn.out_proj.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.self_attn.out_proj.bias", + from_pattern=["text_model.encoder.layers.{block}.self_attn.out_proj.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.layer_norm1.weight", + from_pattern=["text_model.encoder.layers.{block}.layer_norm1.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.layer_norm1.bias", + from_pattern=["text_model.encoder.layers.{block}.layer_norm1.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.mlp.fc1.weight", + from_pattern=["text_model.encoder.layers.{block}.mlp.fc1.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.mlp.fc1.bias", + from_pattern=["text_model.encoder.layers.{block}.mlp.fc1.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.mlp.fc2.weight", + from_pattern=["text_model.encoder.layers.{block}.mlp.fc2.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.mlp.fc2.bias", + from_pattern=["text_model.encoder.layers.{block}.mlp.fc2.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.layer_norm2.weight", + from_pattern=["text_model.encoder.layers.{block}.layer_norm2.weight"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.encoder.layers.{block}.layer_norm2.bias", + from_pattern=["text_model.encoder.layers.{block}.layer_norm2.bias"], + max_blocks=12, + ), + WeightTarget( + to_pattern="text_model.final_layer_norm.weight", + from_pattern=["text_model.final_layer_norm.weight"], + ), + WeightTarget( + to_pattern="text_model.final_layer_norm.bias", + from_pattern=["text_model.final_layer_norm.bias"], + ), + ] diff --git a/src/mflux/models/flux/weights/lora_converter.py b/src/mflux/models/flux/weights/lora_converter.py deleted file mode 100644 index 6552036..0000000 --- a/src/mflux/models/flux/weights/lora_converter.py +++ /dev/null @@ -1,251 +0,0 @@ -import logging - -import mlx.core as mx -import torch -from mlx.utils import tree_unflatten -from safetensors import safe_open - -logger = logging.getLogger(__name__) - - -# This script is based on `convert_flux_lora.py` from `kohya-ss/sd-scripts`. -# For more info, see: https://github.com/kohya-ss/sd-scripts/blob/sd3/networks/convert_flux_lora.py - - -class LoRAConverter: - @staticmethod - def load_weights(lora_path: str) -> dict: - state_dict = LoRAConverter._load_pytorch_weights(lora_path) - state_dict = LoRAConverter._convert_weights_to_diffusers(state_dict) - state_dict = LoRAConverter._convert_to_mlx(state_dict) - state_dict = list(state_dict.items()) - state_dict = tree_unflatten(state_dict) - return state_dict - - @staticmethod - def _load_pytorch_weights(lora_path: str) -> dict: - state_dict = {} - with safe_open(lora_path, framework="pt") as f: - for k in f.keys(): - state_dict[k] = f.get_tensor(k) - return state_dict - - @staticmethod - def _convert_weights_to_diffusers(source: dict) -> dict: - target = {} - for i in range(19): - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_img_attn_proj", - f"transformer.transformer_blocks.{i}.attn.to_out.0", - ) - LoRAConverter._convert_to_diffusers_cat( - source, - target, - f"lora_unet_double_blocks_{i}_img_attn_qkv", - [ - f"transformer.transformer_blocks.{i}.attn.to_q", - f"transformer.transformer_blocks.{i}.attn.to_k", - f"transformer.transformer_blocks.{i}.attn.to_v", - ], - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_img_mlp_0", - f"transformer.transformer_blocks.{i}.ff.net.0.proj", - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_img_mlp_2", - f"transformer.transformer_blocks.{i}.ff.net.2", - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_img_mod_lin", - f"transformer.transformer_blocks.{i}.norm1.linear", - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_txt_attn_proj", - f"transformer.transformer_blocks.{i}.attn.to_add_out", - ) - LoRAConverter._convert_to_diffusers_cat( - source, - target, - f"lora_unet_double_blocks_{i}_txt_attn_qkv", - [ - f"transformer.transformer_blocks.{i}.attn.add_q_proj", - f"transformer.transformer_blocks.{i}.attn.add_k_proj", - f"transformer.transformer_blocks.{i}.attn.add_v_proj", - ], - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_txt_mlp_0", - f"transformer.transformer_blocks.{i}.ff_context.net.0.proj", - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_txt_mlp_2", - f"transformer.transformer_blocks.{i}.ff_context.net.2", - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_double_blocks_{i}_txt_mod_lin", - f"transformer.transformer_blocks.{i}.norm1_context.linear", - ) - - for i in range(38): - LoRAConverter._convert_to_diffusers_cat( - source, - target, - f"lora_unet_single_blocks_{i}_linear1", - [ - f"transformer.single_transformer_blocks.{i}.attn.to_q", - f"transformer.single_transformer_blocks.{i}.attn.to_k", - f"transformer.single_transformer_blocks.{i}.attn.to_v", - f"transformer.single_transformer_blocks.{i}.proj_mlp", - ], - dims=[3072, 3072, 3072, 12288], - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_single_blocks_{i}_linear2", - f"transformer.single_transformer_blocks.{i}.proj_out", - ) - LoRAConverter._convert_to_diffusers( - source, - target, - f"lora_unet_single_blocks_{i}_modulation_lin", - f"transformer.single_transformer_blocks.{i}.norm.linear", - ) - - if len(source) > 0: - logger.warning(f"Unsupported keys for diffusers: {source.keys()}") - return target - - @staticmethod - def _convert_to_diffusers(source: dict, target: dict, source_key: str, target_key: str): - if source_key + ".lora_down.weight" not in source: - return - down_weight = source.pop(source_key + ".lora_down.weight") - - # scale weight by alpha and dim - rank = down_weight.shape[0] - alpha = source.pop(source_key + ".alpha").item() # alpha is scalar - scale = alpha / rank # LoRA is scaled by 'alpha / rank' in forward pass, so we need to scale it back here - - # calculate scale_down and scale_up to keep the same value. if scale is 4, scale_down is 2 and scale_up is 2 - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - - target[target_key + ".lora_A.weight"] = down_weight * scale_down - target[target_key + ".lora_B.weight"] = source.pop(source_key + ".lora_up.weight") * scale_up - - @staticmethod - def _convert_to_diffusers_cat( - source: dict, - target: dict, - source_key: str, - target_keys: list[str], - dims=None, - ): - if source_key + ".lora_down.weight" not in source: - return - down_weight = source.pop(source_key + ".lora_down.weight") - up_weight = source.pop(source_key + ".lora_up.weight") - source_lora_rank = down_weight.shape[0] - - # scale weight by alpha and dim - alpha = source.pop(source_key + ".alpha") - scale = alpha / source_lora_rank - - # calculate scale_down and scale_up - scale_down = scale - scale_up = 1.0 - while scale_down * 2 < scale_up: - scale_down *= 2 - scale_up /= 2 - - down_weight = down_weight * scale_down - up_weight = up_weight * scale_up - - # calculate dims if not provided - num_splits = len(target_keys) - if dims is None: - dims = [up_weight.shape[0] // num_splits] * num_splits - else: - assert sum(dims) == up_weight.shape[0] - - # check up-weight is sparse or not - is_sparse = False - if source_lora_rank % num_splits == 0: - diffusers_rank = source_lora_rank // num_splits - is_sparse = True - i = 0 - for j in range(len(dims)): - for k in range(len(dims)): - if j == k: - continue - is_sparse = is_sparse and torch.all( - up_weight[ - i : i + dims[j], - k * diffusers_rank : (k + 1) * diffusers_rank, - ] - == 0 - ) - i += dims[j] - if is_sparse: - logger.info(f"weight is sparse: {source_key}") - - # make diffusers weight - diffusers_down_keys = [k + ".lora_A.weight" for k in target_keys] - diffusers_up_keys = [k + ".lora_B.weight" for k in target_keys] - if not is_sparse: - # down_weight is copied to each split - target.update({k: down_weight for k in diffusers_down_keys}) - - # up_weight is split to each split - target.update({k: v for k, v in zip(diffusers_up_keys, torch.split(up_weight, dims, dim=0))}) - else: - # down_weight is chunked to each split - target.update( - { - k: v - for k, v in zip( - diffusers_down_keys, - torch.chunk(down_weight, num_splits, dim=0), - ) - } - ) - - # up_weight is sparse: only non-zero values are copied to each split - i = 0 - for j in range(len(dims)): - target[diffusers_up_keys[j]] = up_weight[ - i : i + dims[j], - j * diffusers_rank : (j + 1) * diffusers_rank, - ].contiguous() - i += dims[j] - - @staticmethod - def _convert_to_mlx(torch_dict: dict): - mlx_dict = {} - for key, value in torch_dict.items(): - if isinstance(value, torch.Tensor): - mlx_dict[key] = mx.array(value.detach().cpu()) - else: - mlx_dict[key] = value - return mlx_dict diff --git a/src/mflux/models/flux/weights/weight_handler.py b/src/mflux/models/flux/weights/weight_handler.py deleted file mode 100644 index 3768a7d..0000000 --- a/src/mflux/models/flux/weights/weight_handler.py +++ /dev/null @@ -1,241 +0,0 @@ -from dataclasses import dataclass -from pathlib import Path - -import mlx.core as mx -from mlx.utils import tree_unflatten - -from mflux.models.flux.weights.lora_converter import LoRAConverter -from mflux.models.flux.weights.weight_util import WeightUtil -from mflux.utils.download import snapshot_download - - -@dataclass -class MetaData: - quantization_level: int | None = None - scale: float | None = None - is_lora: bool = False - mflux_version: str | None = None - - -class WeightHandler: - def __init__( - self, - meta_data: MetaData, - clip_encoder: dict | None = None, - t5_encoder: dict | None = None, - vae: dict | None = None, - transformer: dict | None = None, - ): - self.clip_encoder = clip_encoder - self.t5_encoder = t5_encoder - self.vae = vae - self.transformer = transformer - self.meta_data = meta_data - - @staticmethod - 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) - - return WeightHandler( - clip_encoder=clip_encoder, - t5_encoder=t5_encoder, - vae=vae, - transformer=transformer, - meta_data=MetaData( - quantization_level=quantization_level, - scale=None, - is_lora=False, - mflux_version=mflux_version, - ), - ) - - def num_transformer_blocks(self) -> int: - return len(self.transformer["transformer_blocks"]) - - def num_single_transformer_blocks(self) -> int: - return len(self.transformer["single_transformer_blocks"]) - - @staticmethod - def _load_clip_encoder(root_path: Path) -> tuple[dict, int, str | None]: - weights, quantization_level, mflux_version = WeightHandler.get_weights("text_encoder", root_path) - return weights, quantization_level, mflux_version - - @staticmethod - def _load_t5_encoder(root_path: Path) -> tuple[dict, int, str | None]: - weights, quantization_level, mflux_version = WeightHandler.get_weights("text_encoder_2", root_path) - - # Quantized weights (i.e. ones exported from this project) don't need any post-processing. - if quantization_level is not None: - return weights, quantization_level, mflux_version - - # Reshape and process the huggingface weights - weights["final_layer_norm"] = weights["encoder"]["final_layer_norm"] - for block in weights["encoder"]["block"]: - attention = block["layer"][0] - ff = block["layer"][1] - block.pop("layer") - block["attention"] = attention - block["ff"] = ff - - weights["t5_blocks"] = weights["encoder"]["block"] - - # Only the first layer has the weights for "relative_attention_bias", we duplicate them here to keep code simple - relative_attention_bias = weights["t5_blocks"][0]["attention"]["SelfAttention"]["relative_attention_bias"] - for block in weights["t5_blocks"][1:]: - block["attention"]["SelfAttention"]["relative_attention_bias"] = relative_attention_bias - - 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) - - if lora_path: - if "transformer" not in weights: - weights = LoRAConverter.load_weights(lora_path) - weights = weights["transformer"] - - # Quantized weights (i.e. ones exported from this project) don't need any post-processing. - if quantization_level or mflux_version: - return weights, quantization_level, mflux_version - - # Reshape and process the huggingface weights - if "transformer_blocks" in weights: - for block in weights["transformer_blocks"]: - # 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 - def _load_vae(root_path: Path) -> tuple[dict, int, str | None]: - weights, quantization_level, mflux_version = WeightHandler.get_weights("vae", root_path) - - # Quantized weights (i.e. ones exported from this project) don't need any post-processing. - if quantization_level is not None: - return weights, quantization_level, mflux_version - - # Reshape and process the huggingface weights - weights["decoder"]["conv_in"] = {"conv2d": weights["decoder"]["conv_in"]} - weights["decoder"]["conv_out"] = {"conv2d": weights["decoder"]["conv_out"]} - weights["decoder"]["conv_norm_out"] = {"norm": weights["decoder"]["conv_norm_out"]} - weights["encoder"]["conv_in"] = {"conv2d": weights["encoder"]["conv_in"]} - weights["encoder"]["conv_out"] = {"conv2d": weights["encoder"]["conv_out"]} - 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, - root_path: Path | None = None, - lora_path: str | None = None, - ) -> tuple[dict, int, str | None]: - weights = [] - quantization_level = None - mflux_version = None - - if root_path is not None: - 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: - quantization_level = data[1].get("quantization_level") - mflux_version = data[1].get("mflux_version") - weights.extend(weight) - - if lora_path and root_path is None: - data = mx.load(lora_path, return_metadata=True) - weight = list(data[0].items()) - if len(data) > 1: - mflux_version = data[1].get("mflux_version", None) - weights.extend(weight) - - # Non huggingface weights (i.e. ones exported from this project) don't need any reshaping. - if quantization_level is not None or mflux_version is not None: - return tree_unflatten(weights), quantization_level, mflux_version - - # Huggingface weights needs to be reshaped - weights = [WeightUtil.reshape_weights(k, v) for k, v in weights] - weights = WeightUtil.flatten(weights) - unflatten = tree_unflatten(weights) - return unflatten, quantization_level, mflux_version - - @staticmethod - def download_or_get_cached_weights(repo_id: str) -> Path: - return Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=[ - "text_encoder/*.safetensors", - "text_encoder/*.json", - "text_encoder_2/*.safetensors", - "text_encoder_2/*.json", - "transformer/*.safetensors", - "transformer/*.json", - "vae/*.safetensors", - "vae/*.json", - ], - ) - ) - - @staticmethod - def _download_transformer_weights(repo_id: str) -> Path: - return Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=[ - "transformer/*.safetensors", - "*.safetensors", - ], - ) - ) diff --git a/src/mflux/models/flux/weights/weight_handler_lora.py b/src/mflux/models/flux/weights/weight_handler_lora.py deleted file mode 100644 index eb13432..0000000 --- a/src/mflux/models/flux/weights/weight_handler_lora.py +++ /dev/null @@ -1,161 +0,0 @@ -import mlx.core as mx -import mlx.nn as nn -from mlx.utils import tree_flatten, tree_unflatten - -from mflux.models.common.lora.layer.fused_linear_lora_layer import FusedLoRALinear -from mflux.models.common.lora.layer.linear_lora_layer import LoRALinear -from mflux.models.flux.variants.dreambooth.lora_layers.lora_layers import LoRALayers -from mflux.models.flux.weights.weight_handler import MetaData, WeightHandler - - -class WeightHandlerLoRA: - def __init__(self, weight_handlers: list[WeightHandler]): - self.weight_handlers = weight_handlers - - @staticmethod - def load_lora_weights( - transformer: nn.Module, - lora_files: list[str], - lora_scales: list[float] | None = None, - ) -> list["WeightHandler"]: - lora_weights = [] - if lora_files: - lora_scales = WeightHandlerLoRA._validate_lora_scales(lora_files, lora_scales) - for lora_file, lora_scale in zip(lora_files, lora_scales): - weights, _, mflux_version = WeightHandler.load_transformer(lora_path=lora_file) - weights = dict(tree_flatten(weights)) - weights = {key.removesuffix(".weight"): value for key, value in weights.items()} - weights = {f"transformer.{key}": value for key, value in weights.items()} - weights = {key: mx.transpose(value) for key, value in weights.items()} - lora_transformer_dict = LoRALayers.transformer_dict_from_template(weights, transformer, lora_scale) - transformer_weights = tree_unflatten(list(lora_transformer_dict.items()))["transformer"] - weights = WeightHandler( - clip_encoder=None, - t5_encoder=None, - vae=None, - transformer=transformer_weights, - meta_data=MetaData( - quantization_level=None, - scale=lora_scale, - is_lora=True, - mflux_version=mflux_version, - ), - ) - lora_weights.append(weights) - - return lora_weights - - @staticmethod - def set_lora_weights(transformer: nn.Module, loras: list["WeightHandler"]) -> None: - if loras: - lora_transformer_weights = [lora.transformer for lora in loras] - fused_weights = WeightHandlerLoRA._fuse_multiple_lora_dicts(lora_transformer_weights) - fused_weights = WeightHandler( - meta_data=MetaData(), - clip_encoder=None, - t5_encoder=None, - vae=None, - transformer=fused_weights, - ) - - WeightHandlerLoRA.set_lora_layers( - transformer_module=transformer, - lora_layers=LoRALayers(weights=fused_weights), - ) - - @staticmethod - def _fuse_multiple_lora_dicts(dicts: list[dict]) -> dict: - if not dicts: - raise ValueError("No dictionaries provided for fusion.") - if len(dicts) == 1: - return dicts[0] - - # Collect all unique keys across all dictionaries - all_keys = set().union(*dicts) - fused_dict = {} - - for key in all_keys: - # Get all values for this key, filtering out dictionaries that don't have it - values = [d[key] for d in dicts if key in d] - - # Skip if no values found (shouldn't happen due to how we collect keys) - if not values: - continue - - first_value = values[0] - - # Handle nested dictionaries - if all(isinstance(v, dict) and not isinstance(v, LoRALinear) for v in values): - fused_dict[key] = WeightHandlerLoRA._fuse_multiple_lora_dicts(values) - - # Handle LoRALinear layers - elif all(isinstance(v, LoRALinear) for v in values): - fused_dict[key] = FusedLoRALinear(base_linear=first_value.linear, loras=values) - - # Handle lists - elif all(isinstance(v, list) for v in values): - # Get the maximum length of all lists - max_length = max(len(v) for v in values) - - # Initialize the fused list - fused_dict[key] = [] - - # Process each index up to the maximum length - for idx in range(max_length): - # Get elements at current index from lists that are long enough - elements = [v[idx] for v in values if idx < len(v)] - - if not elements: - continue - - # If elements are dicts or LoRALinear, recursively fuse them - if all(isinstance(e, (dict, LoRALinear)) for e in elements): - fused_element = WeightHandlerLoRA._fuse_multiple_lora_dicts([{str(idx): e} for e in elements])[str(idx)] # fmt:off - fused_dict[key].append(fused_element) - else: - # For non-LoRALinear types, keep the first element - fused_dict[key].append(elements[0]) - - else: - types_str = ", ".join(type(v).__name__ for v in values) - raise ValueError(f"Incompatible types for key {key}: {types_str}") - - return fused_dict - - @staticmethod - def _validate_lora_scales(lora_files: list[str], lora_scales: list[float]) -> list[float]: - if len(lora_files) == 1: - if not lora_scales: - lora_scales = [1.0] - if len(lora_scales) > 1: - raise ValueError("Please provide a single scale for the LoRA, or skip it to default to 1") - elif len(lora_files) > 1: - if len(lora_files) != len(lora_scales): - raise ValueError("When providing multiple LoRAs, be sure to specify a scale for each one respectively") - return lora_scales - - @staticmethod - 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): - LoRALayers.set_transformer_block( - transformer_block=transformer_module.transformer_blocks[i], - dictionary=weights, - ) - - # Handle single_transformer_blocks - single_transformer_blocks = transformer.get("single_transformer_blocks", []) - for i, weights in enumerate(single_transformer_blocks): - LoRALayers.set_single_transformer_block( - single_transformer_block=transformer_module.single_transformer_blocks[i], - dictionary=weights, - ) diff --git a/src/mflux/models/flux/weights/weight_util.py b/src/mflux/models/flux/weights/weight_util.py deleted file mode 100644 index 17722c3..0000000 --- a/src/mflux/models/flux/weights/weight_util.py +++ /dev/null @@ -1,118 +0,0 @@ -from typing import TYPE_CHECKING - -import mlx.nn as nn - -from mflux.config.config import Config -from mflux.models.common.quantization.quantization_util import QuantizationUtil - -if TYPE_CHECKING: - from mflux.models.flux.variants.controlnet.weight_handler_controlnet import WeightHandlerControlnet - from mflux.models.flux.variants.redux.weight_handler_redux import WeightHandlerRedux - from mflux.models.flux.weights.weight_handler import WeightHandler - - -class WeightUtil: - @staticmethod - def flatten(params): - return [(k, v) for p in params for (k, v) in p] - - @staticmethod - def reshape_weights(key, value): - if len(value.shape) == 4: - value = value.transpose(0, 2, 3, 1) - value = value.reshape(-1).reshape(value.shape).astype(Config.precision) - return [(key, value)] - - @staticmethod - def set_weights_and_quantize( - quantize_arg: int | None, - weights: "WeightHandler", - vae: nn.Module, - transformer: nn.Module, - t5_text_encoder: nn.Module, - clip_text_encoder: nn.Module, - ) -> int | None: - if weights.meta_data.quantization_level is None and quantize_arg is None: - WeightUtil._set_model_weights(weights, vae, transformer, t5_text_encoder, clip_text_encoder) - return None - - if weights.meta_data.quantization_level is None and quantize_arg is not None: - bits = quantize_arg - WeightUtil._set_model_weights(weights, vae, transformer, t5_text_encoder, clip_text_encoder) - QuantizationUtil.quantize_model(vae, transformer, t5_text_encoder, clip_text_encoder, bits, weights) # fmt:off - return bits - - if weights.meta_data.quantization_level is not None: - bits = weights.meta_data.quantization_level - QuantizationUtil.quantize_model(vae, transformer, t5_text_encoder, clip_text_encoder, bits, weights) # fmt:off - WeightUtil._set_model_weights(weights, vae, transformer, t5_text_encoder, clip_text_encoder) - return bits - - raise Exception("Error setting weights") - - @staticmethod - def set_controlnet_weights_and_quantize( - quantize_arg: int | None, - weights: "WeightHandlerControlnet", - transformer_controlnet: nn.Module, - ) -> int | None: - if weights.meta_data.quantization_level is None and quantize_arg is None: - transformer_controlnet.update(weights.controlnet_transformer, strict=False) - return None - - if weights.meta_data.quantization_level is None and quantize_arg is not None: - bits = quantize_arg - transformer_controlnet.update(weights.controlnet_transformer, strict=False) - QuantizationUtil.quantize_controlnet(bits, weights, transformer_controlnet) - return bits - - if weights.meta_data.quantization_level is not None: - bits = weights.meta_data.quantization_level - QuantizationUtil.quantize_controlnet(bits, weights, transformer_controlnet) - transformer_controlnet.update(weights.controlnet_transformer, strict=False) - return bits - - @staticmethod - def _set_model_weights( - weights: "WeightHandler", - vae: nn.Module, - transformer: nn.Module, - t5_text_encoder: nn.Module, - clip_text_encoder: nn.Module, - ): - vae.update(weights.vae, strict=False) - transformer.update(weights.transformer, strict=False) - t5_text_encoder.update(weights.t5_encoder, strict=False) - clip_text_encoder.update(weights.clip_encoder, strict=False) - - @staticmethod - def _set_redux_model_weights( - weights: "WeightHandlerRedux", - redux_encoder: nn.Module, - siglip_vision_transformer: nn.Module, - ): - redux_encoder.update(weights.redux_encoder, strict=False) - siglip_vision_transformer.update(weights.siglip["vision_model"], strict=False) - - @staticmethod - def set_redux_weights_and_quantize( - quantize_arg: int | None, - weights: "WeightHandlerRedux", - redux_encoder: nn.Module, - siglip_vision_transformer: nn.Module, - ) -> int | None: - if weights.meta_data.quantization_level is None and quantize_arg is None: - WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer) - return None - - if weights.meta_data.quantization_level is None and quantize_arg is not None: - bits = quantize_arg - WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer) - QuantizationUtil.quantize_redux_models(bits, weights, redux_encoder, siglip_vision_transformer) - return bits - - if weights.meta_data.quantization_level is not None: - bits = weights.meta_data.quantization_level - QuantizationUtil.quantize_redux_models(bits, weights, redux_encoder, siglip_vision_transformer) - WeightUtil._set_redux_model_weights(weights, redux_encoder, siglip_vision_transformer) - return bits diff --git a/src/mflux/models/qwen/cli/__init__.py b/src/mflux/models/qwen/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/generate_qwen_edit.py b/src/mflux/models/qwen/cli/qwen_image_edit_generate.py similarity index 59% rename from src/mflux/generate_qwen_edit.py rename to src/mflux/models/qwen/cli/qwen_image_edit_generate.py index a4c22fd..68f59c9 100644 --- a/src/mflux/generate_qwen_edit.py +++ b/src/mflux/models/qwen/cli/qwen_image_edit_generate.py @@ -1,12 +1,12 @@ from pathlib import Path from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -16,13 +16,7 @@ def main(): parser.add_model_arguments(require_model_arg=False) parser.add_lora_arguments() parser.add_image_generator_arguments(supports_metadata_config=True) - parser.add_argument( - "--image-paths", - type=Path, - nargs="+", - required=True, - help="Local paths to one or more init images. For single image editing, provide one path. For multiple image editing, provide multiple paths.", - ) + parser.add_argument("--image-paths", type=Path, nargs="+", required=True, help="Local paths to one or more init images. For single image editing, provide one path. For multiple image editing, provide multiple paths.") # fmt: off parser.add_output_arguments() args = parser.parse_args() @@ -33,36 +27,34 @@ def main(): # 1. Load the model qwen = QwenImageEdit( quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=qwen) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=qwen, + latent_creator=QwenLatentCreator, + ) try: for seed in args.seed: # 3. Prepare image paths image_paths = [str(p) for p in args.image_paths] - # Use first image path for config.image_path (for backward compatibility with Config) - # All image paths are passed to generate_image and stored in metadata - config_image_path = image_paths[0] - # 4. Generate an image for each seed value image = qwen.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=config_image_path, - ), - negative_prompt=PromptUtils.get_effective_negative_prompt(args), + prompt=PromptUtil.read_prompt(args), + negative_prompt=PromptUtil.read_negative_prompt(args), + width=args.width, + height=args.height, + guidance=args.guidance, + image_path=image_paths[0], # Use first image for metadata image_paths=image_paths, + num_inference_steps=args.steps, ) # 5. Save the image diff --git a/src/mflux/generate_qwen.py b/src/mflux/models/qwen/cli/qwen_image_generate.py similarity index 55% rename from src/mflux/generate_qwen.py rename to src/mflux/models/qwen/cli/qwen_image_generate.py index 623e5f5..d9d461e 100644 --- a/src/mflux/generate_qwen.py +++ b/src/mflux/models/qwen/cli/qwen_image_generate.py @@ -1,11 +1,11 @@ from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.utils.dimension_resolver import DimensionResolver from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil def main(): @@ -14,7 +14,7 @@ def main(): parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) parser.add_lora_arguments() - parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_generator_arguments(supports_metadata_config=True, supports_dimension_scale_factor=True) parser.add_image_to_image_arguments(required=False) parser.add_output_arguments() args = parser.parse_args() @@ -25,32 +25,40 @@ def main(): # 1. Load the model qwen = QwenImage( - model_config=ModelConfig.from_name(model_name=args.model or "qwen-image", base_model=args.base_model), quantize=args.quantize, - local_path=args.path, + model_path=args.model_path, lora_paths=args.lora_paths, lora_scales=args.lora_scales, ) # 2. Register callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=qwen) + memory_saver = CallbackManager.register_callbacks( + args=args, + model=qwen, + latent_creator=QwenLatentCreator, + ) try: + # Resolve dimensions (supports ScaleFactor like "2x" when --image-path is provided) + width, height = DimensionResolver.resolve( + width=args.width, + height=args.height, + reference_image_path=args.image_path, + ) + for seed in args.seed: # 3. Generate an image for each seed value image = qwen.generate_image( seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - negative_prompt=PromptUtils.get_effective_negative_prompt(args), - config=Config( - num_inference_steps=args.steps, - height=args.height, - width=args.width, - guidance=args.guidance, - image_path=args.image_path, - image_strength=args.image_strength, - scheduler=args.scheduler, - ), + prompt=PromptUtil.read_prompt(args), + negative_prompt=PromptUtil.read_negative_prompt(args), + width=width, + height=height, + guidance=args.guidance, + scheduler=args.scheduler, + image_path=args.image_path, + num_inference_steps=args.steps, + image_strength=args.image_strength, ) # 4. Save the image image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) diff --git a/src/mflux/models/qwen/latent_creator/qwen_latent_creator.py b/src/mflux/models/qwen/latent_creator/qwen_latent_creator.py index 52e4a3f..6ee4bc5 100644 --- a/src/mflux/models/qwen/latent_creator/qwen_latent_creator.py +++ b/src/mflux/models/qwen/latent_creator/qwen_latent_creator.py @@ -9,5 +9,9 @@ class QwenLatentCreator: return FluxLatentCreator.create_noise(seed, height, width) @staticmethod - def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: - return FluxLatentCreator.pack_latents(latents, height, width) + def pack_latents(latents: mx.array, height: int, width: int, num_channels_latents: int = 16) -> mx.array: + return FluxLatentCreator.pack_latents(latents, height, width, num_channels_latents) + + @staticmethod + def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: + return FluxLatentCreator.unpack_latents(latents, height, width) diff --git a/src/mflux/models/qwen/model/qwen_text_encoder/qwen_prompt_encoder.py b/src/mflux/models/qwen/model/qwen_text_encoder/qwen_prompt_encoder.py index 5eb92cd..29508e5 100644 --- a/src/mflux/models/qwen/model/qwen_text_encoder/qwen_prompt_encoder.py +++ b/src/mflux/models/qwen/model/qwen_text_encoder/qwen_prompt_encoder.py @@ -1,7 +1,7 @@ import mlx.core as mx +from mflux.models.common.tokenizer import Tokenizer from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder -from mflux.models.qwen.tokenizer.qwen_tokenizer import TokenizerQwen class QwenPromptEncoder: @@ -10,7 +10,7 @@ class QwenPromptEncoder: prompt: str, negative_prompt: str, prompt_cache: dict[str, tuple[mx.array, mx.array, mx.array, mx.array]], - qwen_tokenizer: TokenizerQwen, + qwen_tokenizer: Tokenizer, qwen_text_encoder: QwenTextEncoder, ) -> tuple[mx.array, mx.array, mx.array, mx.array]: # 0. Create a cache key that combines both prompts @@ -21,11 +21,13 @@ class QwenPromptEncoder: return prompt_cache[cache_key] # 2. Encode the positive prompt - input_ids, attention_mask = qwen_tokenizer.tokenize(prompt) - neg_input_ids, neg_attention_mask = qwen_tokenizer.tokenize(negative_prompt) - prompt_embeds, prompt_mask = qwen_text_encoder(input_ids=input_ids, attention_mask=attention_mask) + pos_output = qwen_tokenizer.tokenize(prompt) + neg_output = qwen_tokenizer.tokenize(negative_prompt) + prompt_embeds, prompt_mask = qwen_text_encoder( + input_ids=pos_output.input_ids, attention_mask=pos_output.attention_mask + ) neg_prompt_embeds, neg_prompt_mask = qwen_text_encoder( - input_ids=neg_input_ids, attention_mask=neg_attention_mask + input_ids=neg_output.input_ids, attention_mask=neg_output.attention_mask ) # 3. Cache the result (all 4 values) diff --git a/src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_prompt_encoder.py b/src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_prompt_encoder.py index 33ec1ba..3dc7c7f 100644 --- a/src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_prompt_encoder.py +++ b/src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_prompt_encoder.py @@ -4,9 +4,10 @@ import mlx.core as mx import numpy as np from PIL import Image +from mflux.models.common.tokenizer import Tokenizer +from mflux.models.qwen.model.qwen_text_encoder.qwen_prompt_encoder import QwenPromptEncoder from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_language_encoder import QwenVisionLanguageEncoder -from mflux.models.qwen.tokenizer.qwen_tokenizer import TokenizerQwen from mflux.models.qwen.tokenizer.qwen_vision_language_tokenizer import QwenVisionLanguageTokenizer @@ -59,7 +60,7 @@ class QwenVisionLanguagePromptEncoder: negative_prompt: Optional[str] = None, image: Optional[Union[Image.Image, np.ndarray, str]] = None, prompt_cache: Optional[dict] = None, - qwen_tokenizer: Optional[TokenizerQwen] = None, + qwen_tokenizer: Optional[Tokenizer] = None, qwen_text_encoder: Optional[QwenTextEncoder] = None, qwen_vl_tokenizer: Optional[QwenVisionLanguageTokenizer] = None, qwen_vl_encoder: Optional[QwenVisionLanguageEncoder] = None, @@ -84,7 +85,7 @@ class QwenVisionLanguagePromptEncoder: "Text-only components (qwen_tokenizer, qwen_text_encoder) are required when no image is provided" ) - return QwenVisionLanguagePromptEncoder.encode_prompt( + return QwenPromptEncoder.encode_prompt( prompt=prompt, negative_prompt=negative_prompt or "", prompt_cache=prompt_cache, diff --git a/src/mflux/models/qwen/model/qwen_transformer/qwen_transformer.py b/src/mflux/models/qwen/model/qwen_transformer/qwen_transformer.py index 740eded..3041d67 100644 --- a/src/mflux/models/qwen/model/qwen_transformer/qwen_transformer.py +++ b/src/mflux/models/qwen/model/qwen_transformer/qwen_transformer.py @@ -4,7 +4,7 @@ import mlx.core as mx import numpy as np from mlx import nn -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config from mflux.models.flux.model.flux_transformer.ada_layer_norm_continuous import AdaLayerNormContinuous from mflux.models.qwen.model.qwen_transformer.qwen_rope import QwenEmbedRopeMLX from mflux.models.qwen.model.qwen_transformer.qwen_time_text_embed import QwenTimeTextEmbed @@ -37,7 +37,7 @@ class QwenTransformer(nn.Module): def __call__( self, t: int, - config: RuntimeConfig, + config: Config, hidden_states: mx.array, encoder_hidden_states: mx.array, encoder_hidden_states_mask: mx.array, @@ -93,11 +93,8 @@ class QwenTransformer(nn.Module): @staticmethod def _compute_timestep( t: int | float, - config: RuntimeConfig, + config: Config, ) -> mx.array: - """ - Compute timestep tensor from step index or value. - """ if isinstance(t, int): if t < len(config.scheduler.sigmas): timestep_idx = t @@ -123,7 +120,7 @@ class QwenTransformer(nn.Module): def _compute_rotary_embeddings( encoder_hidden_states_mask: mx.array, pos_embed: QwenEmbedRopeMLX, - config: RuntimeConfig, + config: Config, cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None, ) -> tuple[mx.array, mx.array]: latent_height = config.height // 16 diff --git a/src/mflux/models/qwen/qwen_edit_initializer.py b/src/mflux/models/qwen/qwen_edit_initializer.py deleted file mode 100644 index eb019e5..0000000 --- a/src/mflux/models/qwen/qwen_edit_initializer.py +++ /dev/null @@ -1,103 +0,0 @@ -from pathlib import Path - -from transformers import Qwen2TokenizerFast - -from mflux.config.model_config import ModelConfig -from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_language_encoder import QwenVisionLanguageEncoder -from mflux.models.qwen.qwen_initializer import QwenImageInitializer -from mflux.models.qwen.tokenizer.qwen_vision_language_processor import QwenVisionLanguageProcessor -from mflux.models.qwen.tokenizer.qwen_vision_language_tokenizer import QwenVisionLanguageTokenizer -from mflux.utils.download import snapshot_download - - -class QwenImageEditInitializer: - @staticmethod - def init( - qwen_model, - model_config: ModelConfig, - quantize: int | None, - local_path: str | None, - lora_paths: list[str] | None = None, - lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, - ) -> None: - # 1. Initialize the base Qwen Image model (VAE, transformer, text encoder, etc.) - QwenImageInitializer.init( - qwen_model=qwen_model, - model_config=model_config, - quantize=quantize, - local_path=local_path, - lora_paths=lora_paths, - lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, - ) - - # 2. Add vision-language components for edit functionality - QwenImageEditInitializer._init_vision_language_components( - qwen_model=qwen_model, - repo_id=model_config.model_name, - local_path=local_path, - model_config=model_config, - ) - - @staticmethod - def _init_vision_language_components( - qwen_model, - repo_id: str, - local_path: str | None, - model_config: ModelConfig | None = None, - ) -> None: - # 1. Download or get cached tokenizer - root_path = Path(local_path) if local_path else QwenImageEditInitializer._download_vl_processor(repo_id) - - # 2. Load only the tokenizer (we implement image processing ourselves) - tokenizer = QwenImageEditInitializer._load_tokenizer(root_path, repo_id) - - # 3. Create our MLX processor with the tokenizer - processor = QwenVisionLanguageProcessor(tokenizer=tokenizer) - - # 4. Initialize vision-language tokenizer wrapper - qwen_model.qwen_vl_tokenizer = QwenVisionLanguageTokenizer( - processor=processor, - max_length=1024, - use_picture_prefix=True, - ) - - # 5. Initialize vision-language encoder (integrated approach like Diffusers) - qwen_model.qwen_vl_encoder = QwenVisionLanguageEncoder(encoder=qwen_model.text_encoder.encoder) - - @staticmethod - def _load_tokenizer(root_path: Path, repo_id: str) -> Qwen2TokenizerFast: - """Load the tokenizer from local path or download from HuggingFace.""" - tokenizer_path = root_path / "tokenizer" - if not tokenizer_path.exists(): - tokenizer_path = root_path - - try: - return Qwen2TokenizerFast.from_pretrained( - pretrained_model_name_or_path=tokenizer_path, - local_files_only=True, - ) - except OSError: - return Qwen2TokenizerFast.from_pretrained(repo_id) - - @staticmethod - def _download_vl_processor(repo_id: str) -> Path: - return Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=[ - "processor/**", - "preprocessor_config.json", - "tokenizer/**", - "added_tokens.json", - "chat_template.jinja", - "tokenizer.json", - "tokenizer_config.json", - "vocab.json", - "merges.txt", - ], - ) - ) diff --git a/src/mflux/models/qwen/qwen_initializer.py b/src/mflux/models/qwen/qwen_initializer.py index 99718a2..a20e4fe 100644 --- a/src/mflux/models/qwen/qwen_initializer.py +++ b/src/mflux/models/qwen/qwen_initializer.py @@ -1,70 +1,118 @@ -from mflux.config.model_config import ModelConfig -from mflux.models.common.lora.download.lora_huggingface_downloader import LoRAHuggingFaceDownloader +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.models.common.config import ModelConfig from mflux.models.common.lora.mapping.lora_loader import LoRALoader +from mflux.models.common.tokenizer import TokenizerLoader +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_loader import WeightLoader from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder +from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_language_encoder import QwenVisionLanguageEncoder +from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_transformer import VisionTransformer from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE -from mflux.models.qwen.tokenizer.qwen_tokenizer_handler import QwenTokenizerHandler +from mflux.models.qwen.tokenizer.qwen_vision_language_processor import QwenVisionLanguageProcessor +from mflux.models.qwen.tokenizer.qwen_vision_language_tokenizer import QwenVisionLanguageTokenizer from mflux.models.qwen.weights.qwen_lora_mapping import QwenLoRAMapping -from mflux.models.qwen.weights.qwen_weight_handler import QwenWeightHandler -from mflux.models.qwen.weights.qwen_weight_util import QwenWeightUtil +from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition class QwenImageInitializer: @staticmethod def init( - qwen_model, + model, model_config: ModelConfig, quantize: int | None, - local_path: str | None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, ) -> None: - # 0. Set paths, configs, and prompt_cache for later - qwen_model.prompt_cache = {} - qwen_model.model_config = model_config + path = model_path if model_path else model_config.model_name + QwenImageInitializer._init_config(model, model_config) + weights = QwenImageInitializer._load_weights(path) + QwenImageInitializer._init_tokenizers(model, path) + QwenImageInitializer._init_models(model) + QwenImageInitializer._apply_weights(model, weights, quantize) + QwenImageInitializer._apply_lora(model, lora_paths, lora_scales) - # 1. Load the regular weights - weights = QwenWeightHandler.load_regular_weights( - repo_id=model_config.model_name, - local_path=local_path, + @staticmethod + def init_edit( + model, + model_config: ModelConfig, + quantize: int | None, + model_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + ) -> None: + # Use model_path if provided, otherwise fall back to model_config.model_name + path = model_path if model_path else model_config.model_name + QwenImageInitializer._init_config(model, model_config) + weights = QwenImageInitializer._load_weights(path) + QwenImageInitializer._init_tokenizers(model, path) + QwenImageInitializer._init_edit_models(model) + QwenImageInitializer._apply_weights(model, weights, quantize) + QwenImageInitializer._apply_lora(model, lora_paths, lora_scales) + + # Add vision-language tokenizer + raw_tokenizer = model.tokenizers["qwen"].tokenizer + processor = QwenVisionLanguageProcessor(tokenizer=raw_tokenizer) + model.tokenizers["qwen_vl"] = QwenVisionLanguageTokenizer( + processor=processor, + max_length=1024, + use_picture_prefix=True, + ) + model.qwen_vl_encoder = QwenVisionLanguageEncoder(encoder=model.text_encoder.encoder) + + @staticmethod + def _init_config(model, model_config: ModelConfig) -> None: + model.prompt_cache = {} + model.model_config = model_config + model.callbacks = CallbackRegistry() + + @staticmethod + def _load_weights(model_path: str) -> LoadedWeights: + return WeightLoader.load( + weight_definition=QwenWeightDefinition, + model_path=model_path, ) - # 2. Initialize tokenizers - tokenizer_handler = QwenTokenizerHandler( - repo_id=model_config.model_name, - local_path=local_path, + @staticmethod + def _init_tokenizers(model, model_path: str) -> None: + model.tokenizers = TokenizerLoader.load_all( + definitions=QwenWeightDefinition.get_tokenizers(), + model_path=model_path, ) - qwen_model.qwen_tokenizer = tokenizer_handler.qwen - # 3. Initialize all models - qwen_model.vae = QwenVAE() - qwen_model.transformer = QwenTransformer() - qwen_model.text_encoder = QwenTextEncoder() + @staticmethod + def _init_models(model) -> None: + model.vae = QwenVAE() + model.transformer = QwenTransformer() + model.text_encoder = QwenTextEncoder() - # 4. Apply weights and quantize the models - qwen_model.bits = QwenWeightUtil.set_weights_and_quantize( - quantize_arg=quantize, + @staticmethod + def _init_edit_models(model) -> None: + model.vae = QwenVAE() + model.transformer = QwenTransformer() + model.text_encoder = QwenTextEncoder() + model.text_encoder.encoder.visual = VisionTransformer() + + @staticmethod + def _apply_weights(model, weights: LoadedWeights, quantize: int | None) -> None: + model.bits = WeightApplier.apply_and_quantize( weights=weights, - vae=qwen_model.vae, - transformer=qwen_model.transformer, - text_encoder=qwen_model.text_encoder, + quantize_arg=quantize, + weight_definition=QwenWeightDefinition, + models={ + "vae": model.vae, + "transformer": model.transformer, + "text_encoder": model.text_encoder, + }, ) - # 5. Set LoRA weights - hf_lora_paths = LoRAHuggingFaceDownloader.download_loras( - lora_names=lora_names, - repo_id=lora_repo_id, - model_name="Qwen", + @staticmethod + def _apply_lora(model, lora_paths: list[str] | None, lora_scales: list[float] | None) -> None: + model.lora_paths, model.lora_scales = LoRALoader.load_and_apply_lora( + lora_mapping=QwenLoRAMapping.get_mapping(), + transformer=model.transformer, + lora_paths=lora_paths, + lora_scales=lora_scales, ) - qwen_model.lora_paths = (lora_paths or []) + hf_lora_paths - qwen_model.lora_scales = (lora_scales or []) + [1.0] * len(hf_lora_paths) - if qwen_model.lora_paths: - LoRALoader.load_and_apply_lora( - lora_mapping=QwenLoRAMapping.get_mapping(), - transformer=qwen_model.transformer, - lora_files=qwen_model.lora_paths, - lora_scales=qwen_model.lora_scales, - ) diff --git a/src/mflux/models/qwen/tokenizer/__init__.py b/src/mflux/models/qwen/tokenizer/__init__.py index e69de29..9982088 100644 --- a/src/mflux/models/qwen/tokenizer/__init__.py +++ b/src/mflux/models/qwen/tokenizer/__init__.py @@ -0,0 +1,2 @@ +# Qwen Tokenizer - uses unified tokenizer system from mflux.models.common.tokenizer + diff --git a/src/mflux/models/qwen/tokenizer/qwen_tokenizer.py b/src/mflux/models/qwen/tokenizer/qwen_tokenizer.py deleted file mode 100644 index a3589c1..0000000 --- a/src/mflux/models/qwen/tokenizer/qwen_tokenizer.py +++ /dev/null @@ -1,23 +0,0 @@ -import mlx.core as mx -from transformers import Qwen2Tokenizer - - -class TokenizerQwen: - def __init__(self, tokenizer: Qwen2Tokenizer, max_length: int = 1024): - self.tokenizer = tokenizer - self.max_length = max_length - self.prompt_template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" - self.template_start_idx = 34 - - def tokenize(self, prompt: str) -> tuple[mx.array, mx.array]: - formatted_text = self.prompt_template.format(prompt) - tokens = self.tokenizer( - formatted_text, - max_length=self.max_length + self.template_start_idx, - padding=False, - truncation=True, - return_tensors="np", - ) - input_ids = mx.array(tokens["input_ids"]) - attention_mask = mx.array(tokens["attention_mask"]) - return input_ids, attention_mask diff --git a/src/mflux/models/qwen/tokenizer/qwen_tokenizer_handler.py b/src/mflux/models/qwen/tokenizer/qwen_tokenizer_handler.py deleted file mode 100644 index 848df74..0000000 --- a/src/mflux/models/qwen/tokenizer/qwen_tokenizer_handler.py +++ /dev/null @@ -1,39 +0,0 @@ -from pathlib import Path - -import transformers - -from mflux.models.qwen.tokenizer.qwen_tokenizer import TokenizerQwen -from mflux.utils.download import snapshot_download - - -class QwenTokenizerHandler: - def __init__( - self, - repo_id: str, - max_length: int = 1024, - local_path: str | None = None, - ): - root_path = Path(local_path) if local_path else QwenTokenizerHandler._download_or_get_cached_tokenizer(repo_id) - - # Load the Qwen2 tokenizer - tokenizer_path = root_path / "tokenizer" - self.tokenizer_raw = transformers.Qwen2Tokenizer.from_pretrained( - pretrained_model_name_or_path=tokenizer_path, - local_files_only=True, - ) - - # Wrap in our TokenizerQwen class - self.qwen = TokenizerQwen(self.tokenizer_raw, max_length) - - @staticmethod - def _download_or_get_cached_tokenizer(repo_id: str) -> Path: - return Path( - snapshot_download( - repo_id=repo_id, - allow_patterns=[ - "tokenizer/**", - "added_tokens.json", - "chat_template.jinja", - ], - ) - ) diff --git a/src/mflux/models/qwen/variants/edit/utils/qwen_edit_util.py b/src/mflux/models/qwen/variants/edit/qwen_edit_util.py similarity index 95% rename from src/mflux/models/qwen/variants/edit/utils/qwen_edit_util.py rename to src/mflux/models/qwen/variants/edit/qwen_edit_util.py index e362e1d..665b41a 100644 --- a/src/mflux/models/qwen/variants/edit/utils/qwen_edit_util.py +++ b/src/mflux/models/qwen/variants/edit/qwen_edit_util.py @@ -3,7 +3,7 @@ import os import mlx.core as mx from mflux.models.common.latent_creator.latent_creator import LatentCreator -from mflux.utils.array_util import ArrayUtil +from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator class QwenEditUtil: @@ -41,7 +41,7 @@ class QwenEditUtil: width=calc_w, ) - image_latents = ArrayUtil.pack_latents( + image_latents = QwenLatentCreator.pack_latents( latents=input_image, height=calc_h, width=calc_w, diff --git a/src/mflux/models/qwen/variants/edit/qwen_image_edit.py b/src/mflux/models/qwen/variants/edit/qwen_image_edit.py index 0d0689b..84bd26d 100644 --- a/src/mflux/models/qwen/variants/edit/qwen_image_edit.py +++ b/src/mflux/models/qwen/variants/edit/qwen_image_edit.py @@ -1,21 +1,19 @@ import math +from pathlib import Path import mlx.core as mx from mlx import nn from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE -from mflux.models.qwen.qwen_edit_initializer import QwenImageEditInitializer -from mflux.models.qwen.variants.edit.utils.qwen_edit_util import QwenEditUtil +from mflux.models.qwen.qwen_initializer import QwenImageInitializer +from mflux.models.qwen.variants.edit.qwen_edit_util import QwenEditUtil from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage -from mflux.utils.array_util import ArrayUtil from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -29,83 +27,86 @@ class QwenImageEdit(nn.Module): def __init__( self, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, + model_config: ModelConfig = ModelConfig.qwen_image_edit(), ): super().__init__() - QwenImageEditInitializer.init( - qwen_model=self, - model_config=ModelConfig.qwen_image_edit(), + QwenImageInitializer.init_edit( + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + image_paths: list[str], + num_inference_steps: int = 4, + height: int | None = None, + width: int | None = None, + guidance: float = 4.0, + image_path: Path | str | None = None, + scheduler: str = "linear", negative_prompt: str | None = None, - image_paths: list[str] | None = None, ) -> GeneratedImage: - if image_paths is None: - image_paths = [str(config.image_path)] - - runtime_config, vl_width, vl_height, vae_width, vae_height = self._compute_dimensions(config, image_paths) - timesteps = runtime_config.scheduler.timesteps + config, vl_width, vl_height, vae_width, vae_height = self._compute_dimensions( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_paths=image_paths, + num_inference_steps=num_inference_steps, + ) + timesteps = config.scheduler.timesteps time_steps = tqdm(range(len(timesteps))) # 1. Create initial latents latents = QwenLatentCreator.create_noise( seed=seed, - height=runtime_config.height, - width=runtime_config.width, + width=config.width, + height=config.height, ) # 2. Encode the prompt prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = self._encode_prompts_with_images( prompt=prompt, - negative_prompt=negative_prompt, - image_paths=image_paths, - runtime_config=runtime_config, + config=config, vl_width=vl_width, vl_height=vl_height, + image_paths=image_paths, + negative_prompt=negative_prompt, ) # 3. Generate image conditioning latents static_image_latents, qwen_image_ids, cond_h_patches, cond_w_patches, num_images = ( QwenEditUtil.create_image_conditioning_latents( vae=self.vae, - height=vae_height, width=vae_width, - image_paths=image_paths, + height=vae_height, vl_width=vl_width, vl_height=vl_height, + image_paths=image_paths, ) ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - ) + # 4. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) for t in time_steps: try: - # 4.t Concatenate the updated latents with the static image latents + # 5.t Concatenate the updated latents with the static image latents hidden_states = mx.concatenate([latents, static_image_latents], axis=1) hidden_states_neg = mx.concatenate([latents, static_image_latents], axis=1) - # 5.t Predict the noise + # 6.t Predict the noise if num_images > 1: cond_image_grid = [(1, cond_h_patches, cond_w_patches) for _ in range(num_images)] else: @@ -113,7 +114,7 @@ class QwenImageEdit(nn.Module): noise = self.transformer( t=t, - config=runtime_config, + config=config, hidden_states=hidden_states, encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_mask, @@ -122,66 +123,43 @@ class QwenImageEdit(nn.Module): )[:, : latents.shape[1]] noise_negative = self.transformer( t=t, - config=runtime_config, + config=config, hidden_states=hidden_states_neg, encoder_hidden_states=negative_prompt_embeds, encoder_hidden_states_mask=negative_prompt_mask, qwen_image_ids=qwen_image_ids, cond_image_grid=cond_image_grid, )[:, : latents.shape[1]] - guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, runtime_config.guidance) + guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, config.guidance) - # 6.t Take one denoise step - latents = runtime_config.scheduler.step( - model_output=guided_noise, - timestep=t, - sample=latents, - ) + # 7.t Take one denoise step + latents = config.scheduler.step(noise=guided_noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, - ) + # 8.t Call subscribers in-loop + ctx.in_loop(t, latents, 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=runtime_config, - time_steps=time_steps, - ) + ctx.interruption(t, latents, time_steps=time_steps) raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(timesteps)}") - # (Optional) Call subscribers after loop - Callbacks.after_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - ) + # 9. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width) + # 10. Decode the latent array and return the image + latents = QwenLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, - config=runtime_config, + config=config, seed=seed, prompt=prompt, quantization=self.bits, lora_paths=self.lora_paths, lora_scales=self.lora_scales, - image_path=runtime_config.image_path, + image_path=config.image_path, image_paths=image_paths, generation_time=time_steps.format_dict["elapsed"], negative_prompt=negative_prompt, @@ -192,11 +170,11 @@ class QwenImageEdit(nn.Module): prompt: str, negative_prompt: str, image_paths: list[str], - runtime_config, + config, vl_width: int | None = None, vl_height: int | None = None, ) -> tuple[mx.array, mx.array, mx.array, mx.array]: - tokenizer = self.qwen_vl_tokenizer + tokenizer = self.tokenizers["qwen_vl"] pos_input_ids, pos_attention_mask, pos_pixel_values, pos_image_grid_thw = tokenizer.tokenize_with_image( prompt, image_paths, vl_width=vl_width, vl_height=vl_height ) @@ -207,8 +185,6 @@ class QwenImageEdit(nn.Module): pixel_values=pos_pixel_values, image_grid_thw=pos_image_grid_thw, ) - mx.eval(pos_hidden_states[0]) - mx.eval(pos_hidden_states[1]) neg_prompt = negative_prompt if negative_prompt is not None else "" @@ -223,8 +199,6 @@ class QwenImageEdit(nn.Module): pixel_values=neg_pixel_values, image_grid_thw=neg_image_grid_thw, ) - mx.eval(neg_hidden_states[0]) - mx.eval(neg_hidden_states[1]) final_prompt_embeds = pos_hidden_states[0].astype(mx.float16) final_prompt_mask = pos_hidden_states[1].astype(mx.float16) @@ -238,9 +212,14 @@ class QwenImageEdit(nn.Module): def _compute_dimensions( self, - config: Config, image_paths: list[str], - ) -> tuple[RuntimeConfig, int, int, int, int]: + num_inference_steps: int, + height: int | None, + width: int | None, + guidance: float, + image_path: Path | str | None, + scheduler: str, + ) -> tuple[Config, int, int, int, int]: last_image = ImageUtil.load_image(image_paths[-1]).convert("RGB") image_size = last_image.size @@ -251,23 +230,23 @@ class QwenImageEdit(nn.Module): calculated_width = round(calculated_width / 32) * 32 calculated_height = round(calculated_height / 32) * 32 - use_height = config.height or calculated_height - use_width = config.width or calculated_width + use_height = height or int(calculated_height) + use_width = width or int(calculated_width) vae_scale_factor = 8 multiple_of = vae_scale_factor * 2 use_width = use_width // multiple_of * multiple_of use_height = use_height // multiple_of * multiple_of - final_config = Config( - height=use_height, + config = Config( width=use_width, - image_path=config.image_path, - num_inference_steps=config.num_inference_steps, - guidance=config.guidance, - scheduler=config.scheduler_str, + height=use_height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + model_config=self.model_config, + num_inference_steps=num_inference_steps, ) - runtime_config = RuntimeConfig(final_config, self.model_config) CONDITION_IMAGE_SIZE = 384 * 384 condition_ratio = image_size[0] / image_size[1] @@ -283,4 +262,4 @@ class QwenImageEdit(nn.Module): vae_width = round(vae_width / 32) * 32 vae_height = round(vae_height / 32) * 32 - return runtime_config, int(vl_width), int(vl_height), int(vae_width), int(vae_height) + return config, int(vl_width), int(vl_height), int(vae_width), int(vae_height) diff --git a/src/mflux/models/qwen/variants/edit/utils/qwen_prompt_rewriter.py b/src/mflux/models/qwen/variants/edit/qwen_prompt_rewriter.py similarity index 100% rename from src/mflux/models/qwen/variants/edit/utils/qwen_prompt_rewriter.py rename to src/mflux/models/qwen/variants/edit/qwen_prompt_rewriter.py diff --git a/src/mflux/models/qwen/variants/edit/utils/__init__.py b/src/mflux/models/qwen/variants/edit/utils/__init__.py deleted file mode 100644 index 71bf0a9..0000000 --- a/src/mflux/models/qwen/variants/edit/utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Qwen Edit utilities diff --git a/src/mflux/models/qwen/variants/txt2img/qwen_image.py b/src/mflux/models/qwen/variants/txt2img/qwen_image.py index 5a9c0c8..3c1e828 100644 --- a/src/mflux/models/qwen/variants/txt2img/qwen_image.py +++ b/src/mflux/models/qwen/variants/txt2img/qwen_image.py @@ -1,20 +1,19 @@ +from pathlib import Path + import mlx.core as mx from mlx import nn -from tqdm import tqdm -from mflux.callbacks.callbacks import Callbacks -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config import ModelConfig +from mflux.models.common.config.config import Config from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator -from mflux.models.common.weights.model_saver import ModelSaver +from mflux.models.common.weights.saving.model_saver import ModelSaver from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.models.qwen.model.qwen_text_encoder.qwen_prompt_encoder import QwenPromptEncoder from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE from mflux.models.qwen.qwen_initializer import QwenImageInitializer -from mflux.utils.array_util import ArrayUtil +from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition from mflux.utils.exceptions import StopImageGenerationException from mflux.utils.generated_image import GeneratedImage from mflux.utils.image_util import ImageUtil @@ -27,48 +26,58 @@ class QwenImage(nn.Module): def __init__( self, - model_config: ModelConfig, quantize: int | None = None, - local_path: str | None = None, + model_path: str | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, + model_config: ModelConfig = ModelConfig.qwen_image(), ): super().__init__() QwenImageInitializer.init( - qwen_model=self, - model_config=model_config, + model=self, quantize=quantize, - local_path=local_path, + model_path=model_path, lora_paths=lora_paths, lora_scales=lora_scales, - lora_names=lora_names, - lora_repo_id=lora_repo_id, + model_config=model_config, ) def generate_image( self, seed: int, prompt: str, - config: Config, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + guidance: float = 4.0, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", negative_prompt: str | None = None, ) -> GeneratedImage: - # 0. Create a new runtime config based on the model type and input parameters - runtime_config = RuntimeConfig(config, self.model_config) - time_steps = tqdm(range(runtime_config.init_time_step, runtime_config.num_inference_steps)) + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=guidance, + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) # 1. Create the initial latents latents = LatentCreator.create_for_txt2img_or_img2img( seed=seed, - height=runtime_config.height, - width=runtime_config.width, + width=config.width, + height=config.height, img2img=Img2Img( vae=self.vae, latent_creator=QwenLatentCreator, - sigmas=runtime_config.scheduler.sigmas, - init_time_step=runtime_config.init_time_step, - image_path=runtime_config.image_path, + sigmas=config.scheduler.sigmas, + init_time_step=config.init_time_step, + image_path=config.image_path, ), ) @@ -77,93 +86,68 @@ class QwenImage(nn.Module): prompt=prompt, negative_prompt=negative_prompt, prompt_cache=self.prompt_cache, - qwen_tokenizer=self.qwen_tokenizer, + qwen_tokenizer=self.tokenizers["qwen"], qwen_text_encoder=self.text_encoder, ) - # (Optional) Call subscribers for beginning of loop - Callbacks.before_loop( - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - ) + # 3. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) - for t in time_steps: + for t in config.time_steps: try: # Scale model input if needed by the scheduler - latents = runtime_config.scheduler.scale_model_input(latents, t) + latents = config.scheduler.scale_model_input(latents, t) - # 3. Predict the noise + # 4. Predict the noise noise = self.transformer( t=t, - config=runtime_config, + config=config, hidden_states=latents, encoder_hidden_states=prompt_embeds, encoder_hidden_states_mask=prompt_mask, ) noise_negative = self.transformer( t=t, - config=runtime_config, + config=config, hidden_states=latents, encoder_hidden_states=negative_prompt_embeds, encoder_hidden_states_mask=negative_prompt_mask, ) - guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, runtime_config.guidance) + guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, config.guidance) - # 4.t Take one denoise step - latents = runtime_config.scheduler.step( - model_output=guided_noise, - timestep=t, - sample=latents, - ) + # 5.t Take one denoise step + latents = config.scheduler.step(noise=guided_noise, timestep=t, latents=latents) - # (Optional) Call subscribers in-loop - Callbacks.in_loop( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, - ) + # 6.t Call subscribers in-loop + ctx.in_loop(t, latents) # (Optional) Evaluate to enable progress tracking mx.eval(latents) except KeyboardInterrupt: # noqa: PERF203 - Callbacks.interruption( - t=t, - seed=seed, - prompt=prompt, - latents=latents, - config=runtime_config, - time_steps=time_steps, + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_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=runtime_config, - ) + # 7. Call subscribers after loop + ctx.after_loop(latents) - # 7. Decode the latent array and return the image - latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width) + # 8. Decode the latent array and return the image + latents = QwenLatentCreator.unpack_latents(latents=latents, height=config.height, width=config.width) decoded = self.vae.decode(latents) return ImageUtil.to_image( decoded_latents=decoded, - config=runtime_config, + config=config, seed=seed, prompt=prompt, quantization=self.bits, lora_paths=self.lora_paths, lora_scales=self.lora_scales, - image_path=runtime_config.image_path, - image_strength=runtime_config.image_strength, - generation_time=time_steps.format_dict["elapsed"], + image_path=config.image_path, + image_strength=config.image_strength, + generation_time=config.time_steps.format_dict["elapsed"], negative_prompt=negative_prompt, ) @@ -172,14 +156,7 @@ class QwenImage(nn.Module): model=self, bits=self.bits, base_path=base_path, - tokenizers=[ - ("qwen_tokenizer.tokenizer", "tokenizer"), - ], - components=[ - ("vae", "vae"), - ("transformer", "transformer"), - ("text_encoder", "text_encoder"), - ], + weight_definition=QwenWeightDefinition, ) @staticmethod diff --git a/src/mflux/models/qwen/weights/__init__.py b/src/mflux/models/qwen/weights/__init__.py index e69de29..70114f5 100644 --- a/src/mflux/models/qwen/weights/__init__.py +++ b/src/mflux/models/qwen/weights/__init__.py @@ -0,0 +1,4 @@ +from mflux.models.qwen.weights.qwen_weight_definition import QwenWeightDefinition +from mflux.models.qwen.weights.qwen_weight_mapping import QwenWeightMapping + +__all__ = ["QwenWeightDefinition", "QwenWeightMapping"] diff --git a/src/mflux/models/qwen/weights/qwen_weight_definition.py b/src/mflux/models/qwen/weights/qwen_weight_definition.py new file mode 100644 index 0000000..39d62b2 --- /dev/null +++ b/src/mflux/models/qwen/weights/qwen_weight_definition.py @@ -0,0 +1,63 @@ +from typing import List + +import mlx.core as mx + +from mflux.models.common.tokenizer import LanguageTokenizer +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition, TokenizerDefinition +from mflux.models.qwen.weights.qwen_weight_mapping import QwenWeightMapping + + +class QwenWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="vae", + hf_subdir="vae", + loading_mode="single", + mapping_getter=QwenWeightMapping.get_vae_mapping, + ), + ComponentDefinition( + name="transformer", + hf_subdir="transformer", + loading_mode="multi_glob", + mapping_getter=QwenWeightMapping.get_transformer_mapping, + ), + ComponentDefinition( + name="text_encoder", + hf_subdir="text_encoder", + loading_mode="multi_json", + precision=mx.bfloat16, + skip_quantization=True, # Quantization causes significant semantic degradation + mapping_getter=QwenWeightMapping.get_text_encoder_mapping, + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [ + TokenizerDefinition( + name="qwen", + hf_subdir="tokenizer", + tokenizer_class="Qwen2Tokenizer", + encoder_class=LanguageTokenizer, + max_length=1024, + template="<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n", + download_patterns=["tokenizer/**", "added_tokens.json", "chat_template.jinja"], + ), + ] + + @staticmethod + def get_download_patterns() -> List[str]: + return [ + "vae/*.safetensors", + "vae/*.json", + "transformer/*.safetensors", + "transformer/*.json", + "text_encoder/*.safetensors", + "text_encoder/*.json", + ] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + return hasattr(module, "to_quantized") diff --git a/src/mflux/models/qwen/weights/qwen_weight_handler.py b/src/mflux/models/qwen/weights/qwen_weight_handler.py deleted file mode 100644 index 7c06315..0000000 --- a/src/mflux/models/qwen/weights/qwen_weight_handler.py +++ /dev/null @@ -1,190 +0,0 @@ -import json -from collections.abc import Callable -from pathlib import Path - -import mlx.core as mx -import torch -from mlx.utils import tree_unflatten -from safetensors.mlx import load_file as mlx_load_file -from safetensors.torch import load_file as torch_load_file - -from mflux.models.common.weights.mapping.weight_mapper import WeightMapper -from mflux.models.common.weights.mapping.weight_mapping import WeightTarget -from mflux.models.flux.weights.weight_handler import MetaData, WeightHandler -from mflux.models.qwen.weights.qwen_weight_mapping import QwenWeightMapping - - -class QwenWeightHandler: - def __init__( - self, - meta_data: MetaData, - qwen_text_encoder: dict | None = None, - transformer: dict | None = None, - vae: dict | None = None, - ): - self.qwen_text_encoder = qwen_text_encoder - self.transformer = transformer - self.vae = vae - self.meta_data = meta_data - - @staticmethod - def load_regular_weights( - repo_id: str | None = None, - local_path: str | None = None, - ) -> "QwenWeightHandler": - # 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) - - # Determine if we should load visual weights (for Edit model) - load_visual_weights = repo_id and "edit" in repo_id.lower() - - # Load the weights - transformer, quantization_level, mflux_version = QwenWeightHandler._load_transformer(root_path=root_path) - qwen_text_encoder, _, _ = QwenWeightHandler._load_qwen_text_encoder(root_path=root_path, load_visual_weights=load_visual_weights) # fmt: off - vae, _, _ = QwenWeightHandler._load_vae(root_path=root_path) - - return QwenWeightHandler( - qwen_text_encoder=qwen_text_encoder, - transformer=transformer, - vae=vae, - meta_data=MetaData( - quantization_level=quantization_level, - scale=None, - is_lora=False, - mflux_version=mflux_version, - ), - ) - - @staticmethod - def _load_transformer(root_path: Path) -> tuple[dict, int | None, str | None]: - return QwenWeightHandler._load_component( - root_path=root_path, - component_name="transformer", - loading_mode="multi_glob", - mapping_getter=QwenWeightMapping.get_transformer_mapping, - ) - - @staticmethod - def _load_qwen_text_encoder( - root_path: Path, load_visual_weights: bool = False - ) -> tuple[dict, int | None, str | None]: - return QwenWeightHandler._load_component( - root_path=root_path, - component_name="text_encoder", - loading_mode="multi_json", - mapping_getter=QwenWeightMapping.get_text_encoder_mapping, - ) - - @staticmethod - def _load_vae(root_path: Path) -> tuple[dict, int | None, str | None]: - return QwenWeightHandler._load_component( - root_path=root_path, - component_name="vae", - loading_mode="single", - mapping_getter=QwenWeightMapping.get_vae_mapping, - ) - - @staticmethod - def _load_component( - root_path: Path, - component_name: str, - loading_mode: str, - mapping_getter: Callable[[], list[WeightTarget]], - ) -> tuple[dict, int | None, str | None]: - component_path = root_path / component_name - weights = QwenWeightHandler._load_safetensors_shards(component_path, loading_mode=loading_mode) - - # Check if this is a saved model (has metadata) - quantization_level, mflux_version = QwenWeightHandler._detect_metadata(component_path) - if quantization_level is not None or mflux_version is not None: - return QwenWeightHandler._load_saved_model_weights(weights, None, quantization_level, mflux_version) - - # Otherwise, it's HuggingFace weights that need mapping - mapping = mapping_getter() - mapped_weights = WeightMapper.apply_mapping(weights, mapping) - return mapped_weights, None, None - - @staticmethod - def _load_safetensors_shards(path: Path, loading_mode: str = "multi_glob") -> dict[str, mx.array]: - all_weights = {} - - if loading_mode == "single": - # VAE style: Single file loading - safetensors_files = [f for f in path.glob("*.safetensors") if not f.name.startswith("._")] - if not safetensors_files: - raise FileNotFoundError(f"No safetensors files found in {path}") - - weights_file = safetensors_files[0] - data = mx.load(str(weights_file), return_metadata=True) - all_weights = dict(data[0].items()) - - elif loading_mode == "multi_json": - # Text encoder style: Use JSON index to map params to files - index_path = path / "model.safetensors.index.json" - with open(index_path) as f: - index = json.load(f) - - # Group weights by file - files_to_load = {} - for param_name, file_name in index["weight_map"].items(): - if file_name not in files_to_load: - files_to_load[file_name] = [] - files_to_load[file_name].append(param_name) - - # Load weights from each file - for file_name, param_names in files_to_load.items(): - file_path = path / file_name - - # Load the safetensor file with fallback to torch conversion - try: - file_weights = mlx_load_file(str(file_path)) - except Exception: # noqa: BLE001 - # If MLX can't load directly, try with torch and convert - torch_weights = torch_load_file(str(file_path)) - file_weights = {} - for name, tensor in torch_weights.items(): - # Convert to float32 if bfloat16, then to MLX - if tensor.dtype == torch.bfloat16: - tensor = tensor.to(torch.float32) - file_weights[name] = mx.array(tensor.numpy()) - - # Add requested parameters to combined weights - for param_name in param_names: - if param_name in file_weights: - all_weights[param_name] = file_weights[param_name] - else: # "multi_glob" - # Transformer style: Directly glob all safetensors files - shard_files = sorted([f for f in path.glob("*.safetensors") if not f.name.startswith("._")]) - if not shard_files: - raise FileNotFoundError(f"No safetensors found in {path}") - - for shard in shard_files: - data, metadata = mx.load(str(shard), return_metadata=True) - all_weights.update(dict(data.items())) - - return all_weights - - @staticmethod - def _load_saved_model_weights( - weights: dict[str, mx.array] | None, - path: Path | None, - quantization_level: int | None, - mflux_version: str | None, - ) -> tuple[dict, int | None, str | None]: - # If weights already loaded, use them; otherwise load from path - if weights is None: - # For saved models, always use multi_glob (no index.json needed) - weights = QwenWeightHandler._load_safetensors_shards(path, loading_mode="multi_glob") - - return tree_unflatten(list(weights.items())), quantization_level, mflux_version - - @staticmethod - def _detect_metadata(path: Path) -> tuple[int | None, str | None]: - file_glob = sorted([f for f in path.glob("*.safetensors") if not f.name.startswith("._")]) - if file_glob: - data = mx.load(str(file_glob[0]), return_metadata=True) - if len(data) > 1: - quantization_level = data[1].get("quantization_level") - mflux_version = data[1].get("mflux_version") - return quantization_level, mflux_version - return None, None diff --git a/src/mflux/models/qwen/weights/qwen_weight_mapping.py b/src/mflux/models/qwen/weights/qwen_weight_mapping.py index c3b477f..6fabb47 100644 --- a/src/mflux/models/qwen/weights/qwen_weight_mapping.py +++ b/src/mflux/models/qwen/weights/qwen_weight_mapping.py @@ -1,915 +1,854 @@ from typing import List -import mlx.core as mx - from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget - - -def reshape_gamma_to_1d(tensor: mx.array) -> mx.array: - if len(tensor.shape) > 1: - return mx.reshape(tensor, (tensor.shape[0],)) - return tensor - - -def transpose_patch_embed(tensor: mx.array) -> mx.array: - if len(tensor.shape) == 5: - return tensor.transpose(0, 2, 3, 4, 1) - return tensor - - -def transpose_conv3d_weight(tensor: mx.array) -> mx.array: - if len(tensor.shape) == 5: - return tensor.transpose(0, 2, 3, 4, 1) - return tensor - - -def transpose_conv2d_weight(tensor: mx.array) -> mx.array: - if len(tensor.shape) == 4: - return tensor.transpose(0, 2, 3, 1) - return tensor +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms class QwenWeightMapping(WeightMapping): @staticmethod def get_transformer_mapping() -> List[WeightTarget]: return [ - # Top-level mappings WeightTarget( - mlx_path="img_in.weight", - hf_patterns=["img_in.weight"], + to_pattern="img_in.weight", + from_pattern=["img_in.weight"], ), WeightTarget( - mlx_path="img_in.bias", - hf_patterns=["img_in.bias"], + to_pattern="img_in.bias", + from_pattern=["img_in.bias"], ), WeightTarget( - mlx_path="txt_norm.weight", - hf_patterns=["txt_norm.weight"], + to_pattern="txt_norm.weight", + from_pattern=["txt_norm.weight"], ), WeightTarget( - mlx_path="txt_in.weight", - hf_patterns=["txt_in.weight"], + to_pattern="txt_in.weight", + from_pattern=["txt_in.weight"], ), WeightTarget( - mlx_path="txt_in.bias", - hf_patterns=["txt_in.bias"], - ), - # Time text embedder - WeightTarget( - mlx_path="time_text_embed.timestep_embedder.linear_1.weight", - hf_patterns=["time_text_embed.timestep_embedder.linear_1.weight"], + to_pattern="txt_in.bias", + from_pattern=["txt_in.bias"], ), WeightTarget( - mlx_path="time_text_embed.timestep_embedder.linear_1.bias", - hf_patterns=["time_text_embed.timestep_embedder.linear_1.bias"], + to_pattern="time_text_embed.timestep_embedder.linear_1.weight", + from_pattern=["time_text_embed.timestep_embedder.linear_1.weight"], ), WeightTarget( - mlx_path="time_text_embed.timestep_embedder.linear_2.weight", - hf_patterns=["time_text_embed.timestep_embedder.linear_2.weight"], + to_pattern="time_text_embed.timestep_embedder.linear_1.bias", + from_pattern=["time_text_embed.timestep_embedder.linear_1.bias"], ), WeightTarget( - mlx_path="time_text_embed.timestep_embedder.linear_2.bias", - hf_patterns=["time_text_embed.timestep_embedder.linear_2.bias"], - ), - # Output head - WeightTarget( - mlx_path="norm_out.linear.weight", - hf_patterns=["norm_out.linear.weight"], + to_pattern="time_text_embed.timestep_embedder.linear_2.weight", + from_pattern=["time_text_embed.timestep_embedder.linear_2.weight"], ), WeightTarget( - mlx_path="norm_out.linear.bias", - hf_patterns=["norm_out.linear.bias"], + to_pattern="time_text_embed.timestep_embedder.linear_2.bias", + from_pattern=["time_text_embed.timestep_embedder.linear_2.bias"], ), WeightTarget( - mlx_path="proj_out.weight", - hf_patterns=["proj_out.weight"], + to_pattern="norm_out.linear.weight", + from_pattern=["norm_out.linear.weight"], ), WeightTarget( - mlx_path="proj_out.bias", - hf_patterns=["proj_out.bias"], - ), - # Transformer blocks - Attention - WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_q.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_q.weight"], + to_pattern="norm_out.linear.bias", + from_pattern=["norm_out.linear.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_q.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_q.bias"], + to_pattern="proj_out.weight", + from_pattern=["proj_out.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_k.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_k.weight"], + to_pattern="proj_out.bias", + from_pattern=["proj_out.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_k.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_k.bias"], + to_pattern="transformer_blocks.{block}.attn.to_q.weight", + from_pattern=["transformer_blocks.{block}.attn.to_q.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_v.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_v.weight"], + to_pattern="transformer_blocks.{block}.attn.to_q.bias", + from_pattern=["transformer_blocks.{block}.attn.to_q.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_v.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_v.bias"], + to_pattern="transformer_blocks.{block}.attn.to_k.weight", + from_pattern=["transformer_blocks.{block}.attn.to_k.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_q_proj.weight", - hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.weight"], + to_pattern="transformer_blocks.{block}.attn.to_k.bias", + from_pattern=["transformer_blocks.{block}.attn.to_k.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_q_proj.bias", - hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.bias"], + to_pattern="transformer_blocks.{block}.attn.to_v.weight", + from_pattern=["transformer_blocks.{block}.attn.to_v.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_k_proj.weight", - hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.weight"], + to_pattern="transformer_blocks.{block}.attn.to_v.bias", + from_pattern=["transformer_blocks.{block}.attn.to_v.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_k_proj.bias", - hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.bias"], + to_pattern="transformer_blocks.{block}.attn.add_q_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_q_proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_v_proj.weight", - hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.weight"], + to_pattern="transformer_blocks.{block}.attn.add_q_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_q_proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.add_v_proj.bias", - hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.bias"], + to_pattern="transformer_blocks.{block}.attn.add_k_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_k_proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_q.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_q.weight"], + to_pattern="transformer_blocks.{block}.attn.add_k_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_k_proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_k.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_k.weight"], + to_pattern="transformer_blocks.{block}.attn.add_v_proj.weight", + from_pattern=["transformer_blocks.{block}.attn.add_v_proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_added_q.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_added_q.weight"], + to_pattern="transformer_blocks.{block}.attn.add_v_proj.bias", + from_pattern=["transformer_blocks.{block}.attn.add_v_proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.norm_added_k.weight", - hf_patterns=["transformer_blocks.{block}.attn.norm_added_k.weight"], + to_pattern="transformer_blocks.{block}.attn.norm_q.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_q.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.attn_to_out.0.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_out.0.weight"], + to_pattern="transformer_blocks.{block}.attn.norm_k.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_k.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.attn_to_out.0.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_out.0.bias"], + to_pattern="transformer_blocks.{block}.attn.norm_added_q.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_added_q.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_add_out.weight", - hf_patterns=["transformer_blocks.{block}.attn.to_add_out.weight"], + to_pattern="transformer_blocks.{block}.attn.norm_added_k.weight", + from_pattern=["transformer_blocks.{block}.attn.norm_added_k.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.attn.to_add_out.bias", - hf_patterns=["transformer_blocks.{block}.attn.to_add_out.bias"], - ), - # Transformer blocks - Modulation - WeightTarget( - mlx_path="transformer_blocks.{block}.img_mod_linear.weight", - hf_patterns=["transformer_blocks.{block}.img_mod.1.weight"], + to_pattern="transformer_blocks.{block}.attn.attn_to_out.0.weight", + from_pattern=["transformer_blocks.{block}.attn.to_out.0.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.img_mod_linear.bias", - hf_patterns=["transformer_blocks.{block}.img_mod.1.bias"], + to_pattern="transformer_blocks.{block}.attn.attn_to_out.0.bias", + from_pattern=["transformer_blocks.{block}.attn.to_out.0.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.txt_mod_linear.weight", - hf_patterns=["transformer_blocks.{block}.txt_mod.1.weight"], + to_pattern="transformer_blocks.{block}.attn.to_add_out.weight", + from_pattern=["transformer_blocks.{block}.attn.to_add_out.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.txt_mod_linear.bias", - hf_patterns=["transformer_blocks.{block}.txt_mod.1.bias"], - ), - # Transformer blocks - Feed Forward - WeightTarget( - mlx_path="transformer_blocks.{block}.img_ff.mlp_in.weight", - hf_patterns=["transformer_blocks.{block}.img_mlp.net.0.proj.weight"], + to_pattern="transformer_blocks.{block}.attn.to_add_out.bias", + from_pattern=["transformer_blocks.{block}.attn.to_add_out.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.img_ff.mlp_in.bias", - hf_patterns=["transformer_blocks.{block}.img_mlp.net.0.proj.bias"], + to_pattern="transformer_blocks.{block}.img_mod_linear.weight", + from_pattern=["transformer_blocks.{block}.img_mod.1.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.img_ff.mlp_out.weight", - hf_patterns=["transformer_blocks.{block}.img_mlp.net.2.weight"], + to_pattern="transformer_blocks.{block}.img_mod_linear.bias", + from_pattern=["transformer_blocks.{block}.img_mod.1.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.img_ff.mlp_out.bias", - hf_patterns=["transformer_blocks.{block}.img_mlp.net.2.bias"], + to_pattern="transformer_blocks.{block}.txt_mod_linear.weight", + from_pattern=["transformer_blocks.{block}.txt_mod.1.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.txt_ff.mlp_in.weight", - hf_patterns=["transformer_blocks.{block}.txt_mlp.net.0.proj.weight"], + to_pattern="transformer_blocks.{block}.txt_mod_linear.bias", + from_pattern=["transformer_blocks.{block}.txt_mod.1.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.txt_ff.mlp_in.bias", - hf_patterns=["transformer_blocks.{block}.txt_mlp.net.0.proj.bias"], + to_pattern="transformer_blocks.{block}.img_ff.mlp_in.weight", + from_pattern=["transformer_blocks.{block}.img_mlp.net.0.proj.weight"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.txt_ff.mlp_out.weight", - hf_patterns=["transformer_blocks.{block}.txt_mlp.net.2.weight"], + to_pattern="transformer_blocks.{block}.img_ff.mlp_in.bias", + from_pattern=["transformer_blocks.{block}.img_mlp.net.0.proj.bias"], ), WeightTarget( - mlx_path="transformer_blocks.{block}.txt_ff.mlp_out.bias", - hf_patterns=["transformer_blocks.{block}.txt_mlp.net.2.bias"], + to_pattern="transformer_blocks.{block}.img_ff.mlp_out.weight", + from_pattern=["transformer_blocks.{block}.img_mlp.net.2.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.img_ff.mlp_out.bias", + from_pattern=["transformer_blocks.{block}.img_mlp.net.2.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.txt_ff.mlp_in.weight", + from_pattern=["transformer_blocks.{block}.txt_mlp.net.0.proj.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.txt_ff.mlp_in.bias", + from_pattern=["transformer_blocks.{block}.txt_mlp.net.0.proj.bias"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.txt_ff.mlp_out.weight", + from_pattern=["transformer_blocks.{block}.txt_mlp.net.2.weight"], + ), + WeightTarget( + to_pattern="transformer_blocks.{block}.txt_ff.mlp_out.bias", + from_pattern=["transformer_blocks.{block}.txt_mlp.net.2.bias"], ), ] @staticmethod def get_vae_mapping() -> List[WeightTarget]: return [ - # ========== Decoder ========== - # Decoder conv_in WeightTarget( - mlx_path="decoder.conv_in.conv3d.weight", - hf_patterns=["decoder.conv_in.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.conv_in.conv3d.weight", + from_pattern=["decoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.conv_in.conv3d.bias", - hf_patterns=["decoder.conv_in.bias"], + to_pattern="decoder.conv_in.conv3d.bias", + from_pattern=["decoder.conv_in.bias"], ), - # Decoder conv_out WeightTarget( - mlx_path="decoder.conv_out.conv3d.weight", - hf_patterns=["decoder.conv_out.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.conv_out.conv3d.weight", + from_pattern=["decoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.conv_out.conv3d.bias", - hf_patterns=["decoder.conv_out.bias"], + to_pattern="decoder.conv_out.conv3d.bias", + from_pattern=["decoder.conv_out.bias"], ), - # Decoder norm_out (gamma -> weight with reshape) WeightTarget( - mlx_path="decoder.norm_out.weight", - hf_patterns=["decoder.norm_out.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.norm_out.weight", + from_pattern=["decoder.norm_out.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), - # Post quant conv WeightTarget( - mlx_path="post_quant_conv.conv3d.weight", - hf_patterns=["post_quant_conv.weight"], - transform=transpose_conv3d_weight, + to_pattern="post_quant_conv.conv3d.weight", + from_pattern=["post_quant_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="post_quant_conv.conv3d.bias", - hf_patterns=["post_quant_conv.bias"], + to_pattern="post_quant_conv.conv3d.bias", + from_pattern=["post_quant_conv.bias"], ), - # Decoder mid_block resnets WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.mid_block.resnets.{i}.conv1.conv3d.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.bias", - hf_patterns=["decoder.mid_block.resnets.{i}.conv1.bias"], + to_pattern="decoder.mid_block.resnets.{i}.conv1.conv3d.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.mid_block.resnets.{i}.conv2.conv3d.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.bias", - hf_patterns=["decoder.mid_block.resnets.{i}.conv2.bias"], + to_pattern="decoder.mid_block.resnets.{i}.conv2.conv3d.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.norm1.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.mid_block.resnets.{i}.norm2.weight", - hf_patterns=["decoder.mid_block.resnets.{i}.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), - # Decoder mid_block attention WeightTarget( - mlx_path="decoder.mid_block.attentions.0.norm.weight", - hf_patterns=["decoder.mid_block.attentions.0.norm.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.mid_block.attentions.0.norm.weight", + from_pattern=["decoder.mid_block.attentions.0.norm.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.mid_block.attentions.0.to_qkv.weight", - hf_patterns=["decoder.mid_block.attentions.0.to_qkv.weight"], - transform=transpose_conv2d_weight, + to_pattern="decoder.mid_block.attentions.0.to_qkv.weight", + from_pattern=["decoder.mid_block.attentions.0.to_qkv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.attentions.0.to_qkv.bias", - hf_patterns=["decoder.mid_block.attentions.0.to_qkv.bias"], + to_pattern="decoder.mid_block.attentions.0.to_qkv.bias", + from_pattern=["decoder.mid_block.attentions.0.to_qkv.bias"], ), WeightTarget( - mlx_path="decoder.mid_block.attentions.0.proj.weight", - hf_patterns=["decoder.mid_block.attentions.0.proj.weight"], - transform=transpose_conv2d_weight, + to_pattern="decoder.mid_block.attentions.0.proj.weight", + from_pattern=["decoder.mid_block.attentions.0.proj.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="decoder.mid_block.attentions.0.proj.bias", - hf_patterns=["decoder.mid_block.attentions.0.proj.bias"], + to_pattern="decoder.mid_block.attentions.0.proj.bias", + from_pattern=["decoder.mid_block.attentions.0.proj.bias"], ), - # Decoder up_blocks resnets WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.conv1.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_block{block}.resnets.{res}.conv1.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.conv1.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"], + to_pattern="decoder.up_block{block}.resnets.{res}.conv1.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"], ), WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.conv2.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_block{block}.resnets.{res}.conv2.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.conv2.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"], + to_pattern="decoder.up_block{block}.resnets.{res}.conv2.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"], ), WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.norm1.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.up_block{block}.resnets.{res}.norm1.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.norm2.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="decoder.up_block{block}.resnets.{res}.norm2.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), - # Decoder up_blocks optional conv_shortcut -> skip_conv WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"], - transform=transpose_conv3d_weight, - required=False, # Optional weight - ), - WeightTarget( - mlx_path="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"], - required=False, # Optional weight - ), - # Decoder up_blocks upsamplers (blocks 0, 1, 2) - WeightTarget( - mlx_path="decoder.up_block{block}.upsamplers.0.resample_conv.weight", - hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.resample.1.weight"], - transform=transpose_conv2d_weight, - ), - WeightTarget( - mlx_path="decoder.up_block{block}.upsamplers.0.resample_conv.bias", - hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.resample.1.bias"], - ), - # Decoder up_blocks time_conv (blocks 0, 1) - WeightTarget( - mlx_path="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.weight", - hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.time_conv.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.bias", - hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.time_conv.bias"], - ), - # ========== Encoder ========== - # Encoder conv_in - WeightTarget( - mlx_path="encoder.conv_in.conv3d.weight", - hf_patterns=["encoder.conv_in.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.conv_in.conv3d.bias", - hf_patterns=["encoder.conv_in.bias"], - ), - # Encoder conv_out - WeightTarget( - mlx_path="encoder.conv_out.conv3d.weight", - hf_patterns=["encoder.conv_out.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.conv_out.conv3d.bias", - hf_patterns=["encoder.conv_out.bias"], - ), - # Encoder norm_out - WeightTarget( - mlx_path="encoder.norm_out.weight", - hf_patterns=["encoder.norm_out.gamma"], - transform=reshape_gamma_to_1d, - ), - # Encoder mid_block attention - WeightTarget( - mlx_path="encoder.mid_block.attentions.0.norm.weight", - hf_patterns=["encoder.mid_block.attentions.0.norm.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.mid_block.attentions.0.to_qkv.weight", - hf_patterns=["encoder.mid_block.attentions.0.to_qkv.weight"], - transform=transpose_conv2d_weight, - ), - WeightTarget( - mlx_path="encoder.mid_block.attentions.0.to_qkv.bias", - hf_patterns=["encoder.mid_block.attentions.0.to_qkv.bias"], - ), - WeightTarget( - mlx_path="encoder.mid_block.attentions.0.proj.weight", - hf_patterns=["encoder.mid_block.attentions.0.proj.weight"], - transform=transpose_conv2d_weight, - ), - WeightTarget( - mlx_path="encoder.mid_block.attentions.0.proj.bias", - hf_patterns=["encoder.mid_block.attentions.0.proj.bias"], - ), - # Encoder mid_block resnets - WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.conv1.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.bias", - hf_patterns=["encoder.mid_block.resnets.{i}.conv1.bias"], - ), - WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.conv2.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.bias", - hf_patterns=["encoder.mid_block.resnets.{i}.conv2.bias"], - ), - WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.norm1.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.norm1.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.mid_block.resnets.{i}.norm2.weight", - hf_patterns=["encoder.mid_block.resnets.{i}.norm2.gamma"], - transform=reshape_gamma_to_1d, - ), - # Encoder down_blocks - Stage 0 (flat_idx 0,1 -> resnets[0,1]; flat_idx 2 -> downsampler) - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.0.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.0.conv1.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.0.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.0.conv1.bias"], - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.0.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.0.conv2.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.0.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.0.conv2.bias"], - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.0.norm1.weight", - hf_patterns=["encoder.down_blocks.0.norm1.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.0.norm2.weight", - hf_patterns=["encoder.down_blocks.0.norm2.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.1.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.1.conv1.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.1.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.1.conv1.bias"], - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.1.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.1.conv2.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.1.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.1.conv2.bias"], - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.1.norm1.weight", - hf_patterns=["encoder.down_blocks.1.norm1.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.resnets.1.norm2.weight", - hf_patterns=["encoder.down_blocks.1.norm2.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.downsamplers.0.resample_conv.weight", - hf_patterns=["encoder.down_blocks.2.resample.1.weight"], - transform=transpose_conv2d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.0.downsamplers.0.resample_conv.bias", - hf_patterns=["encoder.down_blocks.2.resample.1.bias"], - ), - # Encoder down_blocks - Stage 1 (flat_idx 3,4 -> resnets[0,1]; flat_idx 5 -> downsampler) - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.3.conv1.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.3.conv1.bias"], - ), - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.3.conv2.weight"], - transform=transpose_conv3d_weight, - ), - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.3.conv2.bias"], - ), - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.norm1.weight", - hf_patterns=["encoder.down_blocks.3.norm1.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.norm2.weight", - hf_patterns=["encoder.down_blocks.3.norm2.gamma"], - transform=reshape_gamma_to_1d, - ), - WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.weight", - hf_patterns=["encoder.down_blocks.3.conv_shortcut.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv3d_weight, required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.bias", - hf_patterns=["encoder.down_blocks.3.conv_shortcut.bias"], + to_pattern="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"], required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.1.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.4.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_block{block}.upsamplers.0.resample_conv.weight", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.resample.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.1.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.4.conv1.bias"], + to_pattern="decoder.up_block{block}.upsamplers.0.resample_conv.bias", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.resample.1.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.1.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.4.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.weight", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.time_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.1.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.4.conv2.bias"], + to_pattern="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.bias", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.time_conv.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.1.norm1.weight", - hf_patterns=["encoder.down_blocks.4.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.conv_in.conv3d.weight", + from_pattern=["encoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.1.resnets.1.norm2.weight", - hf_patterns=["encoder.down_blocks.4.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.conv_in.conv3d.bias", + from_pattern=["encoder.conv_in.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.1.downsamplers.0.resample_conv.weight", - hf_patterns=["encoder.down_blocks.5.resample.1.weight"], - transform=transpose_conv2d_weight, + to_pattern="encoder.conv_out.conv3d.weight", + from_pattern=["encoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.1.downsamplers.0.resample_conv.bias", - hf_patterns=["encoder.down_blocks.5.resample.1.bias"], - ), - # Encoder down_blocks - Stage 2 (flat_idx 6,7 -> resnets[0,1]; flat_idx 8 -> downsampler) - WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.6.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.conv_out.conv3d.bias", + from_pattern=["encoder.conv_out.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.6.conv1.bias"], + to_pattern="encoder.norm_out.weight", + from_pattern=["encoder.norm_out.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.6.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.mid_block.attentions.0.norm.weight", + from_pattern=["encoder.mid_block.attentions.0.norm.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.6.conv2.bias"], + to_pattern="encoder.mid_block.attentions.0.to_qkv.weight", + from_pattern=["encoder.mid_block.attentions.0.to_qkv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.norm1.weight", - hf_patterns=["encoder.down_blocks.6.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.mid_block.attentions.0.to_qkv.bias", + from_pattern=["encoder.mid_block.attentions.0.to_qkv.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.norm2.weight", - hf_patterns=["encoder.down_blocks.6.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.mid_block.attentions.0.proj.weight", + from_pattern=["encoder.mid_block.attentions.0.proj.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.weight", - hf_patterns=["encoder.down_blocks.6.conv_shortcut.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.mid_block.attentions.0.proj.bias", + from_pattern=["encoder.mid_block.attentions.0.proj.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv1.conv3d.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv1.conv3d.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv2.conv3d.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv2.conv3d.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.0.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.0.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.0.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.0.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.0.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.0.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.0.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.0.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.0.norm1.weight", + from_pattern=["encoder.down_blocks.0.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.0.norm2.weight", + from_pattern=["encoder.down_blocks.0.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.1.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.1.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.1.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.1.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.1.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.1.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.1.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.1.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.1.norm1.weight", + from_pattern=["encoder.down_blocks.1.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.resnets.1.norm2.weight", + from_pattern=["encoder.down_blocks.1.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.downsamplers.0.resample_conv.weight", + from_pattern=["encoder.down_blocks.2.resample.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.0.downsamplers.0.resample_conv.bias", + from_pattern=["encoder.down_blocks.2.resample.1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.3.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.3.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.3.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.3.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.norm1.weight", + from_pattern=["encoder.down_blocks.3.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.norm2.weight", + from_pattern=["encoder.down_blocks.3.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.weight", + from_pattern=["encoder.down_blocks.3.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv3d_weight, required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.bias", - hf_patterns=["encoder.down_blocks.6.conv_shortcut.bias"], + to_pattern="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.bias", + from_pattern=["encoder.down_blocks.3.conv_shortcut.bias"], required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.1.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.7.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.1.resnets.1.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.4.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.1.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.7.conv1.bias"], + to_pattern="encoder.down_blocks.1.resnets.1.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.4.conv1.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.1.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.7.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.1.resnets.1.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.4.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.1.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.7.conv2.bias"], + to_pattern="encoder.down_blocks.1.resnets.1.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.4.conv2.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.1.norm1.weight", - hf_patterns=["encoder.down_blocks.7.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.down_blocks.1.resnets.1.norm1.weight", + from_pattern=["encoder.down_blocks.4.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.2.resnets.1.norm2.weight", - hf_patterns=["encoder.down_blocks.7.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.down_blocks.1.resnets.1.norm2.weight", + from_pattern=["encoder.down_blocks.4.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.2.downsamplers.0.resample_conv.weight", - hf_patterns=["encoder.down_blocks.8.resample.1.weight"], - transform=transpose_conv2d_weight, + to_pattern="encoder.down_blocks.1.downsamplers.0.resample_conv.weight", + from_pattern=["encoder.down_blocks.5.resample.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.2.downsamplers.0.resample_conv.bias", - hf_patterns=["encoder.down_blocks.8.resample.1.bias"], + to_pattern="encoder.down_blocks.1.downsamplers.0.resample_conv.bias", + from_pattern=["encoder.down_blocks.5.resample.1.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.weight", - hf_patterns=["encoder.down_blocks.8.time_conv.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.2.resnets.0.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.6.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.bias", - hf_patterns=["encoder.down_blocks.8.time_conv.bias"], - ), - # Encoder down_blocks - Stage 3 (flat_idx 9,10 -> resnets[0,1]) - WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.0.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.9.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.2.resnets.0.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.6.conv1.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.0.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.9.conv1.bias"], + to_pattern="encoder.down_blocks.2.resnets.0.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.6.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.0.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.9.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.2.resnets.0.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.6.conv2.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.0.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.9.conv2.bias"], + to_pattern="encoder.down_blocks.2.resnets.0.norm1.weight", + from_pattern=["encoder.down_blocks.6.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.0.norm1.weight", - hf_patterns=["encoder.down_blocks.9.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.down_blocks.2.resnets.0.norm2.weight", + from_pattern=["encoder.down_blocks.6.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.0.norm2.weight", - hf_patterns=["encoder.down_blocks.9.norm2.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.weight", + from_pattern=["encoder.down_blocks.6.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.1.conv1.conv3d.weight", - hf_patterns=["encoder.down_blocks.10.conv1.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.bias", + from_pattern=["encoder.down_blocks.6.conv_shortcut.bias"], + required=False, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.1.conv1.conv3d.bias", - hf_patterns=["encoder.down_blocks.10.conv1.bias"], + to_pattern="encoder.down_blocks.2.resnets.1.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.7.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.1.conv2.conv3d.weight", - hf_patterns=["encoder.down_blocks.10.conv2.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.2.resnets.1.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.7.conv1.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.1.conv2.conv3d.bias", - hf_patterns=["encoder.down_blocks.10.conv2.bias"], + to_pattern="encoder.down_blocks.2.resnets.1.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.7.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.1.norm1.weight", - hf_patterns=["encoder.down_blocks.10.norm1.gamma"], - transform=reshape_gamma_to_1d, + to_pattern="encoder.down_blocks.2.resnets.1.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.7.conv2.bias"], ), WeightTarget( - mlx_path="encoder.down_blocks.3.resnets.1.norm2.weight", - hf_patterns=["encoder.down_blocks.10.norm2.gamma"], - transform=reshape_gamma_to_1d, - ), - # Quant conv - WeightTarget( - mlx_path="quant_conv.conv3d.weight", - hf_patterns=["quant_conv.weight"], - transform=transpose_conv3d_weight, + to_pattern="encoder.down_blocks.2.resnets.1.norm1.weight", + from_pattern=["encoder.down_blocks.7.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, ), WeightTarget( - mlx_path="quant_conv.conv3d.bias", - hf_patterns=["quant_conv.bias"], + to_pattern="encoder.down_blocks.2.resnets.1.norm2.weight", + from_pattern=["encoder.down_blocks.7.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.2.downsamplers.0.resample_conv.weight", + from_pattern=["encoder.down_blocks.8.resample.1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.2.downsamplers.0.resample_conv.bias", + from_pattern=["encoder.down_blocks.8.resample.1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.weight", + from_pattern=["encoder.down_blocks.8.time_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.bias", + from_pattern=["encoder.down_blocks.8.time_conv.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.0.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.9.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.0.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.9.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.0.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.9.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.0.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.9.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.0.norm1.weight", + from_pattern=["encoder.down_blocks.9.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.0.norm2.weight", + from_pattern=["encoder.down_blocks.9.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.1.conv1.conv3d.weight", + from_pattern=["encoder.down_blocks.10.conv1.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.1.conv1.conv3d.bias", + from_pattern=["encoder.down_blocks.10.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.1.conv2.conv3d.weight", + from_pattern=["encoder.down_blocks.10.conv2.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.1.conv2.conv3d.bias", + from_pattern=["encoder.down_blocks.10.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.1.norm1.weight", + from_pattern=["encoder.down_blocks.10.norm1.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="encoder.down_blocks.3.resnets.1.norm2.weight", + from_pattern=["encoder.down_blocks.10.norm2.gamma"], + transform=WeightTransforms.reshape_gamma_to_1d, + ), + WeightTarget( + to_pattern="quant_conv.conv3d.weight", + from_pattern=["quant_conv.weight"], + transform=WeightTransforms.transpose_conv3d_weight, + ), + WeightTarget( + to_pattern="quant_conv.conv3d.bias", + from_pattern=["quant_conv.bias"], ), ] @staticmethod def get_text_encoder_mapping() -> List[WeightTarget]: return [ - # Top-level embeddings WeightTarget( - mlx_path="encoder.embed_tokens.weight", - hf_patterns=["model.embed_tokens.weight"], - ), - # Final norm - WeightTarget( - mlx_path="encoder.norm.weight", - hf_patterns=["model.norm.weight"], - ), - # Encoder layers (28 layers) - WeightTarget( - mlx_path="encoder.layers.{layer}.input_layernorm.weight", - hf_patterns=["model.layers.{layer}.input_layernorm.weight"], + to_pattern="encoder.embed_tokens.weight", + from_pattern=["model.embed_tokens.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.post_attention_layernorm.weight", - hf_patterns=["model.layers.{layer}.post_attention_layernorm.weight"], - ), - # Self attention - WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.q_proj.weight", - hf_patterns=["model.layers.{layer}.self_attn.q_proj.weight"], + to_pattern="encoder.norm.weight", + from_pattern=["model.norm.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.q_proj.bias", - hf_patterns=["model.layers.{layer}.self_attn.q_proj.bias"], + to_pattern="encoder.layers.{layer}.input_layernorm.weight", + from_pattern=["model.layers.{layer}.input_layernorm.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.k_proj.weight", - hf_patterns=["model.layers.{layer}.self_attn.k_proj.weight"], + to_pattern="encoder.layers.{layer}.post_attention_layernorm.weight", + from_pattern=["model.layers.{layer}.post_attention_layernorm.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.k_proj.bias", - hf_patterns=["model.layers.{layer}.self_attn.k_proj.bias"], + to_pattern="encoder.layers.{layer}.self_attn.q_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.q_proj.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.v_proj.weight", - hf_patterns=["model.layers.{layer}.self_attn.v_proj.weight"], + to_pattern="encoder.layers.{layer}.self_attn.q_proj.bias", + from_pattern=["model.layers.{layer}.self_attn.q_proj.bias"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.v_proj.bias", - hf_patterns=["model.layers.{layer}.self_attn.v_proj.bias"], + to_pattern="encoder.layers.{layer}.self_attn.k_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.k_proj.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.self_attn.o_proj.weight", - hf_patterns=["model.layers.{layer}.self_attn.o_proj.weight"], + to_pattern="encoder.layers.{layer}.self_attn.k_proj.bias", + from_pattern=["model.layers.{layer}.self_attn.k_proj.bias"], + ), + WeightTarget( + to_pattern="encoder.layers.{layer}.self_attn.v_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.v_proj.weight"], + ), + WeightTarget( + to_pattern="encoder.layers.{layer}.self_attn.v_proj.bias", + from_pattern=["model.layers.{layer}.self_attn.v_proj.bias"], + ), + WeightTarget( + to_pattern="encoder.layers.{layer}.self_attn.o_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.o_proj.weight"], ), # MLP WeightTarget( - mlx_path="encoder.layers.{layer}.mlp.gate_proj.weight", - hf_patterns=["model.layers.{layer}.mlp.gate_proj.weight"], + to_pattern="encoder.layers.{layer}.mlp.gate_proj.weight", + from_pattern=["model.layers.{layer}.mlp.gate_proj.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.mlp.up_proj.weight", - hf_patterns=["model.layers.{layer}.mlp.up_proj.weight"], + to_pattern="encoder.layers.{layer}.mlp.up_proj.weight", + from_pattern=["model.layers.{layer}.mlp.up_proj.weight"], ), WeightTarget( - mlx_path="encoder.layers.{layer}.mlp.down_proj.weight", - hf_patterns=["model.layers.{layer}.mlp.down_proj.weight"], + to_pattern="encoder.layers.{layer}.mlp.down_proj.weight", + from_pattern=["model.layers.{layer}.mlp.down_proj.weight"], ), - # Visual weights (optional, only present in Edit models) - # Patch embedding (with transpose transform) WeightTarget( - mlx_path="encoder.visual.patch_embed.proj.weight", - hf_patterns=["visual.patch_embed.proj.weight"], - transform=transpose_patch_embed, - required=False, - ), - # Vision transformer blocks (32 blocks) - use {block} placeholder - WeightTarget( - mlx_path="encoder.visual.blocks.{block}.attn.qkv.weight", - hf_patterns=["visual.blocks.{block}.attn.qkv.weight"], + to_pattern="encoder.visual.patch_embed.proj.weight", + from_pattern=["visual.patch_embed.proj.weight"], + transform=WeightTransforms.transpose_patch_embed, required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.attn.qkv.bias", - hf_patterns=["visual.blocks.{block}.attn.qkv.bias"], + to_pattern="encoder.visual.blocks.{block}.attn.qkv.weight", + from_pattern=["visual.blocks.{block}.attn.qkv.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.attn.proj.weight", - hf_patterns=["visual.blocks.{block}.attn.proj.weight"], + to_pattern="encoder.visual.blocks.{block}.attn.qkv.bias", + from_pattern=["visual.blocks.{block}.attn.qkv.bias"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.attn.proj.bias", - hf_patterns=["visual.blocks.{block}.attn.proj.bias"], + to_pattern="encoder.visual.blocks.{block}.attn.proj.weight", + from_pattern=["visual.blocks.{block}.attn.proj.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.mlp.gate_proj.weight", - hf_patterns=["visual.blocks.{block}.mlp.gate_proj.weight"], + to_pattern="encoder.visual.blocks.{block}.attn.proj.bias", + from_pattern=["visual.blocks.{block}.attn.proj.bias"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.mlp.gate_proj.bias", - hf_patterns=["visual.blocks.{block}.mlp.gate_proj.bias"], + to_pattern="encoder.visual.blocks.{block}.mlp.gate_proj.weight", + from_pattern=["visual.blocks.{block}.mlp.gate_proj.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.mlp.up_proj.weight", - hf_patterns=["visual.blocks.{block}.mlp.up_proj.weight"], + to_pattern="encoder.visual.blocks.{block}.mlp.gate_proj.bias", + from_pattern=["visual.blocks.{block}.mlp.gate_proj.bias"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.mlp.up_proj.bias", - hf_patterns=["visual.blocks.{block}.mlp.up_proj.bias"], + to_pattern="encoder.visual.blocks.{block}.mlp.up_proj.weight", + from_pattern=["visual.blocks.{block}.mlp.up_proj.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.mlp.down_proj.weight", - hf_patterns=["visual.blocks.{block}.mlp.down_proj.weight"], + to_pattern="encoder.visual.blocks.{block}.mlp.up_proj.bias", + from_pattern=["visual.blocks.{block}.mlp.up_proj.bias"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.mlp.down_proj.bias", - hf_patterns=["visual.blocks.{block}.mlp.down_proj.bias"], + to_pattern="encoder.visual.blocks.{block}.mlp.down_proj.weight", + from_pattern=["visual.blocks.{block}.mlp.down_proj.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.norm1.weight", - hf_patterns=["visual.blocks.{block}.norm1.weight"], + to_pattern="encoder.visual.blocks.{block}.mlp.down_proj.bias", + from_pattern=["visual.blocks.{block}.mlp.down_proj.bias"], required=False, ), WeightTarget( - mlx_path="encoder.visual.blocks.{block}.norm2.weight", - hf_patterns=["visual.blocks.{block}.norm2.weight"], - required=False, - ), - # Patch merger - WeightTarget( - mlx_path="encoder.visual.merger.ln_q.weight", - hf_patterns=["visual.merger.ln_q.weight"], + to_pattern="encoder.visual.blocks.{block}.norm1.weight", + from_pattern=["visual.blocks.{block}.norm1.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.merger.mlp_0.weight", - hf_patterns=["visual.merger.mlp.0.weight"], + to_pattern="encoder.visual.blocks.{block}.norm2.weight", + from_pattern=["visual.blocks.{block}.norm2.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.merger.mlp_0.bias", - hf_patterns=["visual.merger.mlp.0.bias"], + to_pattern="encoder.visual.merger.ln_q.weight", + from_pattern=["visual.merger.ln_q.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.merger.mlp_1.weight", - hf_patterns=["visual.merger.mlp.2.weight"], + to_pattern="encoder.visual.merger.mlp_0.weight", + from_pattern=["visual.merger.mlp.0.weight"], required=False, ), WeightTarget( - mlx_path="encoder.visual.merger.mlp_1.bias", - hf_patterns=["visual.merger.mlp.2.bias"], + to_pattern="encoder.visual.merger.mlp_0.bias", + from_pattern=["visual.merger.mlp.0.bias"], + required=False, + ), + WeightTarget( + to_pattern="encoder.visual.merger.mlp_1.weight", + from_pattern=["visual.merger.mlp.2.weight"], + required=False, + ), + WeightTarget( + to_pattern="encoder.visual.merger.mlp_1.bias", + from_pattern=["visual.merger.mlp.2.bias"], required=False, ), ] diff --git a/src/mflux/models/qwen/weights/qwen_weight_util.py b/src/mflux/models/qwen/weights/qwen_weight_util.py deleted file mode 100644 index dde5011..0000000 --- a/src/mflux/models/qwen/weights/qwen_weight_util.py +++ /dev/null @@ -1,107 +0,0 @@ -from typing import TYPE_CHECKING - -import mlx.core as mx -import mlx.nn as nn -from mlx.utils import tree_flatten # noqa: F401 - -from mflux.config.config import Config -from mflux.models.common.quantization.quantization_util import QuantizationUtil -from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_transformer import VisionTransformer - -if TYPE_CHECKING: - from mflux.models.qwen.weights.qwen_weight_handler import QwenWeightHandler - - -class QwenWeightUtil: - @staticmethod - def flatten(params): - return [(k, v) for p in params for (k, v) in p] - - @staticmethod - def reshape_weights(key, value): - if len(value.shape) == 4: - value = value.transpose(0, 2, 3, 1) - elif len(value.shape) == 5: - value = value.transpose(0, 2, 3, 4, 1) - value = value.reshape(-1).reshape(value.shape).astype(Config.precision) - return [(key, value)] - - @staticmethod - def set_weights_and_quantize( - quantize_arg: int | None, - weights: "QwenWeightHandler", - vae: nn.Module, - transformer: nn.Module, - text_encoder: nn.Module | None = None, - ) -> int | None: - if weights.meta_data.quantization_level is None and quantize_arg is None: - QwenWeightUtil._set_model_weights(weights, vae, transformer, text_encoder) - return None - - if weights.meta_data.quantization_level is None and quantize_arg is not None: - bits = quantize_arg - QwenWeightUtil._set_model_weights(weights, vae, transformer, text_encoder) - QuantizationUtil.quantize_qwen_models(text_encoder, vae, transformer, bits, weights) - return bits - - if weights.meta_data.quantization_level is not None: - bits = weights.meta_data.quantization_level - QuantizationUtil.quantize_qwen_models(text_encoder, vae, transformer, bits, weights) - QwenWeightUtil._set_model_weights(weights, vae, transformer, text_encoder) - return bits - - raise Exception("Error setting weights") - - @staticmethod - def _convert_weights_to_bf16(weights_dict: dict): - """ - Recursively convert all weight tensors to BF16. - This matches PyTorch's behavior where the text encoder model is loaded in BF16. - """ - - def convert_recursive(obj): - if isinstance(obj, mx.array): - return obj.astype(mx.bfloat16) - elif isinstance(obj, dict): - return {k: convert_recursive(v) for k, v in obj.items()} - elif isinstance(obj, list): - return [convert_recursive(item) for item in obj] - else: - return obj - - return convert_recursive(weights_dict) - - @staticmethod - def _set_model_weights( - weights: "QwenWeightHandler", - vae: nn.Module, - transformer: nn.Module, - text_encoder: nn.Module | None = None, - ): - vae.update(weights.vae, strict=False) - transformer.update(weights.transformer, strict=False) - - # Check if visual weights are present and create visual transformer if needed - if text_encoder is not None and weights.qwen_text_encoder is not None: - has_visual_weights = ( - "encoder" in weights.qwen_text_encoder and "visual" in weights.qwen_text_encoder["encoder"] - ) - if has_visual_weights and text_encoder.encoder.visual is None: - text_encoder.encoder.visual = VisionTransformer( - patch_size=14, - temporal_patch_size=2, - in_channels=3, - embed_dim=1280, - depth=32, - num_heads=16, - mlp_ratio=2.671875, - hidden_size=text_encoder.encoder.hidden_size, - spatial_merge_size=2, # Match HF's spatial merging - ) - - # Convert all text encoder weights to BF16 to match PyTorch's behavior - # PyTorch loads the model in BF16 and performs all computations in BF16 - # This ensures numerical consistency between MLX and PyTorch - weights.qwen_text_encoder = QwenWeightUtil._convert_weights_to_bf16(weights.qwen_text_encoder) - - text_encoder.update(weights.qwen_text_encoder, strict=False) diff --git a/src/mflux/models/z_image/__init__.py b/src/mflux/models/z_image/__init__.py new file mode 100644 index 0000000..c7de8b8 --- /dev/null +++ b/src/mflux/models/z_image/__init__.py @@ -0,0 +1,4 @@ +from mflux.models.z_image.variants.turbo import ZImageTurbo +from mflux.models.z_image.z_image_initializer import ZImageInitializer + +__all__ = ["ZImageTurbo", "ZImageInitializer"] diff --git a/src/mflux/models/z_image/cli/__init__.py b/src/mflux/models/z_image/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mflux/models/z_image/cli/z_image_turbo_generate.py b/src/mflux/models/z_image/cli/z_image_turbo_generate.py new file mode 100644 index 0000000..4b47695 --- /dev/null +++ b/src/mflux/models/z_image/cli/z_image_turbo_generate.py @@ -0,0 +1,65 @@ +from mflux.callbacks.callback_manager import CallbackManager +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.z_image.latent_creator import ZImageLatentCreator +from mflux.models.z_image.variants.turbo.z_image_turbo import ZImageTurbo +from mflux.utils.dimension_resolver import DimensionResolver +from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException +from mflux.utils.prompt_util import PromptUtil + + +def main(): + # 0. Parse command line arguments + parser = CommandLineParser(description="Generate an image using Z-Image Turbo based on a prompt.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True, supports_dimension_scale_factor=True) + parser.add_image_to_image_arguments(required=False) + parser.add_output_arguments() + args = parser.parse_args() + + # 1. Load the model + model = ZImageTurbo( + quantize=args.quantize, + model_path=args.model_path, + lora_paths=args.lora_paths, + lora_scales=args.lora_scales, + ) + + # 2. Register callbacks + memory_saver = CallbackManager.register_callbacks( + args=args, + model=model, + latent_creator=ZImageLatentCreator, + ) + + try: + # Resolve dimensions (supports ScaleFactor like "2x" when --image-path is provided) + width, height = DimensionResolver.resolve( + width=args.width, + height=args.height, + reference_image_path=args.image_path, + ) + + for seed in args.seed: + # 3. Generate an image for each seed value + image = model.generate_image( + seed=seed, + prompt=PromptUtil.read_prompt(args), + width=width, + height=height, + image_path=args.image_path, + num_inference_steps=args.steps, + image_strength=args.image_strength, + ) + # 4. Save the image + image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) + except (StopImageGenerationException, PromptFileReadError) as exc: + print(exc) + finally: + if memory_saver: + print(memory_saver.memory_stats()) + + +if __name__ == "__main__": + main() diff --git a/src/mflux/models/z_image/latent_creator/__init__.py b/src/mflux/models/z_image/latent_creator/__init__.py new file mode 100644 index 0000000..4465972 --- /dev/null +++ b/src/mflux/models/z_image/latent_creator/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.latent_creator.z_image_latent_creator import ZImageLatentCreator + +__all__ = ["ZImageLatentCreator"] diff --git a/src/mflux/models/z_image/latent_creator/z_image_latent_creator.py b/src/mflux/models/z_image/latent_creator/z_image_latent_creator.py new file mode 100644 index 0000000..31d628c --- /dev/null +++ b/src/mflux/models/z_image/latent_creator/z_image_latent_creator.py @@ -0,0 +1,29 @@ +import mlx.core as mx + +from mflux.models.common.config import ModelConfig + + +class ZImageLatentCreator: + @staticmethod + def create_noise(seed: int, height: int, width: int) -> mx.array: + return mx.random.normal( + shape=[ + 16, + 1, + height // 8, + width // 8, + ], + key=mx.random.key(seed), + ).astype(ModelConfig.precision) + + @staticmethod + def pack_latents(latents: mx.array, height: int, width: int) -> mx.array: + latents = mx.expand_dims(latents, axis=2) + latents = mx.squeeze(latents, axis=0) + return latents + + @staticmethod + def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: # noqa: ARG004 + latents = mx.expand_dims(latents, axis=0) + latents = mx.squeeze(latents, axis=2) + return latents diff --git a/src/mflux/models/z_image/model/__init__.py b/src/mflux/models/z_image/model/__init__.py new file mode 100644 index 0000000..cce6886 --- /dev/null +++ b/src/mflux/models/z_image/model/__init__.py @@ -0,0 +1,2 @@ +# Z-Image model components + diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/__init__.py b/src/mflux/models/z_image/model/z_image_text_encoder/__init__.py new file mode 100644 index 0000000..5e5b811 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/__init__.py @@ -0,0 +1,2 @@ +# Z-Image Text Encoder (Qwen3) + diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/attention.py b/src/mflux/models/z_image/model/z_image_text_encoder/attention.py new file mode 100644 index 0000000..7c980d3 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/attention.py @@ -0,0 +1,76 @@ +import mlx.core as mx +from mlx import nn +from mlx.core.fast import scaled_dot_product_attention + + +class Attention(nn.Module): + def __init__( + self, + hidden_size: int = 2560, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + head_dim: int = 128, + ): + super().__init__() + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + self.num_kv_groups = num_attention_heads // num_key_value_heads + self.scale = head_dim**-0.5 + + self.q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.o_proj = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False) + self.q_norm = nn.RMSNorm(head_dim) + self.k_norm = nn.RMSNorm(head_dim) + + def __call__( + self, + hidden_states: mx.array, + attention_mask: mx.array | None = None, + position_embeddings: tuple[mx.array, mx.array] | None = None, + ) -> mx.array: + batch_size, seq_len, _ = hidden_states.shape + + q = self.q_proj(hidden_states).reshape(batch_size, seq_len, self.num_heads, self.head_dim) + k = self.k_proj(hidden_states).reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = self.v_proj(hidden_states).reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim) + + q = self.q_norm(q) + k = self.k_norm(k) + + if position_embeddings is not None: + q, k = Attention._apply_rotary_pos_emb(q, k, *position_embeddings) + + if self.num_kv_groups > 1: + k = mx.repeat(k, self.num_kv_groups, axis=2) + v = mx.repeat(v, self.num_kv_groups, axis=2) + + q = mx.transpose(q, axes=(0, 2, 1, 3)) + k = mx.transpose(k, axes=(0, 2, 1, 3)) + v = mx.transpose(v, axes=(0, 2, 1, 3)) + + attn_output = scaled_dot_product_attention(q, k, v, scale=self.scale, mask=attention_mask) + attn_output = mx.transpose(attn_output, axes=(0, 2, 1, 3)).reshape(batch_size, seq_len, -1) + return self.o_proj(attn_output) + + @staticmethod + def _apply_rotary_pos_emb( + q: mx.array, + k: mx.array, + cos: mx.array, + sin: mx.array, + unsqueeze_dim: int = 2, + ) -> tuple[mx.array, mx.array]: + cos = mx.expand_dims(cos, axis=unsqueeze_dim) + sin = mx.expand_dims(sin, axis=unsqueeze_dim) + q_embed = (q * cos) + (Attention._rotate_half(q) * sin) + k_embed = (k * cos) + (Attention._rotate_half(k) * sin) + return q_embed, k_embed + + @staticmethod + def _rotate_half(x: mx.array) -> mx.array: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return mx.concatenate([-x2, x1], axis=-1) diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/encoder_layer.py b/src/mflux/models/z_image/model/z_image_text_encoder/encoder_layer.py new file mode 100644 index 0000000..8804957 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/encoder_layer.py @@ -0,0 +1,35 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_text_encoder.attention import Attention +from mflux.models.z_image.model.z_image_text_encoder.mlp import MLP + + +class EncoderLayer(nn.Module): + def __init__( + self, + hidden_size: int = 2560, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + intermediate_size: int = 9728, + head_dim: int = 128, + rms_norm_eps: float = 1e-6, + ): + super().__init__() + self.input_layernorm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) + self.self_attn = Attention(hidden_size, num_attention_heads, num_key_value_heads, head_dim) + self.mlp = MLP(hidden_size, intermediate_size) + + def __call__( + self, + hidden_states: mx.array, + attention_mask: mx.array | None = None, + position_embeddings: tuple[mx.array, mx.array] | None = None, + ) -> mx.array: + residual = hidden_states + hidden_states = self.self_attn(self.input_layernorm(hidden_states), attention_mask, position_embeddings) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp(self.post_attention_layernorm(hidden_states)) + return residual + hidden_states diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/mlp.py b/src/mflux/models/z_image/model/z_image_text_encoder/mlp.py new file mode 100644 index 0000000..ee1b0e5 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/mlp.py @@ -0,0 +1,13 @@ +import mlx.core as mx +from mlx import nn + + +class MLP(nn.Module): + def __init__(self, hidden_size: int = 2560, intermediate_size: int = 9728): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/prompt_encoder.py b/src/mflux/models/z_image/model/z_image_text_encoder/prompt_encoder.py new file mode 100644 index 0000000..c190683 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/prompt_encoder.py @@ -0,0 +1,17 @@ +import mlx.core as mx + +from mflux.models.common.tokenizer import Tokenizer +from mflux.models.z_image.model.z_image_text_encoder.text_encoder import TextEncoder + + +class PromptEncoder: + @staticmethod + def encode_prompt( + prompt: str, + tokenizer: Tokenizer, + text_encoder: TextEncoder, + ) -> mx.array: + output = tokenizer.tokenize(prompt) + cap_feats = text_encoder(output.input_ids, output.attention_mask) + num_valid = int(mx.sum(output.attention_mask[0]).item()) + return cap_feats[0, :num_valid, :] diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/rope.py b/src/mflux/models/z_image/model/z_image_text_encoder/rope.py new file mode 100644 index 0000000..59d2ed2 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/rope.py @@ -0,0 +1,16 @@ +import mlx.core as mx +from mlx import nn + + +class RotaryEmbedding(nn.Module): + def __init__(self, dim: int, base: float = 1000000.0): + super().__init__() + self.inv_freq = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)) + + def __call__(self, x: mx.array, position_ids: mx.array) -> tuple[mx.array, mx.array]: + seq_len = position_ids.shape[-1] + freqs = mx.outer(mx.arange(seq_len, dtype=mx.float32), self.inv_freq) + emb = mx.concatenate([freqs, freqs], axis=-1) + cos = mx.cos(emb)[None, :, :] + sin = mx.sin(emb)[None, :, :] + return cos.astype(x.dtype), sin.astype(x.dtype) diff --git a/src/mflux/models/z_image/model/z_image_text_encoder/text_encoder.py b/src/mflux/models/z_image/model/z_image_text_encoder/text_encoder.py new file mode 100644 index 0000000..eeb535a --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_text_encoder/text_encoder.py @@ -0,0 +1,81 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.common.config import ModelConfig +from mflux.models.z_image.model.z_image_text_encoder.encoder_layer import EncoderLayer +from mflux.models.z_image.model.z_image_text_encoder.rope import RotaryEmbedding + + +class TextEncoder(nn.Module): + def __init__( + self, + vocab_size: int = 151936, + hidden_size: int = 2560, + num_hidden_layers: int = 36, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + intermediate_size: int = 9728, + head_dim: int = 128, + max_position_embeddings: int = 40960, + rope_theta: float = 1000000.0, + rms_norm_eps: float = 1e-6, + ): + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, hidden_size) + self.layers = [ + EncoderLayer( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + intermediate_size=intermediate_size, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + ) + for _ in range(num_hidden_layers) + ] + self.norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps) + self.rotary_emb = RotaryEmbedding(dim=head_dim, base=rope_theta) + + def __call__(self, input_ids: mx.array, attention_mask: mx.array | None = None) -> mx.array: + batch_size, seq_len = input_ids.shape + hidden_states = self.embed_tokens(input_ids).astype(mx.float32) + position_ids = mx.broadcast_to(mx.arange(seq_len, dtype=mx.int32)[None, :], (batch_size, seq_len)) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Causal mask + causal_mask = TextEncoder._get_causal_mask(attention_mask, batch_size, hidden_states, seq_len) + + # Forward through layers + all_hidden_states = [hidden_states] + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + attention_mask=causal_mask, + position_embeddings=position_embeddings, + ) + all_hidden_states.append(hidden_states) + + return all_hidden_states[-2].astype(ModelConfig.precision) + + @staticmethod + def _get_causal_mask(attention_mask, batch_size, hidden_states, seq_len): + causal_mask = TextEncoder._create_causal_mask(seq_len, hidden_states.dtype) + if attention_mask is not None: + padding_mask = mx.where( + attention_mask[:, None, None, :] == 1, + mx.zeros((batch_size, 1, 1, seq_len), dtype=hidden_states.dtype), + mx.full((batch_size, 1, 1, seq_len), float("-inf"), dtype=hidden_states.dtype), + ) + causal_mask = causal_mask + padding_mask + return causal_mask + + @staticmethod + def _create_causal_mask(seq_len: int, dtype: mx.Dtype) -> mx.array: + idx = mx.arange(seq_len, dtype=mx.int32) + mask = idx[:, None] >= idx[None, :] + causal_mask = mx.where( + mask, + mx.zeros((seq_len, seq_len), dtype=dtype), + mx.full((seq_len, seq_len), float("-inf"), dtype=dtype), + ) + return causal_mask[None, None, :, :] diff --git a/src/mflux/models/z_image/model/z_image_transformer/__init__.py b/src/mflux/models/z_image/model/z_image_transformer/__init__.py new file mode 100644 index 0000000..07770a3 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.model.z_image_transformer.transformer import ZImageTransformer + +__all__ = ["ZImageTransformer"] diff --git a/src/mflux/models/z_image/model/z_image_transformer/attention.py b/src/mflux/models/z_image/model/z_image_transformer/attention.py new file mode 100644 index 0000000..205151b --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/attention.py @@ -0,0 +1,90 @@ +import mlx.core as mx +from mlx import nn +from mlx.core.fast import scaled_dot_product_attention + + +class ZImageAttention(nn.Module): + def __init__( + self, + dim: int, + n_heads: int, + qk_norm: bool = True, + eps: float = 1e-5, + ): + super().__init__() + self.dim = dim + self.n_heads = n_heads + self.head_dim = dim // n_heads + self.scale = self.head_dim**-0.5 + + # Projections (no bias) + self.to_q = nn.Linear(dim, dim, bias=False) + self.to_k = nn.Linear(dim, dim, bias=False) + self.to_v = nn.Linear(dim, dim, bias=False) + self.to_out = [nn.Linear(dim, dim, bias=False)] + + # Optional QK normalization + if qk_norm: + self.norm_q = nn.RMSNorm(self.head_dim, eps=eps) + self.norm_k = nn.RMSNorm(self.head_dim, eps=eps) + else: + self.norm_q = None + self.norm_k = None + + def __call__( + self, + hidden_states: mx.array, + attention_mask: mx.array | None = None, + freqs_cis: mx.array | None = None, + ) -> mx.array: + batch_size, seq_len, _ = hidden_states.shape + + # Project to Q, K, V + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + # Reshape to (batch, seq_len, heads, head_dim) + query = query.reshape(batch_size, seq_len, self.n_heads, self.head_dim) + key = key.reshape(batch_size, seq_len, self.n_heads, self.head_dim) + value = value.reshape(batch_size, seq_len, self.n_heads, self.head_dim) + + # Apply QK normalization + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + # Apply RoPE + if freqs_cis is not None: + query = ZImageAttention._apply_rotary_emb(query, freqs_cis) + key = ZImageAttention._apply_rotary_emb(key, freqs_cis) + + # Transpose for attention: (batch, heads, seq_len, head_dim) + query = mx.transpose(query, axes=(0, 2, 1, 3)) + key = mx.transpose(key, axes=(0, 2, 1, 3)) + value = mx.transpose(value, axes=(0, 2, 1, 3)) + + # Convert boolean mask to additive mask for SDPA + mask = None + if attention_mask is not None: + mask = mx.where(attention_mask[:, None, None, :], mx.array(0.0), mx.array(float("-inf"))) + + hidden_states = scaled_dot_product_attention(query, key, value, scale=self.scale, mask=mask) + hidden_states = mx.transpose(hidden_states, axes=(0, 2, 1, 3)) + hidden_states = hidden_states.reshape(batch_size, seq_len, self.dim) + hidden_states = self.to_out[0](hidden_states) + return hidden_states + + @staticmethod + def _apply_rotary_emb(x: mx.array, freqs_cis: mx.array) -> mx.array: + batch_size, seq_len, n_heads, head_dim = x.shape + x = x.reshape(batch_size, seq_len, n_heads, head_dim // 2, 2) + freqs_cis = mx.expand_dims(freqs_cis, axis=0) + freqs_cis = mx.expand_dims(freqs_cis, axis=2) + x_real, x_imag = x[..., 0], x[..., 1] + freqs_cos, freqs_sin = freqs_cis[..., 0], freqs_cis[..., 1] + x_out_real = x_real * freqs_cos - x_imag * freqs_sin + x_out_imag = x_real * freqs_sin + x_imag * freqs_cos + x_out = mx.stack([x_out_real, x_out_imag], axis=-1) + return x_out.reshape(batch_size, seq_len, n_heads, head_dim) diff --git a/src/mflux/models/z_image/model/z_image_transformer/context_block.py b/src/mflux/models/z_image/model/z_image_transformer/context_block.py new file mode 100644 index 0000000..7eca319 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/context_block.py @@ -0,0 +1,23 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_transformer.attention import ZImageAttention +from mflux.models.z_image.model.z_image_transformer.feed_forward import FeedForward + + +class ZImageContextBlock(nn.Module): + def __init__(self, dim: int, n_heads: int, norm_eps: float = 1e-5, qk_norm: bool = True): + super().__init__() + self.attention = ZImageAttention(dim=dim, n_heads=n_heads, qk_norm=qk_norm, eps=1e-5) + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + self.attention_norm1 = nn.RMSNorm(dim, eps=norm_eps) + self.attention_norm2 = nn.RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = nn.RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = nn.RMSNorm(dim, eps=norm_eps) + + def __call__(self, x: mx.array, attn_mask: mx.array, freqs_cis: mx.array) -> mx.array: + attn_out = self.attention(self.attention_norm1(x), attention_mask=attn_mask, freqs_cis=freqs_cis) + x = x + self.attention_norm2(attn_out) + ffn_out = self.feed_forward(self.ffn_norm1(x)) + x = x + self.ffn_norm2(ffn_out) + return x diff --git a/src/mflux/models/z_image/model/z_image_transformer/feed_forward.py b/src/mflux/models/z_image/model/z_image_transformer/feed_forward.py new file mode 100644 index 0000000..1e5aaac --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/feed_forward.py @@ -0,0 +1,13 @@ +import mlx.core as mx +from mlx import nn + + +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim, bias=False) + self.w2 = nn.Linear(hidden_dim, dim, bias=False) + self.w3 = nn.Linear(dim, hidden_dim, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.w2(nn.silu(self.w1(x)) * self.w3(x)) diff --git a/src/mflux/models/z_image/model/z_image_transformer/final_layer.py b/src/mflux/models/z_image/model/z_image_transformer/final_layer.py new file mode 100644 index 0000000..5c555b5 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/final_layer.py @@ -0,0 +1,16 @@ +import mlx.core as mx +from mlx import nn + + +class FinalLayer(nn.Module): + def __init__(self, hidden_size: int, out_channels: int): + super().__init__() + self.norm = nn.LayerNorm(hidden_size, eps=1e-6, affine=False) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + self.adaLN_modulation = [nn.Linear(min(hidden_size, 256), hidden_size, bias=True)] + + def __call__(self, x: mx.array, c: mx.array) -> mx.array: + scale = 1.0 + self.adaLN_modulation[0](nn.silu(c)) + scale = mx.expand_dims(scale, axis=1) + x = self.norm(x) * scale + return self.linear(x) diff --git a/src/mflux/models/z_image/model/z_image_transformer/rope_embedder.py b/src/mflux/models/z_image/model/z_image_transformer/rope_embedder.py new file mode 100644 index 0000000..4c6d9f8 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/rope_embedder.py @@ -0,0 +1,37 @@ +import mlx.core as mx + + +class RopeEmbedder: + def __init__( + self, + theta: float = 256.0, + axes_dims: list[int] | None = None, + axes_lens: list[int] | None = None, + ): + if axes_dims is None: + axes_dims = [16, 56, 56] + if axes_lens is None: + axes_lens = [64, 128, 128] + + self.axes_dims = axes_dims + self.freqs_cis = RopeEmbedder._precompute_freqs_cis(axes_dims, axes_lens, theta) + + def __call__(self, ids: mx.array) -> mx.array: + result = [] + for i in range(len(self.axes_dims)): + index = ids[:, i].astype(mx.int32) + result.append(self.freqs_cis[i][index]) + return mx.concatenate(result, axis=1) + + @staticmethod + def _precompute_freqs_cis(axes_dims, axes_lens, theta) -> list[mx.array]: + freqs_cis = [] + for d, e in zip(axes_dims, axes_lens): + freqs = 1.0 / (theta ** (mx.arange(0, d, 2, dtype=mx.float32) / d)) + timestep = mx.arange(e, dtype=mx.float32) + freqs = mx.outer(timestep, freqs) + cos_freqs = mx.cos(freqs) + sin_freqs = mx.sin(freqs) + freqs_cis_i = mx.stack([cos_freqs, sin_freqs], axis=-1) + freqs_cis.append(freqs_cis_i) + return freqs_cis diff --git a/src/mflux/models/z_image/model/z_image_transformer/timestep_embedder.py b/src/mflux/models/z_image/model/z_image_transformer/timestep_embedder.py new file mode 100644 index 0000000..265829f --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/timestep_embedder.py @@ -0,0 +1,37 @@ +import math + +import mlx.core as mx +from mlx import nn + + +class TimestepEmbedder(nn.Module): + def __init__( + self, + out_size: int, + mid_size: int | None = None, + frequency_embedding_size: int = 256, + ): + super().__init__() + if mid_size is None: + mid_size = out_size + + self.frequency_embedding_size = frequency_embedding_size + self.linear1 = nn.Linear(frequency_embedding_size, mid_size, bias=True) + self.linear2 = nn.Linear(mid_size, out_size, bias=True) + + def __call__(self, t: mx.array) -> mx.array: + t_freq = self._timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.linear1(t_freq) + t_emb = nn.silu(t_emb) + t_emb = self.linear2(t_emb) + return t_emb + + @staticmethod + def _timestep_embedding(t: mx.array, dim: int, max_period: float = 10000.0) -> mx.array: + half = dim // 2 + freqs = mx.exp(-math.log(max_period) * mx.arange(0, half, dtype=mx.float32) / half) + args = t[:, None].astype(mx.float32) * freqs[None] + embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1) + if dim % 2: + embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1) + return embedding diff --git a/src/mflux/models/z_image/model/z_image_transformer/transformer.py b/src/mflux/models/z_image/model/z_image_transformer/transformer.py new file mode 100644 index 0000000..8358c1d --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/transformer.py @@ -0,0 +1,188 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_transformer.context_block import ZImageContextBlock +from mflux.models.z_image.model.z_image_transformer.final_layer import FinalLayer +from mflux.models.z_image.model.z_image_transformer.rope_embedder import RopeEmbedder +from mflux.models.z_image.model.z_image_transformer.timestep_embedder import TimestepEmbedder +from mflux.models.z_image.model.z_image_transformer.transformer_block import ZImageTransformerBlock + + +class ZImageTransformer(nn.Module): + def __init__( + self, + patch_size: int = 2, + f_patch_size: int = 1, + in_channels: int = 16, + dim: int = 3840, + n_layers: int = 30, + n_refiner_layers: int = 2, + n_heads: int = 30, + norm_eps: float = 1e-5, + qk_norm: bool = True, + cap_feat_dim: int = 2560, + rope_theta: float = 256.0, + t_scale: float = 1000.0, + axes_dims: list[int] | None = None, + axes_lens: list[int] | None = None, + ): + super().__init__() + + axes_dims = axes_dims or [32, 48, 48] + axes_lens = axes_lens or [1024, 512, 512] + + self.in_channels = in_channels + self.out_channels = in_channels + self.patch_size = patch_size + self.f_patch_size = f_patch_size + self.dim = dim + self.n_heads = n_heads + self.t_scale = t_scale + + key = f"{patch_size}-{f_patch_size}" + embed_dim = f_patch_size * patch_size * patch_size * in_channels + self.all_x_embedder = {key: nn.Linear(embed_dim, dim, bias=True)} + self.all_final_layer = {key: FinalLayer(dim, embed_dim)} + + self.t_embedder = TimestepEmbedder(out_size=min(dim, 256), mid_size=1024) + self.cap_embedder = [nn.RMSNorm(cap_feat_dim, eps=norm_eps), nn.Linear(cap_feat_dim, dim, bias=True)] + self.x_pad_token = mx.zeros((1, dim)) + self.cap_pad_token = mx.zeros((1, dim)) + + self.noise_refiner = [ZImageTransformerBlock(dim, n_heads, norm_eps, qk_norm) for _ in range(n_refiner_layers)] # fmt: off + self.context_refiner = [ZImageContextBlock(dim, n_heads, norm_eps, qk_norm) for _ in range(n_refiner_layers)] # fmt: off + self.layers = [ZImageTransformerBlock(dim, n_heads, norm_eps, qk_norm) for _ in range(n_layers)] # fmt: off + self.rope_embedder = RopeEmbedder(theta=rope_theta, axes_dims=axes_dims, axes_lens=axes_lens) + + def __call__(self, x: mx.array, t: int, sigmas: mx.array, cap_feats: mx.array) -> mx.array: + key = f"{self.patch_size}-{self.f_patch_size}" + + # Time embedding + t_value = mx.array([1.0 - sigmas[t].item()]) + t_emb = self.t_embedder(t_value * self.t_scale) + + # Patchify image and caption + x_emb, cap_emb, x_size, x_pos_ids, cap_pos_ids, x_pad_mask, cap_pad_mask = ZImageTransformer._patchify( + image=x, + cap_feats=cap_feats, + patch_size=self.patch_size, + f_patch_size=self.f_patch_size, + ) + + # Image embedding + x_emb = self.all_x_embedder[key](x_emb) + x_emb = mx.where(x_pad_mask[:, None], self.x_pad_token, x_emb) + x_freqs_cis = self.rope_embedder(x_pos_ids) + x_attn_mask = mx.ones((1, x_emb.shape[0]), dtype=mx.bool_) + x_emb = mx.expand_dims(x_emb, axis=0) + + # Noise refiner + for layer in self.noise_refiner: + x_emb = layer( + x=x_emb, + attn_mask=x_attn_mask, + freqs_cis=x_freqs_cis, + t_emb=t_emb, + ) + + # Caption embedding + cap_emb = self.cap_embedder[1](self.cap_embedder[0](cap_emb)) + cap_emb = mx.where(cap_pad_mask[:, None], self.cap_pad_token, cap_emb) + cap_freqs_cis = self.rope_embedder(cap_pos_ids) + cap_attn_mask = mx.ones((1, cap_emb.shape[0]), dtype=mx.bool_) + cap_emb = mx.expand_dims(cap_emb, axis=0) + + # Context refiner + for layer in self.context_refiner: + cap_emb = layer( + x=cap_emb, + attn_mask=cap_attn_mask, + freqs_cis=cap_freqs_cis, + ) + + # Unify and main layers + x_len = x_emb.shape[1] + unified = mx.concatenate([x_emb, cap_emb], axis=1) + unified_freqs_cis = mx.concatenate([x_freqs_cis, cap_freqs_cis], axis=0) + unified_attn_mask = mx.ones((1, unified.shape[1]), dtype=mx.bool_) + + for layer in self.layers: + unified = layer( + x=unified, + attn_mask=unified_attn_mask, + freqs_cis=unified_freqs_cis, + t_emb=t_emb, + ) + + # Final layer and unpatchify + unified = self.all_final_layer[key](unified, t_emb) + output = ZImageTransformer._unpatchify( + x=unified[0, :x_len], + size=x_size, + patch_size=self.patch_size, + f_patch_size=self.f_patch_size, + out_channels=self.out_channels, + ) + return -output + + @staticmethod + def _patchify( + image: mx.array, + cap_feats: mx.array, + patch_size: int, + f_patch_size: int, + ) -> tuple[mx.array, mx.array, tuple[int, int, int], mx.array, mx.array, mx.array, mx.array]: + pH = pW = patch_size + pF = f_patch_size + + # Caption padding + cap_ori_len = cap_feats.shape[0] + cap_padding_len = (-cap_ori_len) % 32 + cap_pos_ids = ZImageTransformer._create_coord_grid((cap_ori_len + cap_padding_len, 1, 1), (1, 0, 0)).reshape(-1, 3) # fmt: off + cap_pad_mask = mx.concatenate([mx.zeros((cap_ori_len,), dtype=mx.bool_), mx.ones((cap_padding_len,), dtype=mx.bool_)]) # fmt: off + cap_padded = mx.concatenate([cap_feats, mx.repeat(cap_feats[-1:], cap_padding_len, axis=0)], axis=0) if cap_padding_len > 0 else cap_feats # fmt: off + + # Image patchification + C, F, H, W = image.shape + image_size = (F, H, W) + F_tokens, H_tokens, W_tokens = F // pF, H // pH, W // pW + + image = image.reshape(C, F_tokens, pF, H_tokens, pH, W_tokens, pW) + image = mx.transpose(image, axes=(1, 3, 5, 2, 4, 6, 0)) + image = image.reshape(F_tokens * H_tokens * W_tokens, pF * pH * pW * C) + + # Image padding + image_ori_len = image.shape[0] + image_padding_len = (-image_ori_len) % 32 + image_pos_ids = ZImageTransformer._create_coord_grid((F_tokens, H_tokens, W_tokens), (cap_ori_len + cap_padding_len + 1, 0, 0)).reshape(-1, 3) # fmt: off + + if image_padding_len > 0: + image_pos_ids = mx.concatenate([image_pos_ids, mx.zeros((image_padding_len, 3), dtype=mx.int32)], axis=0) + image = mx.concatenate([image, mx.repeat(image[-1:], image_padding_len, axis=0)], axis=0) + + image_pad_mask = mx.concatenate([mx.zeros((image_ori_len,), dtype=mx.bool_), mx.ones((image_padding_len,), dtype=mx.bool_)]) # fmt: off + + return image, cap_padded, image_size, image_pos_ids, cap_pos_ids, image_pad_mask, cap_pad_mask + + @staticmethod + def _unpatchify( + x: mx.array, + size: tuple[int, int, int], + patch_size: int, + f_patch_size: int, + out_channels: int, + ) -> mx.array: + pH = pW = patch_size + pF = f_patch_size + F, H, W = size + ori_len = (F // pF) * (H // pH) * (W // pW) + x = x[:ori_len].reshape(F // pF, H // pH, W // pW, pF, pH, pW, out_channels) + x = mx.transpose(x, axes=(6, 0, 3, 1, 4, 2, 5)) + return x.reshape(out_channels, F, H, W) + + @staticmethod + def _create_coord_grid(size: tuple[int, ...], start: tuple[int, ...] | None = None) -> mx.array: + start = start or tuple(0 for _ in size) + axes = [mx.arange(x0, x0 + span, dtype=mx.int32) for x0, span in zip(start, size)] + grids = mx.meshgrid(*axes, indexing="ij") + return mx.stack(grids, axis=-1) diff --git a/src/mflux/models/z_image/model/z_image_transformer/transformer_block.py b/src/mflux/models/z_image/model/z_image_transformer/transformer_block.py new file mode 100644 index 0000000..2845402 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_transformer/transformer_block.py @@ -0,0 +1,35 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_transformer.attention import ZImageAttention +from mflux.models.z_image.model.z_image_transformer.feed_forward import FeedForward + + +class ZImageTransformerBlock(nn.Module): + def __init__(self, dim: int, n_heads: int, norm_eps: float = 1e-5, qk_norm: bool = True): + super().__init__() + self.attention = ZImageAttention(dim=dim, n_heads=n_heads, qk_norm=qk_norm, eps=1e-5) + self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8)) + self.attention_norm1 = nn.RMSNorm(dim, eps=norm_eps) + self.attention_norm2 = nn.RMSNorm(dim, eps=norm_eps) + self.ffn_norm1 = nn.RMSNorm(dim, eps=norm_eps) + self.ffn_norm2 = nn.RMSNorm(dim, eps=norm_eps) + self.adaLN_modulation = [nn.Linear(min(dim, 256), 4 * dim, bias=True)] + + def __call__(self, x: mx.array, attn_mask: mx.array, freqs_cis: mx.array, t_emb: mx.array) -> mx.array: + # Compute modulation parameters + modulation = mx.expand_dims(self.adaLN_modulation[0](t_emb), axis=1) + scale_msa, gate_msa, scale_mlp, gate_mlp = mx.split(modulation, 4, axis=2) + scale_msa = 1.0 + scale_msa + scale_mlp = 1.0 + scale_mlp + gate_msa = mx.tanh(gate_msa) + gate_mlp = mx.tanh(gate_mlp) + + # Attention with modulation + attn_out = self.attention(self.attention_norm1(x) * scale_msa, attention_mask=attn_mask, freqs_cis=freqs_cis) + x = x + gate_msa * self.attention_norm2(attn_out) + + # FFN with modulation + ffn_out = self.feed_forward(self.ffn_norm1(x) * scale_mlp) + x = x + gate_mlp * self.ffn_norm2(ffn_out) + return x diff --git a/src/mflux/models/z_image/model/z_image_vae/__init__.py b/src/mflux/models/z_image/model/z_image_vae/__init__.py new file mode 100644 index 0000000..6dab1ea --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.model.z_image_vae.vae import VAE + +__all__ = ["VAE"] diff --git a/src/mflux/models/z_image/model/z_image_vae/common/__init__.py b/src/mflux/models/z_image/model/z_image_vae/common/__init__.py new file mode 100644 index 0000000..244c6d4 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/common/__init__.py @@ -0,0 +1,5 @@ +from mflux.models.z_image.model.z_image_vae.common.attention import Attention +from mflux.models.z_image.model.z_image_vae.common.resnet_block_2d import ResnetBlock2D +from mflux.models.z_image.model.z_image_vae.common.unet_mid_block import UNetMidBlock + +__all__ = ["Attention", "ResnetBlock2D", "UNetMidBlock"] diff --git a/src/mflux/models/z_image/model/z_image_vae/common/attention.py b/src/mflux/models/z_image/model/z_image_vae/common/attention.py new file mode 100644 index 0000000..ad9278a --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/common/attention.py @@ -0,0 +1,48 @@ +import mlx.core as mx +from mlx import nn +from mlx.core.fast import scaled_dot_product_attention + +from mflux.models.common.config import ModelConfig + + +class Attention(nn.Module): + def __init__(self, channels: int = 512, num_groups: int = 32): + super().__init__() + self.channels = channels + + self.group_norm = nn.GroupNorm( + num_groups=num_groups, + dims=channels, + eps=1e-6, + affine=True, + pytorch_compatible=True, + ) + self.to_q = nn.Linear(channels, channels) + self.to_k = nn.Linear(channels, channels) + self.to_v = nn.Linear(channels, channels) + self.to_out = [nn.Linear(channels, channels)] + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + B, H, W, C = input_array.shape + + # Group norm + hidden_states = self.group_norm(input_array.astype(mx.float32)).astype(ModelConfig.precision) + + # QKV projections - reshape for attention + queries = self.to_q(hidden_states).reshape(B, H * W, 1, C) + keys = self.to_k(hidden_states).reshape(B, H * W, 1, C) + values = self.to_v(hidden_states).reshape(B, H * W, 1, C) + + # Transpose to (B, num_heads, seq_len, head_dim) + queries = mx.transpose(queries, (0, 2, 1, 3)) + keys = mx.transpose(keys, (0, 2, 1, 3)) + values = mx.transpose(values, (0, 2, 1, 3)) + + # Scaled dot product attention + scale = 1.0 / mx.sqrt(mx.array(queries.shape[-1], dtype=queries.dtype)) + hidden_states = scaled_dot_product_attention(queries, keys, values, scale=scale) + hidden_states = mx.transpose(hidden_states, (0, 2, 1, 3)).reshape(B, H, W, C) + hidden_states = self.to_out[0](hidden_states) + output = input_array + hidden_states + return mx.transpose(output, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/common/resnet_block_2d.py b/src/mflux/models/z_image/model/z_image_vae/common/resnet_block_2d.py new file mode 100644 index 0000000..d3f09ab --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/common/resnet_block_2d.py @@ -0,0 +1,70 @@ +import mlx.core as mx +from mlx import nn + + +class ResnetBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + use_conv_shortcut: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.use_conv_shortcut = use_conv_shortcut + + self.norm1 = nn.GroupNorm( + num_groups=32, + dims=in_channels, + eps=1e-6, + affine=True, + pytorch_compatible=True, + ) + self.conv1 = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + self.norm2 = nn.GroupNorm( + num_groups=32, + dims=out_channels, + eps=1e-6, + affine=True, + pytorch_compatible=True, + ) + self.conv2 = nn.Conv2d( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + # Skip connection with 1x1 conv if channels change + if use_conv_shortcut or in_channels != out_channels: + self.conv_shortcut = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + ) + else: + self.conv_shortcut = None + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + hidden_states = self.norm1(input_array) + hidden_states = nn.silu(hidden_states) + hidden_states = self.conv1(hidden_states) + hidden_states = self.norm2(hidden_states) + hidden_states = nn.silu(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_array = self.conv_shortcut(input_array) + + output = input_array + hidden_states + return mx.transpose(output, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/common/unet_mid_block.py b/src/mflux/models/z_image/model/z_image_vae/common/unet_mid_block.py new file mode 100644 index 0000000..b0f7db3 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/common/unet_mid_block.py @@ -0,0 +1,21 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_vae.common.attention import Attention +from mflux.models.z_image.model.z_image_vae.common.resnet_block_2d import ResnetBlock2D + + +class UNetMidBlock(nn.Module): + def __init__(self, channels: int = 512): + super().__init__() + self.attentions = [Attention(channels=channels)] + self.resnets = [ + ResnetBlock2D(in_channels=channels, out_channels=channels), + ResnetBlock2D(in_channels=channels, out_channels=channels), + ] + + def __call__(self, hidden_states: mx.array) -> mx.array: + hidden_states = self.resnets[0](hidden_states) + hidden_states = self.attentions[0](hidden_states) + hidden_states = self.resnets[1](hidden_states) + return hidden_states diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/__init__.py b/src/mflux/models/z_image/model/z_image_vae/decoder/__init__.py new file mode 100644 index 0000000..5e7fa85 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.model.z_image_vae.decoder.decoder import Decoder + +__all__ = ["Decoder"] diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/conv_in.py b/src/mflux/models/z_image/model/z_image_vae/decoder/conv_in.py new file mode 100644 index 0000000..8a23daa --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/conv_in.py @@ -0,0 +1,19 @@ +import mlx.core as mx +from mlx import nn + + +class ConvIn(nn.Module): + def __init__(self, in_channels: int = 16, out_channels: int = 512): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + hidden_states = self.conv(input_array) + return mx.transpose(hidden_states, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/conv_norm_out.py b/src/mflux/models/z_image/model/z_image_vae/decoder/conv_norm_out.py new file mode 100644 index 0000000..ab4893a --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/conv_norm_out.py @@ -0,0 +1,21 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.common.config import ModelConfig + + +class ConvNormOut(nn.Module): + def __init__(self, channels: int = 128): + super().__init__() + self.norm = nn.GroupNorm( + num_groups=32, + dims=channels, + eps=1e-6, + affine=True, + pytorch_compatible=True, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + hidden_states = self.norm(input_array.astype(mx.float32)).astype(ModelConfig.precision) + return mx.transpose(hidden_states, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/conv_out.py b/src/mflux/models/z_image/model/z_image_vae/decoder/conv_out.py new file mode 100644 index 0000000..001f4fe --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/conv_out.py @@ -0,0 +1,19 @@ +import mlx.core as mx +from mlx import nn + + +class ConvOut(nn.Module): + def __init__(self, in_channels: int = 128, out_channels: int = 3): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + hidden_states = self.conv(input_array) + return mx.transpose(hidden_states, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/decoder.py b/src/mflux/models/z_image/model/z_image_vae/decoder/decoder.py new file mode 100644 index 0000000..2078e04 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/decoder.py @@ -0,0 +1,36 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_vae.common.unet_mid_block import UNetMidBlock +from mflux.models.z_image.model.z_image_vae.decoder.conv_in import ConvIn +from mflux.models.z_image.model.z_image_vae.decoder.conv_norm_out import ConvNormOut +from mflux.models.z_image.model.z_image_vae.decoder.conv_out import ConvOut +from mflux.models.z_image.model.z_image_vae.decoder.up_decoder_block import UpDecoderBlock + + +class Decoder(nn.Module): + def __init__(self): + super().__init__() + self.conv_in = ConvIn(in_channels=16, out_channels=512) + self.mid_block = UNetMidBlock(channels=512) + self.up_blocks = [ + UpDecoderBlock(in_channels=512, out_channels=512, num_layers=3, add_upsample=True), + UpDecoderBlock(in_channels=512, out_channels=512, num_layers=3, add_upsample=True), + UpDecoderBlock(in_channels=512, out_channels=256, num_layers=3, add_upsample=True), + UpDecoderBlock(in_channels=256, out_channels=128, num_layers=3, add_upsample=False), + ] + + self.conv_norm_out = ConvNormOut(channels=128) + self.conv_out = ConvOut(in_channels=128, out_channels=3) + + def __call__(self, latents: mx.array) -> mx.array: + hidden_states = self.conv_in(latents) + hidden_states = self.mid_block(hidden_states) + + for up_block in self.up_blocks: + hidden_states = up_block(hidden_states) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = nn.silu(hidden_states) + hidden_states = self.conv_out(hidden_states) + return hidden_states diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/up_decoder_block.py b/src/mflux/models/z_image/model/z_image_vae/decoder/up_decoder_block.py new file mode 100644 index 0000000..c637d9f --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/up_decoder_block.py @@ -0,0 +1,47 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_vae.common.resnet_block_2d import ResnetBlock2D +from mflux.models.z_image.model.z_image_vae.decoder.up_sampler import UpSampler + + +class UpDecoderBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 3, + add_upsample: bool = True, + ): + super().__init__() + self.add_upsample = add_upsample + + # Build ResNet blocks + resnets = [] + for i in range(num_layers): + input_ch = in_channels if i == 0 else out_channels + use_shortcut = (i == 0) and (in_channels != out_channels) + resnets.append( + ResnetBlock2D( + in_channels=input_ch, + out_channels=out_channels, + use_conv_shortcut=use_shortcut, + ) + ) + self.resnets = resnets + + # Upsampler + if add_upsample: + self.upsamplers = [UpSampler(in_channels=out_channels, out_channels=out_channels)] + else: + self.upsamplers = None + + def __call__(self, hidden_states: mx.array) -> mx.array: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states) + + return hidden_states diff --git a/src/mflux/models/z_image/model/z_image_vae/decoder/up_sampler.py b/src/mflux/models/z_image/model/z_image_vae/decoder/up_sampler.py new file mode 100644 index 0000000..dde2f10 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/decoder/up_sampler.py @@ -0,0 +1,27 @@ +import mlx.core as mx +from mlx import nn + + +class UpSampler(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + hidden_states = self._upsample_nearest(input_array, scale=2) + hidden_states = self.conv(hidden_states) + return mx.transpose(hidden_states, (0, 3, 1, 2)) + + @staticmethod + def _upsample_nearest(x: mx.array, scale: int = 2) -> mx.array: + B, H, W, C = x.shape + x = mx.broadcast_to(x[:, :, None, :, None, :], (B, H, scale, W, scale, C)) + x = x.reshape(B, H * scale, W * scale, C) + return x diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/__init__.py b/src/mflux/models/z_image/model/z_image_vae/encoder/__init__.py new file mode 100644 index 0000000..7c037ba --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.model.z_image_vae.encoder.encoder import Encoder + +__all__ = ["Encoder"] diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/conv_in.py b/src/mflux/models/z_image/model/z_image_vae/encoder/conv_in.py new file mode 100644 index 0000000..1737c95 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/conv_in.py @@ -0,0 +1,19 @@ +import mlx.core as mx +from mlx import nn + + +class ConvIn(nn.Module): + def __init__(self, in_channels: int = 3, out_channels: int = 128): + super().__init__() + self.conv2d = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + output = self.conv2d(input_array) + return mx.transpose(output, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/conv_norm_out.py b/src/mflux/models/z_image/model/z_image_vae/encoder/conv_norm_out.py new file mode 100644 index 0000000..f7ff10f --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/conv_norm_out.py @@ -0,0 +1,15 @@ +import mlx.core as mx +from mlx import nn + + +class ConvNormOut(nn.Module): + def __init__(self, num_channels: int = 512, num_groups: int = 32): + super().__init__() + self.norm = nn.GroupNorm( + num_groups=num_groups, dims=num_channels, eps=1e-6, affine=True, pytorch_compatible=True + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + output = self.norm(input_array) + return mx.transpose(output, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/conv_out.py b/src/mflux/models/z_image/model/z_image_vae/encoder/conv_out.py new file mode 100644 index 0000000..a587c01 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/conv_out.py @@ -0,0 +1,19 @@ +import mlx.core as mx +from mlx import nn + + +class ConvOut(nn.Module): + def __init__(self, in_channels: int = 512, out_channels: int = 32): + super().__init__() + self.conv2d = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + input_array = mx.transpose(input_array, (0, 2, 3, 1)) + output = self.conv2d(input_array) + return mx.transpose(output, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/down_encoder_block.py b/src/mflux/models/z_image/model/z_image_vae/encoder/down_encoder_block.py new file mode 100644 index 0000000..efad39c --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/down_encoder_block.py @@ -0,0 +1,41 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_vae.common.resnet_block_2d import ResnetBlock2D +from mflux.models.z_image.model.z_image_vae.encoder.down_sampler import DownSampler + + +class DownEncoderBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_layers: int = 2, + add_downsample: bool = True, + ): + super().__init__() + + # Create ResNet blocks + self.resnets = [] + for i in range(num_layers): + res_in_channels = in_channels if i == 0 else out_channels + self.resnets.append( + ResnetBlock2D( + in_channels=res_in_channels, + out_channels=out_channels, + ) + ) + + if add_downsample: + self.downsamplers = [DownSampler(in_channels=out_channels, out_channels=out_channels)] + else: + self.downsamplers = None + + def __call__(self, hidden_states: mx.array) -> mx.array: + for resnet in self.resnets: + hidden_states = resnet(hidden_states) + + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + + return hidden_states diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/down_sampler.py b/src/mflux/models/z_image/model/z_image_vae/encoder/down_sampler.py new file mode 100644 index 0000000..2ac23c6 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/down_sampler.py @@ -0,0 +1,20 @@ +import mlx.core as mx +from mlx import nn + + +class DownSampler(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.conv = nn.Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=3, + stride=2, + padding=0, + ) + + def __call__(self, input_array: mx.array) -> mx.array: + hidden_states = mx.pad(input_array, ((0, 0), (0, 0), (0, 1), (0, 1))) + hidden_states = mx.transpose(hidden_states, (0, 2, 3, 1)) + hidden_states = self.conv(hidden_states) + return mx.transpose(hidden_states, (0, 3, 1, 2)) diff --git a/src/mflux/models/z_image/model/z_image_vae/encoder/encoder.py b/src/mflux/models/z_image/model/z_image_vae/encoder/encoder.py new file mode 100644 index 0000000..65145d9 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/encoder/encoder.py @@ -0,0 +1,35 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_vae.common.unet_mid_block import UNetMidBlock +from mflux.models.z_image.model.z_image_vae.encoder.conv_in import ConvIn +from mflux.models.z_image.model.z_image_vae.encoder.conv_norm_out import ConvNormOut +from mflux.models.z_image.model.z_image_vae.encoder.conv_out import ConvOut +from mflux.models.z_image.model.z_image_vae.encoder.down_encoder_block import DownEncoderBlock + + +class Encoder(nn.Module): + def __init__(self): + super().__init__() + self.conv_in = ConvIn(in_channels=3, out_channels=128) + self.down_blocks = [ + DownEncoderBlock(in_channels=128, out_channels=128, num_layers=2, add_downsample=True), + DownEncoderBlock(in_channels=128, out_channels=256, num_layers=2, add_downsample=True), + DownEncoderBlock(in_channels=256, out_channels=512, num_layers=2, add_downsample=True), + DownEncoderBlock(in_channels=512, out_channels=512, num_layers=2, add_downsample=False), + ] + self.mid_block = UNetMidBlock(channels=512) + self.conv_norm_out = ConvNormOut(num_channels=512, num_groups=32) + self.conv_out = ConvOut(in_channels=512, out_channels=32) # 2 * 16 latent channels + + def __call__(self, sample: mx.array) -> mx.array: + sample = self.conv_in(sample) + + for down_block in self.down_blocks: + sample = down_block(sample) + + sample = self.mid_block(sample) + sample = self.conv_norm_out(sample) + sample = nn.silu(sample) + sample = self.conv_out(sample) + return sample diff --git a/src/mflux/models/z_image/model/z_image_vae/vae.py b/src/mflux/models/z_image/model/z_image_vae/vae.py new file mode 100644 index 0000000..bc082f0 --- /dev/null +++ b/src/mflux/models/z_image/model/z_image_vae/vae.py @@ -0,0 +1,24 @@ +import mlx.core as mx +from mlx import nn + +from mflux.models.z_image.model.z_image_vae.decoder.decoder import Decoder +from mflux.models.z_image.model.z_image_vae.encoder.encoder import Encoder + + +class VAE(nn.Module): + scaling_factor: float = 0.3611 + shift_factor: float = 0.1159 + + def __init__(self): + super().__init__() + self.encoder = Encoder() + self.decoder = Decoder() + + def encode(self, image: mx.array) -> mx.array: + h = self.encoder(image) + mean, _ = mx.split(h, 2, axis=1) + return (mean - self.shift_factor) * self.scaling_factor + + def decode(self, latents: mx.array) -> mx.array: + scaled_latents = (latents / self.scaling_factor) + self.shift_factor + return self.decoder(scaled_latents) diff --git a/src/mflux/models/z_image/variants/__init__.py b/src/mflux/models/z_image/variants/__init__.py new file mode 100644 index 0000000..e53172a --- /dev/null +++ b/src/mflux/models/z_image/variants/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.variants.turbo import ZImageTurbo + +__all__ = ["ZImageTurbo"] diff --git a/src/mflux/models/z_image/variants/turbo/__init__.py b/src/mflux/models/z_image/variants/turbo/__init__.py new file mode 100644 index 0000000..e5f3e38 --- /dev/null +++ b/src/mflux/models/z_image/variants/turbo/__init__.py @@ -0,0 +1,3 @@ +from mflux.models.z_image.variants.turbo.z_image_turbo import ZImageTurbo + +__all__ = ["ZImageTurbo"] diff --git a/src/mflux/models/z_image/variants/turbo/z_image_turbo.py b/src/mflux/models/z_image/variants/turbo/z_image_turbo.py new file mode 100644 index 0000000..9a63753 --- /dev/null +++ b/src/mflux/models/z_image/variants/turbo/z_image_turbo.py @@ -0,0 +1,141 @@ +from pathlib import Path + +import mlx.core as mx +from mlx import nn +from PIL import Image + +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator +from mflux.models.common.weights.saving.model_saver import ModelSaver +from mflux.models.z_image.latent_creator import ZImageLatentCreator +from mflux.models.z_image.model.z_image_text_encoder.prompt_encoder import PromptEncoder +from mflux.models.z_image.model.z_image_text_encoder.text_encoder import TextEncoder +from mflux.models.z_image.model.z_image_transformer.transformer import ZImageTransformer +from mflux.models.z_image.model.z_image_vae.vae import VAE +from mflux.models.z_image.weights.z_image_weight_definition import ZImageWeightDefinition +from mflux.models.z_image.z_image_initializer import ZImageInitializer +from mflux.utils.exceptions import StopImageGenerationException +from mflux.utils.image_util import ImageUtil + + +class ZImageTurbo(nn.Module): + vae: VAE + text_encoder: TextEncoder + transformer: ZImageTransformer + + def __init__( + self, + quantize: int | None = None, + model_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + model_config: ModelConfig = ModelConfig.z_image_turbo(), + ): + super().__init__() + ZImageInitializer.init( + model=self, + quantize=quantize, + model_path=model_path, + lora_paths=lora_paths, + lora_scales=lora_scales, + model_config=model_config, + ) + + def generate_image( + self, + seed: int, + prompt: str, + num_inference_steps: int = 4, + height: int = 1024, + width: int = 1024, + image_path: Path | str | None = None, + image_strength: float | None = None, + scheduler: str = "linear", + ) -> Image.Image: + # 0. Create a new config based on the model type and input parameters + config = Config( + width=width, + height=height, + guidance=0.0, # Turbo model uses no guidance + scheduler=scheduler, + image_path=image_path, + image_strength=image_strength, + model_config=self.model_config, + num_inference_steps=num_inference_steps, + ) + + # 1. Create the initial latents + latents = LatentCreator.create_for_txt2img_or_img2img( + seed=seed, + width=config.width, + height=config.height, + img2img=Img2Img( + vae=self.vae, + latent_creator=ZImageLatentCreator, + image_path=config.image_path, + sigmas=config.scheduler.sigmas, + init_time_step=config.init_time_step, + ), + ) + + # 2. Encode the prompt + text_encodings = PromptEncoder.encode_prompt( + prompt=prompt, + tokenizer=self.tokenizers["z_image"], + text_encoder=self.text_encoder, + ) + + # 3. Create callback context and call before_loop + ctx = self.callbacks.start(seed=seed, prompt=prompt, config=config) + ctx.before_loop(latents) + + for t in config.time_steps: + try: + # 4.t Predict the noise + noise = self.transformer( + t=t, + x=latents, + cap_feats=text_encodings, + sigmas=config.scheduler.sigmas, + ) + + # 5.t Take one denoise step + latents = config.scheduler.step(noise=noise, timestep=t, latents=latents) + + # 6.t Call subscribers in-loop + ctx.in_loop(t, latents) + + # (Optional) Evaluate to enable progress tracking + mx.eval(latents) + + except KeyboardInterrupt: # noqa: PERF203 + ctx.interruption(t, latents) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{config.num_inference_steps}" + ) + + # 7. Call subscribers after loop + ctx.after_loop(latents) + + # 8. Decode the latents and return the image + latents = ZImageLatentCreator.unpack_latents(latents, config.height, 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, + generation_time=config.time_steps.format_dict["elapsed"], + ) + + def save_model(self, base_path: str) -> None: + ModelSaver.save_model( + model=self, + bits=self.bits, + base_path=base_path, + weight_definition=ZImageWeightDefinition, + ) diff --git a/src/mflux/models/z_image/weights/__init__.py b/src/mflux/models/z_image/weights/__init__.py new file mode 100644 index 0000000..86abc38 --- /dev/null +++ b/src/mflux/models/z_image/weights/__init__.py @@ -0,0 +1,4 @@ +from mflux.models.z_image.weights.z_image_weight_definition import ZImageWeightDefinition +from mflux.models.z_image.weights.z_image_weight_mapping import ZImageWeightMapping + +__all__ = ["ZImageWeightDefinition", "ZImageWeightMapping"] diff --git a/src/mflux/models/z_image/weights/z_image_lora_mapping.py b/src/mflux/models/z_image/weights/z_image_lora_mapping.py new file mode 100644 index 0000000..1b4c95e --- /dev/null +++ b/src/mflux/models/z_image/weights/z_image_lora_mapping.py @@ -0,0 +1,154 @@ +from mflux.models.common.lora.mapping.lora_mapping import LoRAMapping, LoRATarget + + +class ZImageLoRAMapping(LoRAMapping): + @staticmethod + def get_mapping() -> list[LoRATarget]: + targets = [] + + # Generate mappings for all layer types + for layer_type in ["layers", "noise_refiner", "context_refiner"]: + targets.extend(ZImageLoRAMapping._get_layer_targets(layer_type)) + + return targets + + @staticmethod + def _get_layer_targets(layer_type: str) -> list[LoRATarget]: + return [ + LoRATarget( + model_path=f"{layer_type}.{{block}}.adaLN_modulation.0", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.lora_B.weight", + f"{layer_type}.{{block}}.adaLN_modulation.0.lora_B.weight", + f"{layer_type}.{{block}}.adaLN_modulation.0.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.lora_A.weight", + f"{layer_type}.{{block}}.adaLN_modulation.0.lora_A.weight", + f"{layer_type}.{{block}}.adaLN_modulation.0.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.alpha", + f"{layer_type}.{{block}}.adaLN_modulation.0.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.attention.to_q", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_q.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_q.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_q.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_q.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_q.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_q.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_q.alpha", + f"{layer_type}.{{block}}.attention.to_q.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.attention.to_k", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_k.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_k.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_k.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_k.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_k.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_k.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_k.alpha", + f"{layer_type}.{{block}}.attention.to_k.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.attention.to_v", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_v.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_v.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_v.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_v.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_v.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_v.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_v.alpha", + f"{layer_type}.{{block}}.attention.to_v.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.attention.to_out.0", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_out.0.lora_B.weight", + f"{layer_type}.{{block}}.attention.to_out.0.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_out.0.lora_A.weight", + f"{layer_type}.{{block}}.attention.to_out.0.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.alpha", + f"{layer_type}.{{block}}.attention.to_out.0.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.feed_forward.w1", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.lora_B.weight", + f"{layer_type}.{{block}}.feed_forward.w1.lora_B.weight", + f"{layer_type}.{{block}}.feed_forward.w1.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.lora_A.weight", + f"{layer_type}.{{block}}.feed_forward.w1.lora_A.weight", + f"{layer_type}.{{block}}.feed_forward.w1.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.alpha", + f"{layer_type}.{{block}}.feed_forward.w1.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.feed_forward.w2", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.lora_B.weight", + f"{layer_type}.{{block}}.feed_forward.w2.lora_B.weight", + f"{layer_type}.{{block}}.feed_forward.w2.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.lora_A.weight", + f"{layer_type}.{{block}}.feed_forward.w2.lora_A.weight", + f"{layer_type}.{{block}}.feed_forward.w2.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.alpha", + f"{layer_type}.{{block}}.feed_forward.w2.alpha", + ], + ), + LoRATarget( + model_path=f"{layer_type}.{{block}}.feed_forward.w3", + possible_up_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.lora_B.weight", + f"{layer_type}.{{block}}.feed_forward.w3.lora_B.weight", + f"{layer_type}.{{block}}.feed_forward.w3.lora_up.weight", + ], + possible_down_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.lora_A.weight", + f"{layer_type}.{{block}}.feed_forward.w3.lora_A.weight", + f"{layer_type}.{{block}}.feed_forward.w3.lora_down.weight", + ], + possible_alpha_patterns=[ + f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.alpha", + f"{layer_type}.{{block}}.feed_forward.w3.alpha", + ], + ), + ] diff --git a/src/mflux/models/z_image/weights/z_image_weight_definition.py b/src/mflux/models/z_image/weights/z_image_weight_definition.py new file mode 100644 index 0000000..878843f --- /dev/null +++ b/src/mflux/models/z_image/weights/z_image_weight_definition.py @@ -0,0 +1,65 @@ +from typing import List + +from mflux.models.common.config.model_config import ModelConfig +from mflux.models.common.tokenizer import LanguageTokenizer +from mflux.models.common.weights.loading.weight_definition import ComponentDefinition, TokenizerDefinition +from mflux.models.z_image.weights.z_image_weight_mapping import ZImageWeightMapping + + +class ZImageWeightDefinition: + @staticmethod + def get_components() -> List[ComponentDefinition]: + return [ + ComponentDefinition( + name="vae", + hf_subdir="vae", + num_blocks=4, + precision=ModelConfig.precision, + mapping_getter=ZImageWeightMapping.get_vae_mapping, + ), + ComponentDefinition( + name="transformer", + hf_subdir="transformer", + num_layers=30, + precision=ModelConfig.precision, + mapping_getter=ZImageWeightMapping.get_transformer_mapping, + ), + ComponentDefinition( + name="text_encoder", + hf_subdir="text_encoder", + num_layers=36, + precision=ModelConfig.precision, + mapping_getter=ZImageWeightMapping.get_text_encoder_mapping, + ), + ] + + @staticmethod + def get_tokenizers() -> List[TokenizerDefinition]: + return [ + TokenizerDefinition( + name="z_image", + hf_subdir="tokenizer", + tokenizer_class="AutoTokenizer", + encoder_class=LanguageTokenizer, + max_length=512, + use_chat_template=True, + chat_template_kwargs={"enable_thinking": True}, + download_patterns=["tokenizer/*"], + ), + ] + + @staticmethod + def get_download_patterns() -> List[str]: + return [ + "vae/*.safetensors", + "vae/*.json", + "transformer/*.safetensors", + "transformer/*.json", + "text_encoder/*.safetensors", + "text_encoder/*.json", + "tokenizer/*", + ] + + @staticmethod + def quantization_predicate(path: str, module) -> bool: + return hasattr(module, "to_quantized") diff --git a/src/mflux/models/z_image/weights/z_image_weight_mapping.py b/src/mflux/models/z_image/weights/z_image_weight_mapping.py new file mode 100644 index 0000000..8e4053a --- /dev/null +++ b/src/mflux/models/z_image/weights/z_image_weight_mapping.py @@ -0,0 +1,615 @@ +from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget +from mflux.models.common.weights.mapping.weight_transforms import WeightTransforms + + +class ZImageWeightMapping(WeightMapping): + @staticmethod + def get_text_encoder_mapping() -> list[WeightTarget]: + return [ + WeightTarget( + to_pattern="embed_tokens.weight", + from_pattern=["model.embed_tokens.weight"], + ), + WeightTarget( + to_pattern="norm.weight", + from_pattern=["model.norm.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.input_layernorm.weight", + from_pattern=["model.layers.{layer}.input_layernorm.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.post_attention_layernorm.weight", + from_pattern=["model.layers.{layer}.post_attention_layernorm.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.self_attn.q_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.q_proj.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.self_attn.k_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.k_proj.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.self_attn.v_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.v_proj.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.self_attn.o_proj.weight", + from_pattern=["model.layers.{layer}.self_attn.o_proj.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.self_attn.q_norm.weight", + from_pattern=["model.layers.{layer}.self_attn.q_norm.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.self_attn.k_norm.weight", + from_pattern=["model.layers.{layer}.self_attn.k_norm.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.mlp.gate_proj.weight", + from_pattern=["model.layers.{layer}.mlp.gate_proj.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.mlp.up_proj.weight", + from_pattern=["model.layers.{layer}.mlp.up_proj.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.mlp.down_proj.weight", + from_pattern=["model.layers.{layer}.mlp.down_proj.weight"], + ), + ] + + @staticmethod + def get_vae_mapping() -> list[WeightTarget]: + return [ + WeightTarget( + to_pattern="encoder.conv_in.conv2d.weight", + from_pattern=["encoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.conv_in.conv2d.bias", + from_pattern=["encoder.conv_in.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.weight"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm1.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv1.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.weight"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.norm2.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.norm2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv2.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias", + from_pattern=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"], + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.downsamplers.0.conv.weight", + from_pattern=["encoder.down_blocks.{block}.downsamplers.0.conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="encoder.down_blocks.{block}.downsamplers.0.conv.bias", + from_pattern=["encoder.down_blocks.{block}.downsamplers.0.conv.bias"], + required=False, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm1.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm1.bias", + from_pattern=["encoder.mid_block.resnets.{i}.norm1.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv1.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv1.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv1.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["encoder.mid_block.resnets.{i}.norm2.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.norm2.bias", + from_pattern=["encoder.mid_block.resnets.{i}.norm2.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv2.weight", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.mid_block.resnets.{i}.conv2.bias", + from_pattern=["encoder.mid_block.resnets.{i}.conv2.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.group_norm.weight", + from_pattern=["encoder.mid_block.attentions.0.group_norm.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.group_norm.bias", + from_pattern=["encoder.mid_block.attentions.0.group_norm.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_q.weight", + from_pattern=["encoder.mid_block.attentions.0.to_q.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_q.bias", + from_pattern=["encoder.mid_block.attentions.0.to_q.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_k.weight", + from_pattern=["encoder.mid_block.attentions.0.to_k.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_k.bias", + from_pattern=["encoder.mid_block.attentions.0.to_k.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_v.weight", + from_pattern=["encoder.mid_block.attentions.0.to_v.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_v.bias", + from_pattern=["encoder.mid_block.attentions.0.to_v.bias"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_out.0.weight", + from_pattern=["encoder.mid_block.attentions.0.to_out.0.weight"], + ), + WeightTarget( + to_pattern="encoder.mid_block.attentions.0.to_out.0.bias", + from_pattern=["encoder.mid_block.attentions.0.to_out.0.bias"], + ), + WeightTarget( + to_pattern="encoder.conv_norm_out.norm.weight", + from_pattern=["encoder.conv_norm_out.weight"], + ), + WeightTarget( + to_pattern="encoder.conv_norm_out.norm.bias", + from_pattern=["encoder.conv_norm_out.bias"], + ), + WeightTarget( + to_pattern="encoder.conv_out.conv2d.weight", + from_pattern=["encoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="encoder.conv_out.conv2d.bias", + from_pattern=["encoder.conv_out.bias"], + ), + WeightTarget( + to_pattern="decoder.conv_in.conv.weight", + from_pattern=["decoder.conv_in.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.conv_in.conv.bias", + from_pattern=["decoder.conv_in.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm1.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm1.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm1.bias", + from_pattern=["decoder.mid_block.resnets.{i}.norm1.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv1.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv1.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv1.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm2.weight", + from_pattern=["decoder.mid_block.resnets.{i}.norm2.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.norm2.bias", + from_pattern=["decoder.mid_block.resnets.{i}.norm2.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv2.weight", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.mid_block.resnets.{i}.conv2.bias", + from_pattern=["decoder.mid_block.resnets.{i}.conv2.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.group_norm.weight", + from_pattern=["decoder.mid_block.attentions.0.group_norm.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.group_norm.bias", + from_pattern=["decoder.mid_block.attentions.0.group_norm.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_q.weight", + from_pattern=["decoder.mid_block.attentions.0.to_q.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_q.bias", + from_pattern=["decoder.mid_block.attentions.0.to_q.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_k.weight", + from_pattern=["decoder.mid_block.attentions.0.to_k.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_k.bias", + from_pattern=["decoder.mid_block.attentions.0.to_k.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_v.weight", + from_pattern=["decoder.mid_block.attentions.0.to_v.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_v.bias", + from_pattern=["decoder.mid_block.attentions.0.to_v.bias"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_out.0.weight", + from_pattern=["decoder.mid_block.attentions.0.to_out.0.weight"], + ), + WeightTarget( + to_pattern="decoder.mid_block.attentions.0.to_out.0.bias", + from_pattern=["decoder.mid_block.attentions.0.to_out.0.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.weight"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm1.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm1.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv1.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.weight"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.norm2.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.norm2.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv2.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"], + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias", + from_pattern=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"], + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsamplers.0.conv.weight", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.conv.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + required=False, + ), + WeightTarget( + to_pattern="decoder.up_blocks.{block}.upsamplers.0.conv.bias", + from_pattern=["decoder.up_blocks.{block}.upsamplers.0.conv.bias"], + required=False, + ), + WeightTarget( + to_pattern="decoder.conv_norm_out.norm.weight", + from_pattern=["decoder.conv_norm_out.weight"], + ), + WeightTarget( + to_pattern="decoder.conv_norm_out.norm.bias", + from_pattern=["decoder.conv_norm_out.bias"], + ), + WeightTarget( + to_pattern="decoder.conv_out.conv.weight", + from_pattern=["decoder.conv_out.weight"], + transform=WeightTransforms.transpose_conv2d_weight, + ), + WeightTarget( + to_pattern="decoder.conv_out.conv.bias", + from_pattern=["decoder.conv_out.bias"], + ), + ] + + @staticmethod + def get_transformer_mapping() -> list[WeightTarget]: + return [ + WeightTarget( + to_pattern="t_embedder.linear1.weight", + from_pattern=["t_embedder.mlp.0.weight"], + ), + WeightTarget( + to_pattern="t_embedder.linear1.bias", + from_pattern=["t_embedder.mlp.0.bias"], + ), + WeightTarget( + to_pattern="t_embedder.linear2.weight", + from_pattern=["t_embedder.mlp.2.weight"], + ), + WeightTarget( + to_pattern="t_embedder.linear2.bias", + from_pattern=["t_embedder.mlp.2.bias"], + ), + WeightTarget( + to_pattern="cap_embedder.0.weight", + from_pattern=["cap_embedder.0.weight"], + ), + WeightTarget( + to_pattern="cap_embedder.1.weight", + from_pattern=["cap_embedder.1.weight"], + ), + WeightTarget( + to_pattern="cap_embedder.1.bias", + from_pattern=["cap_embedder.1.bias"], + ), + WeightTarget( + to_pattern="x_pad_token", + from_pattern=["x_pad_token"], + ), + WeightTarget( + to_pattern="cap_pad_token", + from_pattern=["cap_pad_token"], + ), + WeightTarget( + to_pattern="all_x_embedder.2-1.weight", + from_pattern=["all_x_embedder.2-1.weight"], + ), + WeightTarget( + to_pattern="all_x_embedder.2-1.bias", + from_pattern=["all_x_embedder.2-1.bias"], + ), + WeightTarget( + to_pattern="all_final_layer.2-1.linear.weight", + from_pattern=["all_final_layer.2-1.linear.weight"], + ), + WeightTarget( + to_pattern="all_final_layer.2-1.linear.bias", + from_pattern=["all_final_layer.2-1.linear.bias"], + ), + WeightTarget( + to_pattern="all_final_layer.2-1.adaLN_modulation.0.weight", + from_pattern=["all_final_layer.2-1.adaLN_modulation.1.weight"], + ), + WeightTarget( + to_pattern="all_final_layer.2-1.adaLN_modulation.0.bias", + from_pattern=["all_final_layer.2-1.adaLN_modulation.1.bias"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.adaLN_modulation.0.weight", + from_pattern=["noise_refiner.{layer}.adaLN_modulation.0.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.adaLN_modulation.0.bias", + from_pattern=["noise_refiner.{layer}.adaLN_modulation.0.bias"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention.to_q.weight", + from_pattern=["noise_refiner.{layer}.attention.to_q.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention.to_k.weight", + from_pattern=["noise_refiner.{layer}.attention.to_k.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention.to_v.weight", + from_pattern=["noise_refiner.{layer}.attention.to_v.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention.to_out.0.weight", + from_pattern=["noise_refiner.{layer}.attention.to_out.0.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention.norm_q.weight", + from_pattern=["noise_refiner.{layer}.attention.norm_q.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention.norm_k.weight", + from_pattern=["noise_refiner.{layer}.attention.norm_k.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention_norm1.weight", + from_pattern=["noise_refiner.{layer}.attention_norm1.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.attention_norm2.weight", + from_pattern=["noise_refiner.{layer}.attention_norm2.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.ffn_norm1.weight", + from_pattern=["noise_refiner.{layer}.ffn_norm1.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.ffn_norm2.weight", + from_pattern=["noise_refiner.{layer}.ffn_norm2.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.feed_forward.w1.weight", + from_pattern=["noise_refiner.{layer}.feed_forward.w1.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.feed_forward.w2.weight", + from_pattern=["noise_refiner.{layer}.feed_forward.w2.weight"], + ), + WeightTarget( + to_pattern="noise_refiner.{layer}.feed_forward.w3.weight", + from_pattern=["noise_refiner.{layer}.feed_forward.w3.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention.to_q.weight", + from_pattern=["context_refiner.{layer}.attention.to_q.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention.to_k.weight", + from_pattern=["context_refiner.{layer}.attention.to_k.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention.to_v.weight", + from_pattern=["context_refiner.{layer}.attention.to_v.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention.to_out.0.weight", + from_pattern=["context_refiner.{layer}.attention.to_out.0.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention.norm_q.weight", + from_pattern=["context_refiner.{layer}.attention.norm_q.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention.norm_k.weight", + from_pattern=["context_refiner.{layer}.attention.norm_k.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention_norm1.weight", + from_pattern=["context_refiner.{layer}.attention_norm1.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.attention_norm2.weight", + from_pattern=["context_refiner.{layer}.attention_norm2.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.ffn_norm1.weight", + from_pattern=["context_refiner.{layer}.ffn_norm1.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.ffn_norm2.weight", + from_pattern=["context_refiner.{layer}.ffn_norm2.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.feed_forward.w1.weight", + from_pattern=["context_refiner.{layer}.feed_forward.w1.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.feed_forward.w2.weight", + from_pattern=["context_refiner.{layer}.feed_forward.w2.weight"], + ), + WeightTarget( + to_pattern="context_refiner.{layer}.feed_forward.w3.weight", + from_pattern=["context_refiner.{layer}.feed_forward.w3.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.adaLN_modulation.0.weight", + from_pattern=["layers.{layer}.adaLN_modulation.0.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.adaLN_modulation.0.bias", + from_pattern=["layers.{layer}.adaLN_modulation.0.bias"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention.to_q.weight", + from_pattern=["layers.{layer}.attention.to_q.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention.to_k.weight", + from_pattern=["layers.{layer}.attention.to_k.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention.to_v.weight", + from_pattern=["layers.{layer}.attention.to_v.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention.to_out.0.weight", + from_pattern=["layers.{layer}.attention.to_out.0.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention.norm_q.weight", + from_pattern=["layers.{layer}.attention.norm_q.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention.norm_k.weight", + from_pattern=["layers.{layer}.attention.norm_k.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention_norm1.weight", + from_pattern=["layers.{layer}.attention_norm1.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.attention_norm2.weight", + from_pattern=["layers.{layer}.attention_norm2.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.ffn_norm1.weight", + from_pattern=["layers.{layer}.ffn_norm1.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.ffn_norm2.weight", + from_pattern=["layers.{layer}.ffn_norm2.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.feed_forward.w1.weight", + from_pattern=["layers.{layer}.feed_forward.w1.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.feed_forward.w2.weight", + from_pattern=["layers.{layer}.feed_forward.w2.weight"], + ), + WeightTarget( + to_pattern="layers.{layer}.feed_forward.w3.weight", + from_pattern=["layers.{layer}.feed_forward.w3.weight"], + ), + ] diff --git a/src/mflux/models/z_image/z_image_initializer.py b/src/mflux/models/z_image/z_image_initializer.py new file mode 100644 index 0000000..64c87c7 --- /dev/null +++ b/src/mflux/models/z_image/z_image_initializer.py @@ -0,0 +1,78 @@ +from mflux.callbacks.callback_registry import CallbackRegistry +from mflux.models.common.config import ModelConfig +from mflux.models.common.lora.mapping.lora_loader import LoRALoader +from mflux.models.common.tokenizer import TokenizerLoader +from mflux.models.common.weights.loading.loaded_weights import LoadedWeights +from mflux.models.common.weights.loading.weight_applier import WeightApplier +from mflux.models.common.weights.loading.weight_loader import WeightLoader +from mflux.models.z_image.model.z_image_text_encoder.text_encoder import TextEncoder +from mflux.models.z_image.model.z_image_transformer.transformer import ZImageTransformer +from mflux.models.z_image.model.z_image_vae.vae import VAE +from mflux.models.z_image.weights.z_image_lora_mapping import ZImageLoRAMapping +from mflux.models.z_image.weights.z_image_weight_definition import ZImageWeightDefinition + + +class ZImageInitializer: + @staticmethod + def init( + model, + model_config: ModelConfig, + quantize: int | None, + model_path: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + ) -> None: + path = model_path if model_path else model_config.model_name + ZImageInitializer._init_config(model, model_config) + weights = ZImageInitializer._load_weights(path) + ZImageInitializer._init_tokenizers(model, path) + ZImageInitializer._init_models(model) + ZImageInitializer._apply_weights(model, weights, quantize) + ZImageInitializer._apply_lora(model, lora_paths, lora_scales) + + @staticmethod + def _init_config(model, model_config: ModelConfig) -> None: + model.model_config = model_config + model.callbacks = CallbackRegistry() + + @staticmethod + def _load_weights(model_path: str) -> LoadedWeights: + return WeightLoader.load( + weight_definition=ZImageWeightDefinition, + model_path=model_path, + ) + + @staticmethod + def _init_tokenizers(model, model_path: str) -> None: + model.tokenizers = TokenizerLoader.load_all( + definitions=ZImageWeightDefinition.get_tokenizers(), + model_path=model_path, + ) + + @staticmethod + def _init_models(model) -> None: + model.vae = VAE() + model.transformer = ZImageTransformer() + model.text_encoder = TextEncoder() + + @staticmethod + def _apply_weights(model, weights: LoadedWeights, quantize: int | None) -> None: + model.bits = WeightApplier.apply_and_quantize( + weights=weights, + quantize_arg=quantize, + weight_definition=ZImageWeightDefinition, + models={ + "vae": model.vae, + "transformer": model.transformer, + "text_encoder": model.text_encoder, + }, + ) + + @staticmethod + def _apply_lora(model, lora_paths: list[str] | None, lora_scales: list[float] | None) -> None: + model.lora_paths, model.lora_scales = LoRALoader.load_and_apply_lora( + lora_mapping=ZImageLoRAMapping.get_mapping(), + transformer=model.transformer, + lora_paths=lora_paths, + lora_scales=lora_scales, + ) diff --git a/src/mflux/ui/box_values.py b/src/mflux/ui/box_values.py deleted file mode 100644 index c1808ad..0000000 --- a/src/mflux/ui/box_values.py +++ /dev/null @@ -1,80 +0,0 @@ -from dataclasses import dataclass - - -@dataclass -class AbsoluteBoxValues: - top: int - right: int - bottom: int - left: int - - -@dataclass -class BoxValues: - top: int | str - right: int | str - bottom: int | str - left: int | str - - def normalize_to_dimensions(self, width, height) -> AbsoluteBoxValues: - parts = [] - dimension_base = [height, width, height, width] - for index, part in enumerate([self.top, self.right, self.bottom, self.left]): - if isinstance(part, str) and part.endswith("%"): - parts.append(int(int(part.strip("%")) / 100 * dimension_base[index])) - else: - # simple integer value - parts.append(int(part)) - return AbsoluteBoxValues(*parts) - - -class BoxValueError(ValueError): - pass - - -def parse_box_value(value, delimiter=","): - """ - Parse CSS-style padding values into a dictionary. - - Accepts formats: - - Single value (all sides) - - Two values (vertical | horizontal) - - Three values (top | horizontal | bottom) - - Four values (top | right | bottom | left) - - Args: - value: A string containing CSS-style box values - - Returns: - BoxValues dataclass with 'top', 'right', 'bottom', and 'left' attributes - """ - parts = [] - for part_value in value.strip().split(delimiter): - try: - part = int(part_value.strip()) - parts.append(part) - except ValueError: # noqa: PERF203 - # not an int - is it a %? - if (part_value := part_value.strip()).endswith("%"): - parts.append(part_value) - else: - raise BoxValueError(f"Invalid padding value: {part_value}") - - if len(parts) == 1: - # If only one value is provided, apply to all sides - return BoxValues(top=parts[0], right=parts[0], bottom=parts[0], left=parts[0]) - elif len(parts) == 2: - # If two values: first is top/bottom, second is left/right - return BoxValues(top=parts[0], right=parts[1], bottom=parts[0], left=parts[1]) - elif len(parts) == 3: - # If three values: top, left/right, bottom - return BoxValues(top=parts[0], right=parts[1], bottom=parts[2], left=parts[1]) - elif len(parts) == 4: - # If four values: top, right, bottom, left - return BoxValues(top=parts[0], right=parts[1], bottom=parts[2], left=parts[3]) - else: - raise BoxValueError( - "Invalid outpaint padding box value format: {value} " - "Expected: 1 (all-sides), 2 (top/bottom, left/right) " - "or 4 (top, right, bottom, left)values of int or percentages. e.g. 10px, 20%" - ) diff --git a/src/mflux/ui/cli/completions/README.md b/src/mflux/ui/cli/completions/README.md deleted file mode 100644 index 7d5eed0..0000000 --- a/src/mflux/ui/cli/completions/README.md +++ /dev/null @@ -1,144 +0,0 @@ -# mflux Shell Completions - -This module provides ZSH shell completion support for all mflux CLI commands. - -## Quick Start - -1. **Install completions:** - ```bash - mflux-completions - ``` - -2. **Reload your shell:** - ```bash - exec zsh - ``` - -3. **Test it:** - ```bash - mflux-generate -- - ``` - -## Features - -- **Auto-completion for all 15 mflux commands**: Generate, train, save, upscale, and more -- **Smart completions**: - - Model names (dev, schnell, dev-fill, or HuggingFace repos) - - Quantization levels (3, 4, 5, 6, 8) - - LoRA styles (storyboard, portrait, etc.) - - File paths with native completion - - Common percentage values for battery limit -- **Dynamic generation**: Completions stay in sync with code changes -- **Mutually exclusive options**: Properly handles -m/--model style options - -## Installation - -### Automatic Installation - -The simplest way is to let mflux-completions auto-detect the best location: - -```bash -mflux-completions -``` - -### Custom Directory - -Install to a specific directory: - -```bash -mflux-completions --dir ~/.config/zsh/completions -``` - -### Update Existing Installation - -Update an existing completion file: - -```bash -mflux-completions --update -``` - -## Troubleshooting - -### Check Installation - -Verify that completions are properly installed: - -```bash -mflux-completions --check -``` - -This will: -- Check if the completion file exists in your $fpath -- Verify that compinit is available -- Provide troubleshooting steps if needed - -### Common Issues - -1. **Completions not working after installation:** - - Start a new shell: `exec zsh` - -2. **Completions not updated after mflux upgrade:** - - Update completions: `mflux-completions --update` - - Then force reload: `compinit -u` - -3. **Permission denied:** - - Use `--dir` to specify a user-writable directory - -4. **Completion shows \} or other artifacts:** - - Update to the latest version: `mflux-completions --update` - -For Zsh-specific configuration issues (fpath, compinit, etc.), please refer to the [official Zsh documentation](https://zsh.sourceforge.io/Doc/Release/Completion-System.html). - -### Manual Installation - -If automatic installation fails: - -```bash -# Generate completion script -mflux-completions --print > _mflux - -# Move to a directory in your $fpath -mkdir -p ~/.zsh/completions -mv _mflux ~/.zsh/completions/ - -# Restart your shell -exec zsh -``` - -## How It Works - -The completion system: - -1. **Introspects argparse**: Dynamically extracts all command-line options from the parsers -2. **Generates ZSH functions**: Creates completion functions for each mflux command -3. **Provides context-aware completions**: Different completion strategies based on argument type -4. **Handles special cases**: Custom completions for models, LoRA styles, etc. - -## Development - -### Adding New Commands - -When adding a new mflux CLI command: - -1. Create the command with argparse as usual -2. Add it to the `commands` list in `generator.py` -3. Add the parser creation logic to `create_parser_for_command()` -4. Completions will be automatically generated! - -### Testing Completions - -Test the generator without installing: - -```bash -python -m mflux.completions.generator -``` - -### File Structure - -- `generator.py`: Core logic for generating ZSH completion scripts -- `install.py`: Installation utility with user-friendly CLI -- `README.md`: This documentation - -## Platform Support - -Currently supports ZSH (the default macOS shell since 2019). Supporting other shells is possible but is considered out of scope for the official repository. diff --git a/src/mflux/ui/cli/completions/__init__.py b/src/mflux/ui/cli/completions/__init__.py deleted file mode 100644 index 7dfa6a7..0000000 --- a/src/mflux/ui/cli/completions/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""ZSH completion generator for mflux CLI commands.""" diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py deleted file mode 100644 index 80d4205..0000000 --- a/src/mflux/ui/defaults.py +++ /dev/null @@ -1,84 +0,0 @@ -import logging -import os -import shutil -from pathlib import Path - -import platformdirs - -logger = logging.getLogger(__name__) - -BATTERY_PERCENTAGE_STOP_LIMIT = 5 -CONTROLNET_STRENGTH = 0.4 -DEFAULT_DEV_FILL_GUIDANCE = 30 -DEFAULT_DEPTH_GUIDANCE = 10 -DIMENSION_STEP_PIXELS = 16 -GUIDANCE_SCALE = 3.5 -GUIDANCE_SCALE_KONTEXT = 2.5 -HEIGHT, WIDTH = 1024, 1024 -MAX_PIXELS_WARNING_THRESHOLD = 2048 * 2048 -IMAGE_STRENGTH = 0.4 -MODEL_CHOICES = ["dev", "schnell", "krea-dev", "dev-krea", "qwen"] -MODEL_INFERENCE_STEPS = { - "dev": 25, - "krea-dev": 25, - "schnell": 4, - "qwen": 25, -} -QUANTIZE_CHOICES = [3, 5, 4, 6, 8] - - -def _migrate_legacy_cache(new_cache_dir: Path) -> None: - """Migrate legacy ~/.cache/mflux to new location if needed.""" - legacy_cache = Path.home() / ".cache" / "mflux" - - # Skip if legacy path doesn't exist or is already a symlink - if not legacy_cache.exists() or legacy_cache.is_symlink(): - return - - # Skip if we're already using the legacy path - if new_cache_dir == legacy_cache: - return - - try: - logger.warning(f"Migrating cache from {legacy_cache} to {new_cache_dir}") - - # Create new directory - new_cache_dir.mkdir(parents=True, exist_ok=True) - - # Move all contents from old to new location - for item in legacy_cache.iterdir(): - src = legacy_cache / item.name - dst = new_cache_dir / item.name - - if dst.exists(): - logger.warning(f" Skipping {item.name} (already exists in destination)") - continue - - logger.warning(f" Moving {item.name}") - shutil.move(str(src), str(dst)) - - # Remove the now-empty old directory - legacy_cache.rmdir() - - # Create symlink from old location to new location for backward compatibility - legacy_cache.parent.mkdir(parents=True, exist_ok=True) - legacy_cache.symlink_to(new_cache_dir) - logger.info(f"Created symlink: {legacy_cache} -> {new_cache_dir}") - - except (OSError, IOError, shutil.Error) as e: - logger.warning(f"Cache migration failed: {e}") - logger.info("Continuing with existing location") - - -# Determine cache directory -if os.environ.get("MFLUX_CACHE_DIR"): - # User specified cache directory (e.g. external storage) - MFLUX_CACHE_DIR = Path(os.environ["MFLUX_CACHE_DIR"]).resolve() -else: - # macOS-idiomatic cache directory @ /Users/username/Library/Caches/mflux - MFLUX_CACHE_DIR = Path(platformdirs.user_cache_dir(appname="mflux")) - -# Perform one-time migration if needed -_migrate_legacy_cache(MFLUX_CACHE_DIR) - -MFLUX_LORA_CACHE_DIR = MFLUX_CACHE_DIR / "loras" diff --git a/src/mflux/upscale.py b/src/mflux/upscale.py deleted file mode 100644 index c8086ad..0000000 --- a/src/mflux/upscale.py +++ /dev/null @@ -1,106 +0,0 @@ -import sys - -import PIL.Image - -from mflux.callbacks.callback_manager import CallbackManager -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet -from mflux.ui import defaults as ui_defaults -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils -from mflux.ui.scale_factor import ScaleFactor -from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException - - -def main(): - # 0. Parse command line arguments - parser = CommandLineParser(description="Upscale an image.") - 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, supports_dimension_scale_factor=True) - parser.add_controlnet_arguments(require_image=True) - parser.add_output_arguments() - args = parser.parse_args() - - # 1. Load the model - flux = Flux1Controlnet( - model_config=ModelConfig.dev_controlnet_upscaler(), - quantize=args.quantize, - local_path=args.path, - lora_paths=args.lora_paths, - lora_scales=args.lora_scales, - ) - - # 2. Register the optional callbacks - memory_saver = CallbackManager.register_callbacks(args=args, model=flux) - - try: - # Calculate output dimensions and handle safety warnings - width, height = _calculate_output_dimensions(args) - - for seed in args.seed: - # 3. Generate an upscaled image for each seed value - image = flux.generate_image( - seed=seed, - prompt=PromptUtils.get_effective_prompt(args), - controlnet_image_path=args.controlnet_image_path, - config=Config( - num_inference_steps=args.steps, - height=height, - width=width, - controlnet_strength=args.controlnet_strength, - ), - ) - - # 4. Save the image - image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata) - except (StopImageGenerationException, PromptFileReadError) as exc: - print(exc) - finally: - if memory_saver: - print(memory_saver.memory_stats()) - - -def _calculate_output_dimensions(args) -> tuple[int, int]: - """Calculate output dimensions from args, handling scale factors and safety warnings.""" - # Image.open is lazy/efficient, just need the dimension metadata - orig_image = PIL.Image.open(args.controlnet_image_path) - output_width, output_height = orig_image.size - - if isinstance(args.height, ScaleFactor): - output_height: int = args.height.get_scaled_value(orig_image.height) # type: ignore - - else: - output_height = args.height # type: ignore - - if isinstance(args.width, ScaleFactor): - output_width: int = args.width.get_scaled_value(orig_image.width) # type: ignore - - else: - output_width = args.width # type: ignore - - # Check if dimensions exceed safe limits - total_pixels = output_height * output_width - - if total_pixels > ui_defaults.MAX_PIXELS_WARNING_THRESHOLD: - print( - f"โš ๏ธ WARNING: The requested dimensions {output_width}x{output_height} " - f"({total_pixels:,} pixels) exceed max recommended ({ui_defaults.MAX_PIXELS_WARNING_THRESHOLD:,} pixels)." - ) - print("This generation is likely to exceed the capabilities of this computer and may:") - print(" โณ Take a very long time to complete") - print(" ๐Ÿ”ฅ Run out of memory") - print(" ๐Ÿ’ฅ Cause the program and your Mac to crash") - - user_input = input("\nPress Enter to continue at your own risk, or type 'n' to cancel: ") - if user_input.lower() in ["n", "no"]: - print("๐Ÿ›‘ Generation cancelled by user.") - sys.exit(1) - - return output_width, output_height - - -if __name__ == "__main__": - main() diff --git a/src/mflux/utils/array_util.py b/src/mflux/utils/array_util.py deleted file mode 100644 index c0dccea..0000000 --- a/src/mflux/utils/array_util.py +++ /dev/null @@ -1,17 +0,0 @@ -import mlx.core as mx - - -class ArrayUtil: - @staticmethod - def unpack_latents(latents: mx.array, height: int, width: int) -> mx.array: - latents = mx.reshape(latents, (1, height // 16, width // 16, 16, 2, 2)) - latents = mx.transpose(latents, (0, 3, 1, 4, 2, 5)) - latents = mx.reshape(latents, (1, 16, height // 16 * 2, width // 16 * 2)) - return latents - - @staticmethod - def pack_latents(latents: mx.array, height: int, width: int, num_channels_latents: int = 16) -> mx.array: - latents = mx.reshape(latents, (1, num_channels_latents, height // 16, 2, width // 16, 2)) - latents = mx.transpose(latents, (0, 2, 4, 1, 3, 5)) - latents = mx.reshape(latents, (1, (width // 16) * (height // 16), num_channels_latents * 4)) - return latents diff --git a/src/mflux/utils/box_values.py b/src/mflux/utils/box_values.py new file mode 100644 index 0000000..ebf652f --- /dev/null +++ b/src/mflux/utils/box_values.py @@ -0,0 +1,65 @@ +from dataclasses import dataclass + + +@dataclass +class AbsoluteBoxValues: + top: int + right: int + bottom: int + left: int + + +class BoxValueError(ValueError): + pass + + +@dataclass +class BoxValues: + top: int | str + right: int | str + bottom: int | str + left: int | str + + def normalize_to_dimensions(self, width, height) -> AbsoluteBoxValues: + parts = [] + dimension_base = [height, width, height, width] + for index, part in enumerate([self.top, self.right, self.bottom, self.left]): + if isinstance(part, str) and part.endswith("%"): + parts.append(int(int(part.strip("%")) / 100 * dimension_base[index])) + else: + # simple integer value + parts.append(int(part)) + return AbsoluteBoxValues(*parts) + + @staticmethod + def parse(value, delimiter=",") -> "BoxValues": + parts = [] + for part_value in value.strip().split(delimiter): + try: + part = int(part_value.strip()) + parts.append(part) + except ValueError: # noqa: PERF203 + # not an int - is it a %? + if (part_value := part_value.strip()).endswith("%"): + parts.append(part_value) + else: + raise BoxValueError(f"Invalid padding value: {part_value}") + + if len(parts) == 1: + # If only one value is provided, apply to all sides + return BoxValues(top=parts[0], right=parts[0], bottom=parts[0], left=parts[0]) + elif len(parts) == 2: + # If two values: first is top/bottom, second is left/right + return BoxValues(top=parts[0], right=parts[1], bottom=parts[0], left=parts[1]) + elif len(parts) == 3: + # If three values: top, left/right, bottom + return BoxValues(top=parts[0], right=parts[1], bottom=parts[2], left=parts[1]) + elif len(parts) == 4: + # If four values: top, right, bottom, left + return BoxValues(top=parts[0], right=parts[1], bottom=parts[2], left=parts[3]) + else: + raise BoxValueError( + "Invalid outpaint padding box value format: {value} " + "Expected: 1 (all-sides), 2 (top/bottom, left/right) " + "or 4 (top, right, bottom, left)values of int or percentages. e.g. 10px, 20%" + ) diff --git a/src/mflux/utils/dimension_resolver.py b/src/mflux/utils/dimension_resolver.py new file mode 100644 index 0000000..d714a41 --- /dev/null +++ b/src/mflux/utils/dimension_resolver.py @@ -0,0 +1,45 @@ +from pathlib import Path + +import PIL.Image + +from mflux.cli.defaults import defaults as ui_defaults +from mflux.utils.scale_factor import ScaleFactor + + +class DimensionResolver: + @staticmethod + def resolve( + height: int | ScaleFactor, + width: int | ScaleFactor, + reference_image_path: Path | str | None = None, + ) -> tuple[int, int]: + height_is_scale = isinstance(height, ScaleFactor) + width_is_scale = isinstance(width, ScaleFactor) + + # If neither dimension uses ScaleFactor, just return as-is + if not height_is_scale and not width_is_scale: + return int(width), int(height) + + # ScaleFactor requires a reference image - fall back to defaults if not provided + if reference_image_path is None: + resolved_width = ui_defaults.WIDTH if width_is_scale else int(width) + resolved_height = ui_defaults.HEIGHT if height_is_scale else int(height) + return resolved_width, resolved_height + + # Open image lazily - PIL.Image.open only reads metadata, not pixel data + with PIL.Image.open(reference_image_path) as orig_image: + orig_width, orig_height = orig_image.size + + # Resolve height + if height_is_scale: + resolved_height = height.get_scaled_value(orig_height) + else: + resolved_height = int(height) + + # Resolve width + if width_is_scale: + resolved_width = width.get_scaled_value(orig_width) + else: + resolved_width = int(width) + + return resolved_width, resolved_height diff --git a/src/mflux/utils/download.py b/src/mflux/utils/download.py deleted file mode 100644 index 96744de..0000000 --- a/src/mflux/utils/download.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Download utilities for mflux weights with cache-first behavior.""" - -from huggingface_hub import snapshot_download as hf_snapshot_download - - -def snapshot_download(repo_id: str, **kwargs) -> str: - """Download repo files with cache-first behavior. - - This wrapper function is a wrapper for upstream hf_snapshot_download - - 2025-07-08 HOT FIX: just pass the args through for now, see #235 discussion - """ - return hf_snapshot_download(repo_id, **kwargs) diff --git a/src/mflux/utils/exceptions.py b/src/mflux/utils/exceptions.py index bdcc694..e4f42fc 100644 --- a/src/mflux/utils/exceptions.py +++ b/src/mflux/utils/exceptions.py @@ -1,34 +1,32 @@ class MFluxException(Exception): - """base class for all custom exceptions in mflux package.""" + pass class ImageSavingException(MFluxException): - """error occurred while attempting to save image to storage.""" + pass class MetadataEmbedException(MFluxException): - """error occurred while attempting to embed metadata in image""" + pass class MFluxUserException(MFluxException): - """an exception raised by user behavior or intention.""" + pass class PromptFileReadError(MFluxUserException): - """Exception raised for errors in reading the prompt file.""" + pass class StopImageGenerationException(MFluxUserException): - """user has requested to stop a image generation in progress.""" + pass class StopTrainingException(MFluxUserException): - """user has requested to stop the training process in progress.""" + pass class CommandExecutionError(MFluxException): - """Raised when a subprocess command invoked by mflux tooling fails.""" - def __init__(self, cmd: list[str], return_code: int, stdout: str | None, stderr: str | None): self.cmd = cmd self.return_code = return_code @@ -40,12 +38,12 @@ class CommandExecutionError(MFluxException): class ReferenceVsOutputImageError(AssertionError): - """Raised when reference and output images don't match within the allowed threshold.""" + pass class ModelConfigError(ValueError): - """User error in model config.""" + pass class InvalidBaseModel(ModelConfigError): - """Invalid base model, cannot infer model properties.""" + pass diff --git a/src/mflux/utils/generated_image.py b/src/mflux/utils/generated_image.py index e165fde..f5bcc38 100644 --- a/src/mflux/utils/generated_image.py +++ b/src/mflux/utils/generated_image.py @@ -6,7 +6,7 @@ from pathlib import Path import mlx.core as mx import PIL.Image -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.concept_attention.attention_data import ConceptHeatmap from mflux.utils.version_util import VersionUtil diff --git a/src/mflux/utils/image_compare.py b/src/mflux/utils/image_compare.py index a898f91..68257b0 100644 --- a/src/mflux/utils/image_compare.py +++ b/src/mflux/utils/image_compare.py @@ -7,10 +7,7 @@ from PIL import Image class ImageCompare: - # How we determined DEFAULT_MISMATCH_THRESHOLD value: by eye test - # successive mlx/mflux versions have generated images that are visually close enough - # to consider as a valid upgrade. Minor visual differences are attributable to mlx updates - DEFAULT_MISMATCH_THRESHOLD = 0.15 + DEFAULT_MISMATCH_THRESHOLD = 0.15 # Set by eye test ENV_MISMATCH_THRESHOLD = float(os.environ.get("MFLUX_IMAGE_MISMATCH_THRESHOLD", DEFAULT_MISMATCH_THRESHOLD)) @staticmethod diff --git a/src/mflux/utils/image_util.py b/src/mflux/utils/image_util.py index ff9e726..b9102db 100644 --- a/src/mflux/utils/image_util.py +++ b/src/mflux/utils/image_util.py @@ -9,9 +9,9 @@ import PIL.Image import PIL.ImageDraw from PIL._typing import StrOrBytesPath -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config from mflux.models.flux.variants.concept_attention.attention_data import ConceptHeatmap -from mflux.ui.box_values import AbsoluteBoxValues, BoxValues +from mflux.utils.box_values import AbsoluteBoxValues, BoxValues from mflux.utils.generated_image import GeneratedImage from mflux.utils.metadata_builder import MetadataBuilder @@ -22,7 +22,7 @@ class ImageUtil: @staticmethod def to_image( decoded_latents: mx.array, - config: RuntimeConfig, + config: Config, seed: int, prompt: str, quantization: int, @@ -141,14 +141,11 @@ class ImageUtil: left: int | str = 0, fill_color: tuple = (255, 255, 255), ) -> PIL.Image.Image: - """ - Expand the image by padding it with the top/right/bottom/left box values specified - in either pixels or percentages relative to original image dimensions. - """ if box_values is None: box_values = BoxValues(top=top, right=right, bottom=bottom, left=left).normalize_to_dimensions( - image.width, image.height - ) # Create new image with expanded dimensions, paste the original image into it + width=image.width, + height=image.height, + ) new_width = image.width + box_values.left + box_values.right new_height = image.height + box_values.top + box_values.bottom @@ -158,10 +155,6 @@ class ImageUtil: @staticmethod def create_outpaint_mask_image(orig_width: int, orig_height: int, **create_bordered_image_kwargs): - """ - Create an outpaint mask image that is black in the middle representing the original image dimensions - and a white border on the outside paddings representing the areas to be painted over. - """ return ImageUtil.create_bordered_image( orig_width, orig_height, @@ -182,9 +175,6 @@ class ImageUtil: bottom: int | str = 0, left: int | str = 0, ) -> PIL.Image.Image: - """ - Create an image with border color and a content/fill-colored center based on CSS box model values. - """ if box_values is None: box_values = BoxValues(top=top, right=right, bottom=bottom, left=left).normalize_to_dimensions( orig_width, orig_height @@ -257,7 +247,6 @@ class ImageUtil: @staticmethod def _embed_metadata(metadata: dict, path: str | Path) -> None: - """Original EXIF metadata embedding - preserved for compatibility""" try: # Convert metadata dictionary to a string metadata_str = json.dumps(metadata) diff --git a/src/mflux/utils/info_util.py b/src/mflux/utils/info_util.py new file mode 100644 index 0000000..298fc6f --- /dev/null +++ b/src/mflux/utils/info_util.py @@ -0,0 +1,87 @@ +from datetime import datetime +from pathlib import Path + + +class InfoUtil: + @staticmethod + def format_metadata(metadata: dict) -> str: + exif = metadata.get("exif", {}) + if not exif: + return "No metadata found" + + lines = ["=" * 60, "MFLUX Image Information", "=" * 60] + + # Prompt + if prompt := exif.get("prompt"): + lines.append(f"\nPrompt: {prompt}") + + if negative_prompt := exif.get("negative_prompt"): + lines.append(f"Negative Prompt: {negative_prompt}") + + # Model information + lines.append("") + if model := exif.get("model"): + lines.append(f"Model: {model}") + + # Image dimensions + if width := exif.get("width"): + lines.append(f"Width: {width}") + if height := exif.get("height"): + lines.append(f"Height: {height}") + + # Generation parameters + lines.append("") + if seed := exif.get("seed"): + lines.append(f"Seed: {seed}") + if steps := exif.get("steps"): + lines.append(f"Steps: {steps}") + if guidance := exif.get("guidance"): + lines.append(f"Guidance: {guidance}") + + # Technical settings + if quantize := exif.get("quantize"): + lines.append(f"Quantization: {quantize}-bit") + if precision := exif.get("precision"): + lines.append(f"Precision: {precision}") + + # LoRA information + if lora_paths := exif.get("lora_paths"): + lines.append("") + lines.append(f"LoRAs ({len(lora_paths)}):") + lora_scales = exif.get("lora_scales") or [] + for i, lora in enumerate(lora_paths): + scale = lora_scales[i] if i < len(lora_scales) else 1.0 + lora_name = Path(lora).name + lines.append(f" - {lora_name} (scale: {scale})") + + # Image-to-image parameters + if image_path := exif.get("image_path"): + lines.append("") + lines.append(f"Source Image: {Path(image_path).name}") + if image_strength := exif.get("image_strength"): + lines.append(f"Image Strength: {image_strength}") + + # ControlNet parameters + if controlnet_path := exif.get("controlnet_image_path"): + lines.append("") + lines.append(f"ControlNet Image: {Path(controlnet_path).name}") + if controlnet_strength := exif.get("controlnet_strength"): + lines.append(f"ControlNet Strength: {controlnet_strength}") + + # Generation metadata + lines.append("") + if gen_time := exif.get("generation_time_seconds"): + lines.append(f"Generation Time: {gen_time:.2f}s") + + if created_at := exif.get("created_at"): + try: + dt = datetime.fromisoformat(created_at) + lines.append(f"Created: {dt.strftime('%Y-%m-%d %H:%M:%S')}") + except (ValueError, AttributeError): + lines.append(f"Created: {created_at}") + + if version := exif.get("mflux_version"): + lines.append(f"MFLUX Version: {version}") + + lines.append("=" * 60) + return "\n".join(lines) diff --git a/src/mflux/utils/lora_library_util.py b/src/mflux/utils/lora_library_util.py new file mode 100644 index 0000000..0c51895 --- /dev/null +++ b/src/mflux/utils/lora_library_util.py @@ -0,0 +1,100 @@ +import os +import sys +from collections import defaultdict +from pathlib import Path + +from mflux.models.common.resolution.lora_resolution import LoraResolution + + +class LoraLibraryUtil: + @staticmethod + def epilog() -> str: + return """ +Environment Variables: + LORA_LIBRARY_PATH Colon-separated list of directories containing .safetensors files + Example: /path/to/loras:/another/path/to/loras + +Examples: + # List all discovered LoRA files using LORA_LIBRARY_PATH + mflux-lora-library list + + # With environment variable set + LORA_LIBRARY_PATH=/home/user/loras:/opt/shared/loras mflux-lora-library list + + # Override with specific paths + mflux-lora-library list --paths /path/to/loras /another/path/to/loras +""" + + @staticmethod + def list_loras(paths: list[str] | None = None) -> int: + if paths: + # Use provided paths + library_paths = [Path(p.strip()) for p in paths] + else: + # Use environment variable + library_path_env = os.environ.get("LORA_LIBRARY_PATH") + + if not library_path_env: + print("LORA_LIBRARY_PATH environment variable is not set.", file=sys.stderr) + print( + "Set it to one or more colon-separated directories containing .safetensors files.", file=sys.stderr + ) + print("Alternatively, use --paths to specify directories directly.", file=sys.stderr) + return 1 + + # Parse library paths from environment + library_paths = [Path(p.strip()) for p in library_path_env.split(":") if p.strip()] + + # Check which paths exist + valid_paths = [] + for path in library_paths: + if path.exists() and path.is_dir(): + valid_paths.append(path) + else: + print(f"Warning: Path does not exist or is not a directory: {path}", file=sys.stderr) + + if not valid_paths: + print("No valid directories found in LORA_LIBRARY_PATH.", file=sys.stderr) + return 1 + + # Discover all LoRA files + lora_registry = LoraResolution.discover_files(valid_paths) + + if not lora_registry: + print("No .safetensors files found in the specified directories.") + return 0 + + # Sort by basename for consistent output + sorted_items = sorted(lora_registry.items()) + + # Print all discovered LoRAs + print("Discovered LoRA files:") + print("-" * 80) + for basename, full_path in sorted_items: + print(f"{basename} -> {full_path}") + + # Calculate statistics per top-level directory + stats = defaultdict(int) + for full_path in lora_registry.values(): + # Find which library path this file belongs to + for lib_path in valid_paths: + try: + # Check if full_path is relative to lib_path + full_path.relative_to(lib_path.resolve()) + stats[str(lib_path)] += 1 + break + except ValueError: + # Not relative to this lib_path, continue + continue + + # Print summary + print("-" * 80) + print(f"\nTotal LoRA files found: {len(lora_registry)}") + + if len(valid_paths) > 1: + print("\nBreakdown by library path:") + for lib_path in valid_paths: + count = stats.get(str(lib_path), 0) + print(f" {lib_path}: {count} files") + + return 0 diff --git a/src/mflux/utils/metadata_builder.py b/src/mflux/utils/metadata_builder.py index 4d4e9c3..e2f3f04 100644 --- a/src/mflux/utils/metadata_builder.py +++ b/src/mflux/utils/metadata_builder.py @@ -1,5 +1,3 @@ -"""Metadata builder for XMP and IPTC formats.""" - import logging from pathlib import Path @@ -9,21 +7,8 @@ log = logging.getLogger(__name__) class MetadataBuilder: - """Builds XMP and IPTC metadata packets for image embedding.""" - @staticmethod def embed_metadata(metadata: dict, path: str | Path) -> None: - """ - Embed XMP and IPTC metadata into an image file without touching existing EXIF. - Only supports PNG format. - - Args: - metadata: Dictionary containing image generation metadata - path: Path to the image file to embed metadata into - - Raises: - Exception: If there's an error during metadata embedding - """ # Check if file is PNG format path_obj = Path(path) if isinstance(path, str) else path if path_obj.suffix.lower() != ".png": @@ -71,15 +56,6 @@ class MetadataBuilder: @staticmethod def build_xmp_packet(metadata: dict) -> str: - """ - Build an XMP metadata packet from the provided metadata dictionary. - - Args: - metadata: Dictionary containing image generation metadata - - Returns: - XMP packet as a string - """ # Escape prompt for XML prompt_escaped = metadata.get("prompt", "").replace("&", "&").replace("<", "<").replace(">", ">") @@ -128,15 +104,6 @@ class MetadataBuilder: @staticmethod def build_iptc_binary(metadata: dict) -> bytes: - """ - Build IPTC metadata in binary format from the provided metadata dictionary. - - Args: - metadata: Dictionary containing image generation metadata - - Returns: - IPTC binary data - """ # Build LoRA info for IPTC lora_info = MetadataBuilder._build_lora_string(metadata) @@ -205,15 +172,6 @@ class MetadataBuilder: @staticmethod def _build_lora_string(metadata: dict) -> str: - """ - Build a LoRA information string from metadata. - - Args: - metadata: Dictionary containing lora_paths and lora_scales - - Returns: - Comma-separated string of LoRA names and scales, or empty string - """ # Check if lora_paths exists and is not None/empty if "lora_paths" not in metadata or metadata["lora_paths"] is None or not metadata["lora_paths"]: return "" diff --git a/src/mflux/utils/metadata_reader.py b/src/mflux/utils/metadata_reader.py index f3406c3..1a8ec0f 100644 --- a/src/mflux/utils/metadata_reader.py +++ b/src/mflux/utils/metadata_reader.py @@ -1,5 +1,3 @@ -"""Metadata reader for extracting and parsing image metadata.""" - import json import logging from pathlib import Path @@ -11,19 +9,8 @@ log = logging.getLogger(__name__) class MetadataReader: - """Reads and parses metadata from MFLUX generated images.""" - @staticmethod def read_exif_metadata(image_path: str | Path) -> dict | None: - """ - Extract EXIF metadata from an image. - - Args: - image_path: Path to the image file - - Returns: - Dictionary containing parsed EXIF metadata, or None if not found - """ try: img = PIL.Image.open(image_path) exif_bytes = img.info.get("exif") @@ -50,15 +37,6 @@ class MetadataReader: @staticmethod def read_xmp_metadata(image_path: str | Path) -> dict | None: - """ - Extract XMP metadata from an image. - - Args: - image_path: Path to the image file - - Returns: - Dictionary containing parsed XMP metadata, or None if not found - """ try: img = PIL.Image.open(image_path) xmp_data = img.info.get("XML:com.adobe.xmp") @@ -110,15 +88,6 @@ class MetadataReader: @staticmethod def read_all_metadata(image_path: str | Path) -> dict: - """ - Read all available metadata from an image. - - Args: - image_path: Path to the image file - - Returns: - Dictionary with 'exif' and 'xmp' keys containing metadata - """ return { "exif": MetadataReader.read_exif_metadata(image_path), "xmp": MetadataReader.read_xmp_metadata(image_path), diff --git a/src/mflux/ui/prompt_utils.py b/src/mflux/utils/prompt_util.py similarity index 89% rename from src/mflux/ui/prompt_utils.py rename to src/mflux/utils/prompt_util.py index 54015de..638c60c 100644 --- a/src/mflux/ui/prompt_utils.py +++ b/src/mflux/utils/prompt_util.py @@ -8,9 +8,9 @@ from mflux.utils.exceptions import PromptFileReadError logger = logging.getLogger(__name__) -class PromptUtils: +class PromptUtil: @staticmethod - def get_effective_prompt(args: Namespace) -> str: + def read_prompt(args: Namespace) -> str: # Handle stdin input when prompt is "-" if args.prompt == "-": try: @@ -24,14 +24,14 @@ class PromptUtils: # Handle prompt file if args.prompt_file is not None: - prompt = PromptUtils._read_prompt_file(args.prompt_file) + prompt = PromptUtil._read_prompt_file(args.prompt_file) return prompt # Return regular prompt return args.prompt @staticmethod - def get_effective_negative_prompt(args: Namespace) -> str: + def read_negative_prompt(args: Namespace) -> str: # Return negative prompt or empty string if not provided return getattr(args, "negative_prompt", "") diff --git a/src/mflux/ui/scale_factor.py b/src/mflux/utils/scale_factor.py similarity index 52% rename from src/mflux/ui/scale_factor.py rename to src/mflux/utils/scale_factor.py index 711ec22..f148a09 100644 --- a/src/mflux/ui/scale_factor.py +++ b/src/mflux/utils/scale_factor.py @@ -2,27 +2,7 @@ import re from dataclasses import dataclass from typing import Union -from mflux.ui.defaults import DIMENSION_STEP_PIXELS - - -@dataclass -class ScaleFactor: - value: Union[int, float] - - def __post_init__(self): - """Validate that the scale factor is positive""" - if self.value <= 0: - raise ValueError("Scale factor must be positive") - - def __str__(self): - """String representation as multiplier""" - if isinstance(self.value, int) or self.value.is_integer(): - return f"{int(self.value)}x" - return f"{self.value}x" - - def get_scaled_value(self, orig_value, pixel_steps=DIMENSION_STEP_PIXELS) -> int: - return int(self.value * orig_value - (self.value * orig_value) % pixel_steps) - +from mflux.cli.defaults.defaults import DIMENSION_STEP_PIXELS # Regex pattern for scale factors SCALE_FACTOR_PATTERN = re.compile( @@ -31,18 +11,34 @@ SCALE_FACTOR_PATTERN = re.compile( ) -def parse_scale_factor(text: str) -> ScaleFactor: - """Parse a scale factor string into a ScaleFactor dataclass""" - match = SCALE_FACTOR_PATTERN.match(text.strip()) - if not match: - raise ValueError(f"Invalid scale factor format: '{text}'. Expected format: '2x', '1.5x', etc.") +@dataclass +class ScaleFactor: + value: Union[int, float] - value_str = match.group(1) + def __post_init__(self): + if self.value <= 0: + raise ValueError("Scale factor must be positive") - # Convert to int if it's a whole number, otherwise float - if "." in value_str: - value = float(value_str) - else: - value = int(value_str) + def __str__(self): + if isinstance(self.value, int) or self.value.is_integer(): + return f"{int(self.value)}x" + return f"{self.value}x" - return ScaleFactor(value) + def get_scaled_value(self, orig_value, pixel_steps=DIMENSION_STEP_PIXELS) -> int: + return int(self.value * orig_value - (self.value * orig_value) % pixel_steps) + + @staticmethod + def parse(text: str) -> "ScaleFactor": + match = SCALE_FACTOR_PATTERN.match(text.strip()) + if not match: + raise ValueError(f"Invalid scale factor format: '{text}'. Expected format: '2x', '1.5x', etc.") + + value_str = match.group(1) + + # Convert to int if it's a whole number, otherwise float + if "." in value_str: + value = float(value_str) + else: + value = int(value_str) + + return ScaleFactor(value) diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index d4b0e25..77c8ad8 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/tests/arg_parser/test_cli_argparser.py @@ -5,9 +5,9 @@ from unittest.mock import patch import pytest -from mflux.ui import defaults as ui_defaults -from mflux.ui.box_values import BoxValues -from mflux.ui.cli.parsers import CommandLineParser +from mflux.cli.defaults import defaults as ui_defaults +from mflux.cli.parser.parsers import CommandLineParser +from mflux.utils.box_values import BoxValues def _create_mflux_generate_parser(with_controlnet=False, require_model_arg=False) -> CommandLineParser: @@ -214,12 +214,14 @@ def mflux_in_context_edit_minimal_argv() -> list[str]: ] +@pytest.mark.fast 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"): assert pytest.raises(SystemExit, mflux_generate_parser.parse_args) +@pytest.mark.fast def test_model_arg_not_in_file(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): metadata_file = temp_dir / "model.json" with metadata_file.open("wt") as m: @@ -245,6 +247,7 @@ def test_model_arg_not_in_file(mflux_generate_parser, mflux_generate_minimal_arg assert args.base_model is None +@pytest.mark.fast def test_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): metadata_file = temp_dir / "model.json" with metadata_file.open("wt") as m: @@ -260,6 +263,7 @@ def test_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_argv, b assert args.model == "schnell" +@pytest.mark.fast def test_base_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): metadata_file = temp_dir / "model.json" with metadata_file.open("wt") as m: @@ -279,6 +283,7 @@ def test_base_model_arg_in_file(mflux_generate_parser, mflux_generate_minimal_ar assert args.base_model == "schnell" +@pytest.mark.fast def test_prompt_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): metadata_file = temp_dir / "prompt.json" file_prompt = "origin of the universe" @@ -296,6 +301,7 @@ def test_prompt_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_met assert args.prompt == cli_prompt +@pytest.mark.fast def test_prompt_file_arg(mflux_generate_parser, mflux_generate_minimal_argv, temp_dir): # Create a prompt file prompt_content = "prompt from a file being re-read for each generation" @@ -310,6 +316,7 @@ def test_prompt_file_arg(mflux_generate_parser, mflux_generate_minimal_argv, tem assert args.prompt is None # prompt should be None since we're using prompt-file +@pytest.mark.fast def test_prompt_and_prompt_file_mutually_exclusive(mflux_generate_parser, temp_dir): # Create a prompt file prompt_file = temp_dir / "prompt.txt" @@ -322,6 +329,7 @@ def test_prompt_and_prompt_file_mutually_exclusive(mflux_generate_parser, temp_d mflux_generate_parser.parse_args() +@pytest.mark.fast def test_guidance_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "guidance.json" with metadata_file.open("wt") as m: @@ -337,6 +345,7 @@ def test_guidance_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_m assert args.guidance == pytest.approx(5.0) +@pytest.mark.fast def test_quantize_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "quantize.json" with metadata_file.open("wt") as m: @@ -352,6 +361,7 @@ def test_quantize_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_m assert args.quantize == 8 +@pytest.mark.fast def test_seed_arg(mflux_generate_parser, mflux_generate_minimal_model_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "seed.json" with metadata_file.open("wt") as m: @@ -383,6 +393,7 @@ def test_seed_arg(mflux_generate_parser, mflux_generate_minimal_model_argv, base assert "_seed_{seed}" not in args.output +@pytest.mark.fast def test_auto_seeds_arg(mflux_generate_parser, mflux_generate_minimal_model_argv): with patch("sys.argv", mflux_generate_minimal_model_argv + ["--seed", "24", "48", "--auto-seeds", "5"]): args = mflux_generate_parser.parse_args() @@ -401,6 +412,7 @@ def test_auto_seeds_arg(mflux_generate_parser, mflux_generate_minimal_model_argv assert isinstance(_, int) +@pytest.mark.fast def test_steps_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "steps.json" with metadata_file.open("wt") as m: @@ -428,6 +440,7 @@ def test_steps_arg(mflux_generate_parser, mflux_generate_minimal_argv, base_meta assert args.steps == 12 +@pytest.mark.fast def test_lora_args(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off test_paths = ["/some/lora/1.safetensors", "/some/lora/2.safetensors"] metadata_file = temp_dir / "lora_args.json" @@ -442,8 +455,8 @@ def test_lora_args(mflux_generate_parser, mflux_generate_minimal_argv, base_meta assert args.lora_paths is None assert args.lora_scales is None - # Mock get_lora_path to bypass file validation for test purposes - with patch("mflux.ui.cli.parsers.get_lora_path", side_effect=lambda x: x): + # Mock LoraResolution.resolve to bypass file validation for test purposes + with patch("mflux.cli.parser.parsers.LoraResolution.resolve", side_effect=lambda x: x): # test metadata config accepted with patch('sys.argv', mflux_generate_minimal_argv + ['--config-from-metadata', metadata_file.as_posix()]): # fmt: off args = mflux_generate_parser.parse_args() @@ -467,6 +480,7 @@ def test_lora_args(mflux_generate_parser, mflux_generate_minimal_argv, base_meta assert args.lora_scales == [pytest.approx(v) for v in [0.3, 0.7, 0.1, 0.9]] +@pytest.mark.fast def test_image_to_image_args(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "image_to_image.json" test_path = "/some/awesome/image.png" @@ -499,6 +513,7 @@ def test_image_to_image_args(mflux_generate_parser, mflux_generate_minimal_argv, assert args.image_strength == 0.4 # default +@pytest.mark.fast def test_image_outpaint_args(mflux_generate_parser, mflux_generate_minimal_argv, base_metadata_dict, temp_dir): # fmt: off metadata_file = temp_dir / "image_outpaint.json" test_padding = "10,20,30,40" @@ -533,6 +548,7 @@ def test_image_outpaint_args(mflux_generate_parser, mflux_generate_minimal_argv, assert args.image_outpaint_padding == BoxValues("10%", 50, "20%", 50) +@pytest.mark.fast def test_controlnet_args(mflux_generate_controlnet_parser, mflux_generate_controlnet_minimal_argv, base_metadata_dict, temp_dir): # fmt: off test_path = "/some/cnet/1.safetensors" metadata_file = temp_dir / "cnet_args.json" @@ -572,6 +588,7 @@ def test_controlnet_args(mflux_generate_controlnet_parser, mflux_generate_contro assert args.controlnet_save_canny is False +@pytest.mark.fast def test_save_args(mflux_save_parser): with patch("sys.argv", ["mflux-save", "--model", "dev"]): # required --path not provided, exits to error @@ -582,6 +599,7 @@ def test_save_args(mflux_save_parser): assert args.path == "/some/model/folder" +@pytest.mark.fast def test_fill_args(mflux_fill_parser, mflux_fill_minimal_argv): # Test required arguments with patch("sys.argv", mflux_fill_minimal_argv): @@ -589,7 +607,7 @@ def test_fill_args(mflux_fill_parser, mflux_fill_minimal_argv): assert args.prompt == "meaning of life" assert args.image_path == Path("image.png") assert args.masked_image_path == Path("mask.png") - # Default guidance for fill should be None (will be set to 30 in generate_fill.py) + # Default guidance for fill should be None (will be set to 30 in flux_generate_fill.py) assert args.guidance is None # Parser doesn't set default, app does # Test with missing required arguments @@ -611,6 +629,7 @@ def test_fill_args(mflux_fill_parser, mflux_fill_minimal_argv): assert args.width == 512 +@pytest.mark.fast def test_fill_args_with_metadata(mflux_fill_parser, mflux_fill_minimal_argv, base_metadata_dict, temp_dir): metadata_file = temp_dir / "fill_metadata.json" # Set up metadata with fill-related values @@ -647,8 +666,9 @@ def test_fill_args_with_metadata(mflux_fill_parser, mflux_fill_minimal_argv, bas assert args.masked_image_path == Path("cli_mask.png") # From CLI +@pytest.mark.fast def test_fill_default_guidance(): - # Create a parser just like in generate_fill.py + # Create a parser just like in flux_generate_fill.py parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.") parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) @@ -666,7 +686,7 @@ def test_fill_default_guidance(): # Verify initial guidance value is None (no default set by parser) assert args.guidance is None - # Simulate what happens in generate_fill.py + # Simulate what happens in flux_generate_fill.py if args.guidance is None: args.guidance = ui_defaults.DEFAULT_DEV_FILL_GUIDANCE @@ -674,6 +694,7 @@ def test_fill_default_guidance(): assert args.guidance == ui_defaults.DEFAULT_DEV_FILL_GUIDANCE +@pytest.mark.fast def test_scheduler_by_name(mflux_generate_parser, mflux_generate_minimal_model_argv): with patch("sys.argv", mflux_generate_minimal_model_argv + ["--scheduler", "linear"]): args = mflux_generate_parser.parse_args() @@ -681,14 +702,15 @@ def test_scheduler_by_name(mflux_generate_parser, mflux_generate_minimal_model_a assert args.scheduler == "linear" with patch("sys.argv", mflux_generate_minimal_model_argv + ["--scheduler", "foobar"]): - # let the CLI will accept "foobar" - later RuntimeConfig will validate + # let the CLI will accept "foobar" - later Config will validate args = mflux_generate_parser.parse_args() with patch("sys.argv", mflux_generate_minimal_model_argv + ["--scheduler", "mflux.someplugin.SchedulerBaz"]): - # let the CLI will accept external import paths - later RuntimeConfig will validate + # let the CLI will accept external import paths - later Config will validate args = mflux_generate_parser.parse_args() +@pytest.mark.fast def test_save_depth_args(mflux_save_depth_parser, mflux_save_depth_minimal_argv): # Test required arguments with patch("sys.argv", mflux_save_depth_minimal_argv): @@ -709,6 +731,7 @@ def test_save_depth_args(mflux_save_depth_parser, mflux_save_depth_minimal_argv) assert args.output == "depth_map.png" +@pytest.mark.fast def test_redux_args(mflux_redux_parser, mflux_redux_minimal_argv): # Test required arguments with patch("sys.argv", mflux_redux_minimal_argv): @@ -754,6 +777,7 @@ def test_redux_args(mflux_redux_parser, mflux_redux_minimal_argv): assert args.output == "redux_result.png" +@pytest.mark.fast def test_concept_attention_args(mflux_concept_parser, mflux_concept_minimal_argv): # Test required arguments with patch("sys.argv", mflux_concept_minimal_argv): @@ -791,6 +815,7 @@ def test_concept_attention_args(mflux_concept_parser, mflux_concept_minimal_argv assert args.heatmap_timesteps == [0, 1, 2] +@pytest.mark.fast def test_catvton_args(mflux_catvton_parser, mflux_catvton_minimal_argv): # Test required arguments with patch("sys.argv", mflux_catvton_minimal_argv): @@ -826,6 +851,7 @@ def test_catvton_args(mflux_catvton_parser, mflux_catvton_minimal_argv): assert args.vae_tiling_split == "horizontal" # Default value from parser +@pytest.mark.fast 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): @@ -879,6 +905,7 @@ def test_in_context_edit_args(mflux_in_context_edit_parser, mflux_in_context_edi assert args.vae_tiling_split == "horizontal" # Default value from parser +@pytest.mark.fast 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): @@ -888,3 +915,334 @@ def test_in_context_args(mflux_catvton_parser, mflux_catvton_minimal_argv): with patch("sys.argv", mflux_catvton_minimal_argv + ["--save-full-image"]): args = mflux_catvton_parser.parse_args() assert args.save_full_image is True + + +# ============================================================================ +# Qwen Image Tests +# ============================================================================ + + +@pytest.fixture +def mflux_qwen_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image using Qwen Image model.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments(required=False) + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_qwen_minimal_argv() -> list[str]: + return ["mflux-generate-qwen", "--prompt", "a beautiful sunset"] + + +@pytest.mark.fast +def test_qwen_args(mflux_qwen_parser, mflux_qwen_minimal_argv): + # Test minimal arguments + with patch("sys.argv", mflux_qwen_minimal_argv): + args = mflux_qwen_parser.parse_args() + assert args.prompt == "a beautiful sunset" + + # Test with image-to-image + with patch("sys.argv", mflux_qwen_minimal_argv + ["--image-path", "input.png", "--image-strength", "0.5"]): + args = mflux_qwen_parser.parse_args() + assert args.prompt == "a beautiful sunset" + assert args.image_path == Path("input.png") + assert args.image_strength == pytest.approx(0.5) + + # Test with quantization + with patch("sys.argv", mflux_qwen_minimal_argv + ["--quantize", "8"]): + args = mflux_qwen_parser.parse_args() + assert args.quantize == 8 + + +# ============================================================================ +# Qwen Image Edit Tests +# ============================================================================ + + +@pytest.fixture +def mflux_qwen_edit_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image using Qwen Image Edit.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_argument("--image-paths", type=Path, nargs="+", required=True, help="Local paths to init images") + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_qwen_edit_minimal_argv() -> list[str]: + return ["mflux-generate-qwen-edit", "--prompt", "make it blue", "--image-paths", "image1.png"] + + +@pytest.mark.fast +def test_qwen_edit_args(mflux_qwen_edit_parser, mflux_qwen_edit_minimal_argv): + # Test minimal arguments + with patch("sys.argv", mflux_qwen_edit_minimal_argv): + args = mflux_qwen_edit_parser.parse_args() + assert args.prompt == "make it blue" + assert len(args.image_paths) == 1 + assert args.image_paths[0] == Path("image1.png") + + # Test with multiple images + with patch( + "sys.argv", + ["mflux-generate-qwen-edit", "--prompt", "combine these", "--image-paths", "img1.png", "img2.png", "img3.png"], + ): + args = mflux_qwen_edit_parser.parse_args() + assert args.prompt == "combine these" + assert len(args.image_paths) == 3 + + # Test missing required image-paths + with patch("sys.argv", ["mflux-generate-qwen-edit", "--prompt", "test"]): + pytest.raises(SystemExit, mflux_qwen_edit_parser.parse_args) + + +# ============================================================================ +# FIBO Tests +# ============================================================================ + + +@pytest.fixture +def mflux_fibo_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image using FIBO model.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments(required=False) + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_fibo_minimal_argv() -> list[str]: + return ["mflux-generate-fibo", "--prompt", '{"scene": "a forest"}'] + + +@pytest.mark.fast +def test_fibo_args(mflux_fibo_parser, mflux_fibo_minimal_argv): + # Test minimal arguments + with patch("sys.argv", mflux_fibo_minimal_argv): + args = mflux_fibo_parser.parse_args() + assert args.prompt == '{"scene": "a forest"}' + + # Test with negative prompt + with patch("sys.argv", mflux_fibo_minimal_argv + ["--negative-prompt", "blurry, low quality"]): + args = mflux_fibo_parser.parse_args() + assert args.negative_prompt == "blurry, low quality" + + # Test with image-to-image + with patch("sys.argv", mflux_fibo_minimal_argv + ["--image-path", "input.png"]): + args = mflux_fibo_parser.parse_args() + assert args.image_path == Path("input.png") + + +# ============================================================================ +# Z-Image Turbo Tests +# ============================================================================ + + +@pytest.fixture +def mflux_z_image_turbo_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image using Z-Image Turbo.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments(required=False) + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_z_image_turbo_minimal_argv() -> list[str]: + return ["mflux-generate-z-image-turbo", "--prompt", "a cat in space"] + + +@pytest.mark.fast +def test_z_image_turbo_args(mflux_z_image_turbo_parser, mflux_z_image_turbo_minimal_argv): + # Test minimal arguments + with patch("sys.argv", mflux_z_image_turbo_minimal_argv): + args = mflux_z_image_turbo_parser.parse_args() + assert args.prompt == "a cat in space" + + # Test with quantization + with patch("sys.argv", mflux_z_image_turbo_minimal_argv + ["--quantize", "8"]): + args = mflux_z_image_turbo_parser.parse_args() + assert args.quantize == 8 + + # Test with image-to-image + with patch("sys.argv", mflux_z_image_turbo_minimal_argv + ["--image-path", "input.png", "--image-strength", "0.6"]): + args = mflux_z_image_turbo_parser.parse_args() + assert args.image_path == Path("input.png") + assert args.image_strength == pytest.approx(0.6) + + # Mock LoraResolution.resolve to bypass file validation for test purposes + with patch("mflux.cli.parser.parsers.LoraResolution.resolve", side_effect=lambda x: x): + # Test with LoRA + with patch( + "sys.argv", + mflux_z_image_turbo_minimal_argv + ["--lora-paths", "some/lora.safetensors", "--lora-scales", "0.8"], + ): + args = mflux_z_image_turbo_parser.parse_args() + assert args.lora_paths == ["some/lora.safetensors"] + assert args.lora_scales == [pytest.approx(0.8)] + + +# ============================================================================ +# Kontext Tests +# ============================================================================ + + +@pytest.fixture +def mflux_kontext_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image using Flux Kontext.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_lora_arguments() + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments(required=True) + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_kontext_minimal_argv() -> list[str]: + return ["mflux-generate-kontext", "--prompt", "make the background blue", "--image-path", "input.png"] + + +@pytest.mark.fast +def test_kontext_args(mflux_kontext_parser, mflux_kontext_minimal_argv): + # Test minimal arguments + with patch("sys.argv", mflux_kontext_minimal_argv): + args = mflux_kontext_parser.parse_args() + assert args.prompt == "make the background blue" + assert args.image_path == Path("input.png") + + # Test missing required image-path + with patch("sys.argv", ["mflux-generate-kontext", "--prompt", "test"]): + pytest.raises(SystemExit, mflux_kontext_parser.parse_args) + + # Test with custom parameters + with patch("sys.argv", mflux_kontext_minimal_argv + ["--steps", "15", "--guidance", "3.5"]): + args = mflux_kontext_parser.parse_args() + assert args.steps == 15 + assert args.guidance == pytest.approx(3.5) + + +# ============================================================================ +# Depth Tests +# ============================================================================ + + +@pytest.fixture +def mflux_depth_parser() -> CommandLineParser: + parser = CommandLineParser(description="Generate an image using Flux Depth.") + 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) + parser.add_depth_arguments() + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_depth_minimal_argv() -> list[str]: + return ["mflux-generate-depth", "--prompt", "a portrait", "--depth-image-path", "depth.png"] + + +@pytest.mark.fast +def test_depth_args(mflux_depth_parser, mflux_depth_minimal_argv): + # Test minimal arguments with depth image + with patch("sys.argv", mflux_depth_minimal_argv): + args = mflux_depth_parser.parse_args() + assert args.prompt == "a portrait" + assert args.depth_image_path == Path("depth.png") + + # Test with source image (depth will be generated from it) + with patch("sys.argv", ["mflux-generate-depth", "--prompt", "test", "--image-path", "source.png"]): + args = mflux_depth_parser.parse_args() + assert args.image_path == Path("source.png") + assert args.depth_image_path is None + + # Test with save-depth-map flag + with patch("sys.argv", mflux_depth_minimal_argv + ["--save-depth-map"]): + args = mflux_depth_parser.parse_args() + assert args.save_depth_map is True + + +# ============================================================================ +# Info Tests +# ============================================================================ + + +@pytest.fixture +def mflux_info_parser() -> CommandLineParser: + parser = CommandLineParser(description="Show information about an image.") + parser.add_info_arguments() + return parser + + +@pytest.fixture +def mflux_info_minimal_argv() -> list[str]: + return ["mflux-info", "output.png"] # positional argument, not --image-path + + +@pytest.mark.fast +def test_info_args(mflux_info_parser, mflux_info_minimal_argv): + # Test minimal arguments (positional image_path) + with patch("sys.argv", mflux_info_minimal_argv): + args = mflux_info_parser.parse_args() + assert args.image_path == "output.png" + + # Test missing required positional argument + with patch("sys.argv", ["mflux-info"]): + pytest.raises(SystemExit, mflux_info_parser.parse_args) + + +# ============================================================================ +# Upscale Tests +# ============================================================================ + + +@pytest.fixture +def mflux_upscale_parser() -> CommandLineParser: + parser = CommandLineParser(description="Upscale an image using Flux Controlnet.") + parser.add_general_arguments() + parser.add_model_arguments(require_model_arg=False) + parser.add_image_generator_arguments(supports_metadata_config=True) + parser.add_image_to_image_arguments(required=True) + parser.add_output_arguments() + return parser + + +@pytest.fixture +def mflux_upscale_minimal_argv() -> list[str]: + return ["mflux-upscale", "--prompt", "enhance details", "--image-path", "low_res.png"] + + +@pytest.mark.fast +def test_upscale_args(mflux_upscale_parser, mflux_upscale_minimal_argv): + # Test minimal arguments + with patch("sys.argv", mflux_upscale_minimal_argv): + args = mflux_upscale_parser.parse_args() + assert args.prompt == "enhance details" + assert args.image_path == Path("low_res.png") + + # Test missing required image-path + with patch("sys.argv", ["mflux-upscale", "--prompt", "test"]): + pytest.raises(SystemExit, mflux_upscale_parser.parse_args) + + # Test with custom dimensions + with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "1024", "--width", "1024"]): + args = mflux_upscale_parser.parse_args() + assert args.height == 1024 + assert args.width == 1024 diff --git a/tests/arg_parser/test_lora_library_integration.py b/tests/arg_parser/test_lora_library_integration.py index 976e9d3..8fbe3b5 100644 --- a/tests/arg_parser/test_lora_library_integration.py +++ b/tests/arg_parser/test_lora_library_integration.py @@ -6,13 +6,12 @@ from unittest import mock import pytest -from mflux.models.common.lora.download import lora_library -from mflux.ui.cli.parsers import CommandLineParser +from mflux.cli.parser.parsers import CommandLineParser +from mflux.models.common.resolution import lora_resolution @pytest.fixture def temp_lora_library(): - """Create a temporary lora library for testing.""" with tempfile.TemporaryDirectory() as temp_dir: lib_path = Path(temp_dir) / "lora_library" lib_path.mkdir() @@ -26,12 +25,12 @@ def temp_lora_library(): yield lib_path +@pytest.mark.fast def test_parser_resolves_lora_paths_from_library(temp_lora_library): - """Test that parser resolves lora names to full paths from library.""" # Set up the environment variable with mock.patch.dict(os.environ, {"LORA_LIBRARY_PATH": str(temp_lora_library)}): # Re-initialize the registry - lora_library._initialize_registry() + lora_resolution.LoraResolution._initialize_registry() parser = CommandLineParser() parser.add_model_arguments() @@ -63,8 +62,8 @@ def test_parser_resolves_lora_paths_from_library(temp_lora_library): assert args.lora_scales == [0.5, 0.8] +@pytest.mark.fast def test_parser_preserves_full_paths(temp_lora_library): - """Test that parser preserves full paths that already exist.""" # Create a lora file outside the library with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as tmp_file: external_lora = tmp_file.name @@ -85,15 +84,15 @@ def test_parser_preserves_full_paths(temp_lora_library): os.unlink(external_lora) +@pytest.mark.fast def test_parser_mixed_lora_paths(temp_lora_library): - """Test parser with mix of library names and full paths.""" # Create an external lora file with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as tmp_file: external_lora = tmp_file.name try: with mock.patch.dict(os.environ, {"LORA_LIBRARY_PATH": str(temp_lora_library)}): - lora_library._initialize_registry() + lora_resolution.LoraResolution._initialize_registry() parser = CommandLineParser() parser.add_model_arguments() @@ -122,11 +121,11 @@ def test_parser_mixed_lora_paths(temp_lora_library): os.unlink(external_lora) +@pytest.mark.fast def test_parser_unknown_lora_names_error(): - """Test that unknown lora names raise an error.""" # No library path set with mock.patch.dict(os.environ, {}, clear=True): - lora_library._initialize_registry() + lora_resolution.LoraResolution._initialize_registry() parser = CommandLineParser() parser.add_model_arguments() diff --git a/tests/arg_parser/test_stdin_prompt.py b/tests/arg_parser/test_stdin_prompt.py index 8fe179f..922fdf7 100644 --- a/tests/arg_parser/test_stdin_prompt.py +++ b/tests/arg_parser/test_stdin_prompt.py @@ -4,8 +4,8 @@ from unittest.mock import patch import pytest -from mflux.ui.cli.parsers import CommandLineParser -from mflux.ui.prompt_utils import PromptUtils +from mflux.cli.parser.parsers import CommandLineParser +from mflux.utils.prompt_util import PromptUtil @pytest.fixture @@ -26,8 +26,8 @@ def temp_output_dir(tmp_path_factory) -> Path: return tmp_path_factory.mktemp("mflux_stdin_test") +@pytest.mark.fast def test_prompt_from_stdin(mflux_generate_parser): - """Test that --prompt - reads from stdin correctly.""" stdin_content = "A beautiful sunset over the ocean" # Simulate stdin input @@ -35,35 +35,34 @@ def test_prompt_from_stdin(mflux_generate_parser): with patch("sys.argv", ["mflux-generate", "--prompt", "-", "--model", "dev"]): args = mflux_generate_parser.parse_args() - # The parser returns the raw args, get_effective_prompt handles stdin + # The parser returns the raw args, read_prompt handles stdin assert args.prompt == "-" +@pytest.mark.fast def test_prompt_stdin_vs_regular(mflux_generate_parser): - """Test that regular prompt still works when not using stdin.""" regular_prompt = "A regular prompt not from stdin" with patch("sys.argv", ["mflux-generate", "--prompt", regular_prompt, "--model", "dev"]): args = mflux_generate_parser.parse_args() assert args.prompt == regular_prompt - assert PromptUtils.get_effective_prompt(args) == regular_prompt + assert PromptUtil.read_prompt(args) == regular_prompt +@pytest.mark.fast def test_prompt_stdin_with_whitespace(mflux_generate_parser): - """Test that stdin prompt with surrounding whitespace is properly stripped.""" stdin_content = "\n\n A prompt with whitespace \n\n" expected_prompt = "A prompt with whitespace" with patch("sys.stdin", StringIO(stdin_content)): with patch("sys.argv", ["mflux-generate", "--prompt", "-", "--model", "dev"]): args = mflux_generate_parser.parse_args() - effective_prompt = PromptUtils.get_effective_prompt(args) + effective_prompt = PromptUtil.read_prompt(args) assert effective_prompt == expected_prompt +@pytest.mark.fast def test_prompt_file_takes_precedence_over_stdin(mflux_generate_parser, temp_output_dir): - """Test that --prompt-file still works and takes precedence over stdin detection. - because --prompt is not used in this scenario.""" # Create a prompt file prompt_file = temp_output_dir / "prompt.txt" file_prompt = "Prompt from file" @@ -73,5 +72,5 @@ def test_prompt_file_takes_precedence_over_stdin(mflux_generate_parser, temp_out with patch("sys.stdin", StringIO(stdin_content)): with patch("sys.argv", ["mflux-generate", "--prompt-file", str(prompt_file), "--model", "dev"]): args = mflux_generate_parser.parse_args() - effective_prompt = PromptUtils.get_effective_prompt(args) + effective_prompt = PromptUtil.read_prompt(args) assert effective_prompt == file_prompt diff --git a/tests/arg_parser/test_upscale_argparser.py b/tests/arg_parser/test_upscale_argparser.py index 1877775..dd83445 100644 --- a/tests/arg_parser/test_upscale_argparser.py +++ b/tests/arg_parser/test_upscale_argparser.py @@ -3,12 +3,11 @@ from unittest.mock import patch import pytest -from mflux.ui.cli.parsers import CommandLineParser, int_or_special_value -from mflux.ui.scale_factor import ScaleFactor +from mflux.cli.parser.parsers import CommandLineParser, int_or_special_value +from mflux.utils.scale_factor import ScaleFactor def _create_custom_upscale_parser() -> CommandLineParser: - """Create parser with custom dimension scale factor support""" parser = CommandLineParser(description="Generate an upscaled image from a source image") parser.add_general_arguments() parser.add_model_arguments(require_model_arg=False) @@ -48,8 +47,8 @@ def mflux_upscale_minimal_argv() -> list[str]: return ["mflux-upscale", "--prompt", "upscaled image", "--controlnet-image-path", "image.png"] +@pytest.mark.fast def test_scale_factor_auto(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test that 'auto' gets parsed as a ScaleFactor with value 1""" with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "auto", "--width", "auto"]): args = mflux_upscale_parser.parse_args() assert isinstance(args.height, ScaleFactor) @@ -58,8 +57,8 @@ def test_scale_factor_auto(mflux_upscale_parser, mflux_upscale_minimal_argv): assert args.width.value == 1 +@pytest.mark.fast def test_scale_factor_multiplier_format(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test scale factor formats like '1x', '2x', '3.5x'""" # Test integer scale factor with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "2x", "--width", "3x"]): args = mflux_upscale_parser.parse_args() @@ -85,8 +84,8 @@ def test_scale_factor_multiplier_format(mflux_upscale_parser, mflux_upscale_mini assert args.width.value == 0.5 +@pytest.mark.fast def test_plain_integer_dimensions(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test plain integer values for dimensions""" with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "1024", "--width", "768"]): args = mflux_upscale_parser.parse_args() assert isinstance(args.height, int) @@ -95,8 +94,8 @@ def test_plain_integer_dimensions(mflux_upscale_parser, mflux_upscale_minimal_ar assert args.width == 768 +@pytest.mark.fast def test_mixed_dimension_types(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test mixing scale factors and integers""" with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "2x", "--width", "1024"]): args = mflux_upscale_parser.parse_args() assert isinstance(args.height, ScaleFactor) @@ -112,8 +111,8 @@ def test_mixed_dimension_types(mflux_upscale_parser, mflux_upscale_minimal_argv) assert args.width.value == 1.5 +@pytest.mark.fast def test_default_dimensions(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test default values are 'auto' for upscale parser""" with patch("sys.argv", mflux_upscale_minimal_argv): args = mflux_upscale_parser.parse_args() # Default "auto" gets parsed into ScaleFactor(value=1) @@ -123,8 +122,8 @@ def test_default_dimensions(mflux_upscale_parser, mflux_upscale_minimal_argv): assert args.width.value == 1 +@pytest.mark.fast def test_invalid_scale_factor_format(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test invalid scale factor formats raise errors""" # Invalid format without 'x' with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "2.5"]): with pytest.raises(SystemExit): @@ -146,8 +145,8 @@ def test_invalid_scale_factor_format(mflux_upscale_parser, mflux_upscale_minimal mflux_upscale_parser.parse_args() +@pytest.mark.fast def test_case_insensitive_scale_factor(mflux_upscale_parser, mflux_upscale_minimal_argv): - """Test that scale factors are case insensitive""" with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "2X", "--width", "1.5X"]): args = mflux_upscale_parser.parse_args() assert isinstance(args.height, ScaleFactor) @@ -156,8 +155,8 @@ def test_case_insensitive_scale_factor(mflux_upscale_parser, mflux_upscale_minim assert args.width.value == 1.5 +@pytest.mark.fast def test_upscale_with_all_arguments(mflux_upscale_parser): - """Test upscale parser with all arguments""" full_argv = [ "mflux-upscale", "--prompt", diff --git a/tests/callbacks/test_battery_saver.py b/tests/callbacks/test_battery_saver.py index 479345c..6a97bb7 100644 --- a/tests/callbacks/test_battery_saver.py +++ b/tests/callbacks/test_battery_saver.py @@ -2,31 +2,32 @@ from unittest.mock import MagicMock, patch import pytest -from mflux.callbacks.instances.battery_saver import BatterySaver, get_battery_percentage +from mflux.callbacks.instances.battery_saver import BatterySaver from mflux.utils.exceptions import StopImageGenerationException +@pytest.mark.fast def test_get_battery_percentage_while_charging(): - """Test that the function returns None when the output doesn't match the expected pattern.""" - with patch("subprocess.run") as mock_run: + with ( + patch.object(BatterySaver, "_is_machine_battery_powered", return_value=True), + patch("subprocess.run") as mock_run, + ): # Set up mock to return an output that doesn't match the expected pattern mock_result = MagicMock() mock_result.stdout = "Now drawing from 'AC Power'" mock_run.return_value = mock_result - # Call the function - percentage = get_battery_percentage() + # Call the method + battery_saver = BatterySaver() + percentage = battery_saver._get_battery_percentage() # Assert the function returns None when no match is found assert percentage is None +@pytest.mark.fast def test_battery_saver_below_limit(): - """Test that BatterySaver raises an exception when the battery is below the limit.""" - with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get: - # Configure mock to return a battery percentage below the limit - mock_get.return_value = 5 - + with patch.object(BatterySaver, "_get_battery_percentage", return_value=5): # Create a BatterySaver instance with a limit of 10% battery_saver = BatterySaver(battery_percentage_stop_limit=10) @@ -39,12 +40,9 @@ def test_battery_saver_below_limit(): assert "5%" in str(excinfo.value) +@pytest.mark.fast def test_battery_saver_above_limit(): - """Test that BatterySaver does not raise an exception when the battery is above the limit.""" - with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get: - # Configure mock to return a battery percentage above the limit - mock_get.return_value = 20 - + with patch.object(BatterySaver, "_get_battery_percentage", return_value=20): # Create a BatterySaver instance with a limit of 10% battery_saver = BatterySaver(battery_percentage_stop_limit=10) @@ -52,12 +50,9 @@ def test_battery_saver_above_limit(): battery_saver.call_before_loop() +@pytest.mark.fast def test_battery_saver_none_percentage(): - """Test that BatterySaver does not raise an exception when percentage is None.""" - with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get: - # Configure mock to return None (e.g., on non-battery systems) - mock_get.return_value = None - + with patch.object(BatterySaver, "_get_battery_percentage", return_value=None): # Create a BatterySaver instance battery_saver = BatterySaver() @@ -65,12 +60,9 @@ def test_battery_saver_none_percentage(): battery_saver.call_before_loop() +@pytest.mark.fast def test_battery_saver_equal_to_limit(): - """Test that BatterySaver raises an exception when the battery equals the limit.""" - with patch("mflux.callbacks.instances.battery_saver.get_battery_percentage") as mock_get: - # Configure mock to return a battery percentage equal to the limit - mock_get.return_value = 10 - + with patch.object(BatterySaver, "_get_battery_percentage", return_value=10): # Create a BatterySaver instance with a limit of 10% battery_saver = BatterySaver(battery_percentage_stop_limit=10) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..9a8a1ce --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,3 @@ +def pytest_configure(config): + config.addinivalue_line("markers", "fast: marks tests as fast (no image generation)") + config.addinivalue_line("markers", "slow: marks tests as slow (generates images)") diff --git a/tests/depth/test_depth_pro.py b/tests/depth/test_depth_pro.py index 0ac10c1..55754a0 100644 --- a/tests/depth/test_depth_pro.py +++ b/tests/depth/test_depth_pro.py @@ -1,11 +1,14 @@ import os from pathlib import Path -from mflux.models.depth_pro.depth_pro import DepthPro +import pytest + +from mflux.models.depth_pro.model.depth_pro import DepthPro from mflux.utils.image_compare import ImageCompare class TestDepthPro: + @pytest.mark.slow def test_depth_pro_generation(self): # Resolve paths resource_dir = Path(__file__).parent.parent / "resources" @@ -32,5 +35,5 @@ class TestDepthPro: finally: # Cleanup - if os.path.exists(output_image_path): + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: os.remove(output_image_path) diff --git a/tests/dreambooth/test_dreambooth_dataset_iterator.py b/tests/dreambooth/test_dreambooth_dataset_iterator.py index db286e9..42819a5 100644 --- a/tests/dreambooth/test_dreambooth_dataset_iterator.py +++ b/tests/dreambooth/test_dreambooth_dataset_iterator.py @@ -34,8 +34,8 @@ def batch_size(request): return request.param +@pytest.mark.fast def test_batch_size_consistency(dataset, batch_size): - """Test that batches are of the specified size except for the last one.""" iterator = Iterator(dataset, batch_size) batch_sizes = [] @@ -55,8 +55,8 @@ def test_batch_size_consistency(dataset, batch_size): assert batch_sizes[-1] == dataset.size() % batch_size or batch_size +@pytest.mark.fast def test_complete_coverage(dataset, batch_size): - """Test that all examples are seen exactly once before reset.""" iterator = Iterator(dataset, batch_size) seen_examples: Set[int] = set() @@ -70,8 +70,8 @@ def test_complete_coverage(dataset, batch_size): assert seen_examples == expected_examples +@pytest.mark.fast def test_state_restoration(dataset, batch_size): - """Test that iterator can be saved and restored to the same state.""" iterator1 = Iterator(dataset, batch_size) first_batches = [] @@ -91,8 +91,8 @@ def test_state_restoration(dataset, batch_size): assert ids1 == ids2 +@pytest.mark.fast def test_state_restoration_across_epochs(dataset, batch_size): - """Test that state restoration works when crossing epoch boundaries.""" iterator1 = Iterator(dataset, batch_size) num_iterations = (dataset.size() + batch_size - 1) // batch_size + 1 @@ -109,8 +109,8 @@ def test_state_restoration_across_epochs(dataset, batch_size): assert ids1 == ids2 +@pytest.mark.fast def test_randomization(dataset, batch_size): - """Test that different seeds produce different sequences.""" iterator1 = Iterator(dataset, batch_size=2, seed=1) iterator2 = Iterator(dataset, batch_size=2, seed=2) @@ -123,8 +123,8 @@ def test_randomization(dataset, batch_size): assert ids1 != ids2 +@pytest.mark.fast def test_epoch_transition(dataset, batch_size): - """Test that epoch transition creates a new permutation.""" iterator = Iterator(dataset, batch_size) num_batches_in_epoch = (dataset.size() + batch_size - 1) // batch_size @@ -140,8 +140,8 @@ def test_epoch_transition(dataset, batch_size): assert len(second_epoch_ids) == min(batch_size, dataset.size()) +@pytest.mark.fast def test_state_consistency(dataset, batch_size): - """Test that state remains consistent across iterations within the same epoch.""" iterator = Iterator(dataset, batch_size) initial_state = iterator.to_dict() @@ -160,8 +160,8 @@ def test_state_consistency(dataset, batch_size): assert restored_state["current_permutation"] == new_state["current_permutation"] +@pytest.mark.fast def test_fixed_num_epochs(dataset, batch_size): - """Test that iterator stops after specified number of epochs.""" num_epochs = 100 iterator = Iterator(dataset, batch_size, num_epochs=num_epochs) @@ -188,8 +188,8 @@ def test_fixed_num_epochs(dataset, batch_size): assert total_examples_seen == dataset.size() * num_epochs +@pytest.mark.fast def test_iteration_counting(dataset, batch_size): - """Test that iteration counting works correctly.""" iterator = Iterator(dataset, batch_size) num_iterations = 5 @@ -199,8 +199,8 @@ def test_iteration_counting(dataset, batch_size): assert iterator.num_iterations == num_iterations +@pytest.mark.fast def test_iteration_counting_across_epochs(dataset, batch_size): - """Test that iteration counting works correctly across epoch boundaries.""" iterator = Iterator(dataset, batch_size) batches_per_epoch = (dataset.size() + batch_size - 1) // batch_size @@ -212,8 +212,8 @@ def test_iteration_counting_across_epochs(dataset, batch_size): assert iterator.num_iterations == total_iterations +@pytest.mark.fast def test_random_seed_consistency(dataset, batch_size): - """Test that the same seed produces the same sequence of batches.""" seed = 42 iterator1 = Iterator(dataset, batch_size, seed=seed) iterator2 = Iterator(dataset, batch_size, seed=seed) @@ -227,8 +227,8 @@ def test_random_seed_consistency(dataset, batch_size): assert ids1 == ids2 +@pytest.mark.fast def test_state_serialization(dataset, batch_size): - """Test that the iterator state can be serialized and deserialized.""" iterator = Iterator(dataset, batch_size) next(iterator) @@ -239,8 +239,8 @@ def test_state_serialization(dataset, batch_size): assert len(batch.examples) > 0 +@pytest.mark.fast def test_state_restoration_after_exhaustion(dataset, batch_size): - """Test that state restoration works correctly after iterator exhaustion.""" iterator = Iterator(dataset, batch_size, num_epochs=10) try: while True: diff --git a/tests/dreambooth/test_resume_training.py b/tests/dreambooth/test_resume_training.py index d2e92f5..c87f4d2 100644 --- a/tests/dreambooth/test_resume_training.py +++ b/tests/dreambooth/test_resume_training.py @@ -3,9 +3,9 @@ import shutil import mlx.core as mx import numpy as np +import pytest -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.dreambooth.dreambooth import DreamBooth from mflux.models.flux.variants.dreambooth.dreambooth_initializer import DreamBoothInitializer from mflux.models.flux.variants.dreambooth.state.zip_util import ZipUtil @@ -19,23 +19,24 @@ LORA_FILE = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint/0000005_adapte class TestResumeTraining: + @pytest.mark.slow def test_resume_training(self): # Clean up any existing temporary directories from previous test runs TestResumeTraining.delete_folder_if_exists("tests/dreambooth/tmp") try: # Given: A small training run from scratch for 5 steps (as described in the config)... - fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + fluxA, config, training_spec, training_state = DreamBoothInitializer.initialize( config_path="tests/dreambooth/config/train.json", checkpoint_path=None, ) DreamBooth.train( flux=fluxA, - runtime_config=runtime_config, + config=config, training_spec=training_spec, training_state=training_state, ) - del fluxA, runtime_config, training_spec, training_state + del fluxA, config, training_spec, training_state # ...where we can inspect the training state after 5 runs... adapter_after_5_steps = ZipUtil.unzip( zip_path=CHECKPOINT_5, @@ -47,17 +48,17 @@ class TestResumeTraining: TestResumeTraining.delete_file(CHECKPOINT_5) # When: Resuming the training from step 3... - fluxB, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + fluxB, config, training_spec, training_state = DreamBoothInitializer.initialize( config_path=None, checkpoint_path=CHECKPOINT_3, ) DreamBooth.train( flux=fluxB, - runtime_config=runtime_config, + config=config, training_spec=training_spec, training_state=training_state, ) - del fluxB, runtime_config, training_spec, training_state + del fluxB, config, training_spec, training_state # ...where we can inspect the training state after 2 additional runs... adapter_after_5_steps_resumed = ZipUtil.unzip( zip_path=CHECKPOINT_5, @@ -86,11 +87,9 @@ class TestResumeTraining: image = flux_with_resumed_lora.generate_image( seed=42, prompt="test", - config=Config( - num_inference_steps=1, - height=128, - width=128, - ), + num_inference_steps=1, + height=128, + width=128, ) # Basic sanity check that we got a valid image diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py index fba2584..2141348 100644 --- a/tests/dreambooth/test_train_and_load_weights.py +++ b/tests/dreambooth/test_train_and_load_weights.py @@ -2,9 +2,9 @@ import os import shutil import numpy as np +import pytest -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.dreambooth.dreambooth import DreamBooth from mflux.models.flux.variants.dreambooth.dreambooth_initializer import DreamBoothInitializer from mflux.models.flux.variants.dreambooth.state.zip_util import ZipUtil @@ -16,19 +16,20 @@ LORA_FILE = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint/0000005_adapte class TestTrainAndLoadWeights: + @pytest.mark.slow def test_train_and_load_weights(self): # Clean up any existing temporary directories from previous test runs TestTrainAndLoadWeights.delete_folder_if_exists("tests/dreambooth/tmp") try: # Given: A small training run from scratch for 5 steps (as described in the config)... - fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize( + fluxA, config, training_spec, training_state = DreamBoothInitializer.initialize( config_path="tests/dreambooth/config/train.json", checkpoint_path=None, ) DreamBooth.train( flux=fluxA, - runtime_config=runtime_config, + config=config, training_spec=training_spec, training_state=training_state, ) @@ -36,13 +37,11 @@ class TestTrainAndLoadWeights: image1 = fluxA.generate_image( seed=42, prompt="test", - config=Config( - num_inference_steps=20, - height=128, - width=128, - ), + num_inference_steps=20, + height=128, + width=128, ) - del fluxA, runtime_config, training_spec, training_state + del fluxA, config, training_spec, training_state # unzip so that LoRA adapter can be read later... ZipUtil.extract_all(zip_path=CHECKPOINT, output_dir=OUTPUT_DIR) @@ -58,11 +57,9 @@ class TestTrainAndLoadWeights: image2 = fluxB.generate_image( seed=42, prompt="test", - config=Config( - num_inference_steps=20, - height=128, - width=128, - ), + num_inference_steps=20, + height=128, + width=128, ) # Then: We want to confirm that the images *exactly* match diff --git a/tests/image_generation/helpers/image_generation_concept_test_helper.py b/tests/image_generation/helpers/image_generation_concept_test_helper.py index 28657d1..87c8432 100644 --- a/tests/image_generation/helpers/image_generation_concept_test_helper.py +++ b/tests/image_generation/helpers/image_generation_concept_test_helper.py @@ -1,8 +1,7 @@ import os from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.concept_attention.flux_concept import Flux1Concept from mflux.models.flux.variants.concept_attention.flux_concept_from_image import Flux1ConceptFromImage from mflux.utils.image_compare import ImageCompare @@ -43,13 +42,11 @@ class ImageGenerationConceptTestHelper: seed=seed, prompt=prompt, concept=concept, + num_inference_steps=steps, + height=height or 1024, + width=width or 1024, heatmap_layer_indices=heatmap_layer_indices, heatmap_timesteps=heatmap_timesteps, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - ), ) # Save only the heatmap (we don't need the original image for testing) image.save_concept_heatmap(path=output_heatmap_path, overwrite=True) @@ -104,13 +101,11 @@ class ImageGenerationConceptTestHelper: prompt=prompt, concept=concept, image_path=str(input_image_path), + num_inference_steps=steps, + height=height or 1024, + width=width or 1024, heatmap_layer_indices=heatmap_layer_indices, heatmap_timesteps=heatmap_timesteps, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - ), ) # Save only the heatmap (we don't need the original image for testing) image.save_concept_heatmap(path=output_heatmap_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index 5967772..3ff3189 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py @@ -1,7 +1,6 @@ import os -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet from mflux.utils.image_compare import ImageCompare from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper @@ -43,12 +42,10 @@ class ImageGeneratorControlnetTestHelper: seed=seed, prompt=prompt, controlnet_image_path=controlnet_image_path, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - controlnet_strength=controlnet_strength, - ), + num_inference_steps=steps, + height=height, + width=width, + controlnet_strength=controlnet_strength, ) image.save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_depth_test_helper.py b/tests/image_generation/helpers/image_generation_depth_test_helper.py index 546bb9e..261788b 100644 --- a/tests/image_generation/helpers/image_generation_depth_test_helper.py +++ b/tests/image_generation/helpers/image_generation_depth_test_helper.py @@ -1,8 +1,7 @@ import os from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.depth.flux_depth import Flux1Depth from mflux.utils.image_compare import ImageCompare @@ -33,13 +32,11 @@ class ImageGeneratorDepthTestHelper: image = flux.generate_image( seed=seed, prompt=prompt, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - image_path=image_path, - depth_image_path=depth_image_path, - ), + num_inference_steps=steps, + height=height or 1024, + width=width or 1024, + image_path=image_path, + depth_image_path=depth_image_path, ) image.save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_edit_test_helper.py b/tests/image_generation/helpers/image_generation_edit_test_helper.py index fb1b9b6..7643b81 100644 --- a/tests/image_generation/helpers/image_generation_edit_test_helper.py +++ b/tests/image_generation/helpers/image_generation_edit_test_helper.py @@ -2,10 +2,9 @@ import os from pathlib import Path from typing import Any, Type -from mflux.callbacks.callback_registry import CallbackRegistry from mflux.callbacks.instances.stepwise_handler import StepwiseHandler -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig +from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator from mflux.utils.image_compare import ImageCompare @@ -26,8 +25,8 @@ class ImageGeneratorEditTestHelper: negative_prompt: str | None = None, quantize: int = 8, image_paths: list[str] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, mismatch_threshold: float | None = None, ): # resolve paths @@ -48,30 +47,23 @@ class ImageGeneratorEditTestHelper: # given model_kwargs = { "quantize": quantize, + "lora_paths": lora_paths, + "lora_scales": lora_scales, } - # Add HuggingFace LoRA parameters if provided - if lora_names is not None: - model_kwargs["lora_names"] = lora_names - if lora_repo_id is not None: - model_kwargs["lora_repo_id"] = lora_repo_id - model = model_class(**model_kwargs) # when - config_kwargs = { + generate_kwargs = { + "seed": seed, + "prompt": prompt, "num_inference_steps": steps, "height": height, "width": width, "guidance": guidance, "image_path": image_path, "scheduler": "flow_match_euler_discrete", # Match debug script - } - generate_kwargs = { - "seed": seed, - "prompt": prompt, "negative_prompt": negative_prompt, - "config": Config(**config_kwargs), } # Qwen Edit uses image_paths instead of image_path in config @@ -86,10 +78,8 @@ class ImageGeneratorEditTestHelper: import tempfile temp_dir = tempfile.mkdtemp() - handler = StepwiseHandler(model=model, output_dir=temp_dir) - CallbackRegistry.register_before_loop(handler) - CallbackRegistry.register_in_loop(handler) - CallbackRegistry.register_interrupt(handler) + handler = StepwiseHandler(model=model, output_dir=temp_dir, latent_creator=QwenLatentCreator) + model.callbacks.register(handler) image = model.generate_image(**generate_kwargs) image.save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_fibo_test_helper.py b/tests/image_generation/helpers/image_generation_fibo_test_helper.py index 7369b0d..24ce85d 100644 --- a/tests/image_generation/helpers/image_generation_fibo_test_helper.py +++ b/tests/image_generation/helpers/image_generation_fibo_test_helper.py @@ -2,8 +2,6 @@ import os from pathlib import Path from typing import Optional -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig from mflux.models.fibo.variants.txt2img.fibo import FIBO from mflux.utils.image_compare import ImageCompare @@ -30,25 +28,20 @@ class ImageGeneratorFiboTestHelper: try: # Step 1: Create FIBO model model = FIBO( - model_config=ModelConfig.fibo(), quantize=quantize, - local_path=None, + model_path=None, ) # Step 2: Generate image from prompt image = model.generate_image( seed=seed, prompt=prompt, + num_inference_steps=steps, + height=height, + width=width, + guidance=guidance, + scheduler="flow_match_euler_discrete", negative_prompt=negative_prompt, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - guidance=guidance, - image_path=None, - image_strength=None, - scheduler="flow_match_euler_discrete", - ), ) # Step 3: Save output image diff --git a/tests/image_generation/helpers/image_generation_fill_test_helper.py b/tests/image_generation/helpers/image_generation_fill_test_helper.py index db4942b..60dfe6f 100644 --- a/tests/image_generation/helpers/image_generation_fill_test_helper.py +++ b/tests/image_generation/helpers/image_generation_fill_test_helper.py @@ -1,9 +1,8 @@ import os -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.cli.defaults import defaults as ui_defaults +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.fill.flux_fill import Flux1Fill -from mflux.ui import defaults as ui_defaults from mflux.utils.image_compare import ImageCompare from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper @@ -43,14 +42,12 @@ class ImageGeneratorFillTestHelper: image = flux.generate_image( seed=seed, prompt=prompt, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - image_path=image_path, - masked_image_path=masked_image_path, - guidance=ui_defaults.DEFAULT_DEV_FILL_GUIDANCE, - ), + image_path=image_path, + masked_image_path=masked_image_path, + num_inference_steps=steps, + height=height, + width=width, + guidance=ui_defaults.DEFAULT_DEV_FILL_GUIDANCE, ) image.save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py b/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py index 4931f85..2c67cc2 100644 --- a/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py +++ b/tests/image_generation/helpers/image_generation_ic_edit_test_helper.py @@ -1,10 +1,9 @@ import os from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill -from mflux.models.flux.variants.in_context.utils.in_context_loras import prepare_ic_edit_loras +from mflux.models.flux.variants.in_context.utils.in_context_loras import IC_EDIT_LORA_SCALE, get_ic_edit_lora_path from mflux.utils.image_compare import ImageCompare @@ -28,13 +27,20 @@ class ImageGeneratorICEditTestHelper: reference_image = ImageGeneratorICEditTestHelper.resolve_path(reference_image) lora_paths = [str(ImageGeneratorICEditTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None + # Build lora_paths: IC-Edit LoRA (required) + user-provided LoRAs + all_lora_paths = [get_ic_edit_lora_path()] + all_lora_scales = [IC_EDIT_LORA_SCALE] + if lora_paths: + all_lora_paths.extend(lora_paths) + all_lora_scales.extend(lora_scales or [1.0] * len(lora_paths)) + try: # given flux = Flux1InContextFill( model_config=ModelConfig.dev_fill(), quantize=8, - lora_paths=prepare_ic_edit_loras(lora_paths), - lora_scales=lora_scales, + lora_paths=all_lora_paths, + lora_scales=all_lora_scales, ) # when @@ -42,12 +48,10 @@ class ImageGeneratorICEditTestHelper: seed=seed, prompt=prompt, left_image_path=str(reference_image), + num_inference_steps=steps, + height=height or 1024, + width=width or 1024, right_image_path=None, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - ), ) # Save the result diff --git a/tests/image_generation/helpers/image_generation_in_context_test_helper.py b/tests/image_generation/helpers/image_generation_in_context_test_helper.py index 5129ced..dae8978 100644 --- a/tests/image_generation/helpers/image_generation_in_context_test_helper.py +++ b/tests/image_generation/helpers/image_generation_in_context_test_helper.py @@ -1,10 +1,9 @@ import os from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.in_context.flux_in_context_dev import Flux1InContextDev -from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename +from mflux.models.flux.variants.in_context.utils.in_context_loras import get_lora_path from mflux.utils.image_compare import ImageCompare @@ -32,26 +31,32 @@ class ImageGeneratorInContextTestHelper: [str(ImageGeneratorInContextTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None ) + # Build lora_paths: style LoRA (if specified) + user-provided LoRAs + all_lora_paths = [] + all_lora_scales = [] + if lora_style: + all_lora_paths.append(get_lora_path(lora_style)) + all_lora_scales.append(1.0) + if lora_paths: + all_lora_paths.extend(lora_paths) + all_lora_scales.extend(lora_scales or [1.0] * len(lora_paths)) + try: # given flux = Flux1InContextDev( model_config=model_config, quantize=8, - lora_names=[get_lora_filename(lora_style)] if lora_style else None, - lora_repo_id=LORA_REPO_ID if lora_style else None, - lora_paths=lora_paths, - lora_scales=lora_scales, + lora_paths=all_lora_paths or None, + lora_scales=all_lora_scales or None, ) # when image = flux.generate_image( seed=seed, prompt=prompt, - config=Config( - num_inference_steps=steps, - image_path=image_path, - height=height, - width=width, - ), + num_inference_steps=steps, + image_path=image_path, + height=height or 1024, + width=width or 1024, ) # Save only the right half of the image (the generated part) image.get_right_half().save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_kontext_test_helper.py b/tests/image_generation/helpers/image_generation_kontext_test_helper.py index e2f3475..ecc0111 100644 --- a/tests/image_generation/helpers/image_generation_kontext_test_helper.py +++ b/tests/image_generation/helpers/image_generation_kontext_test_helper.py @@ -1,8 +1,7 @@ import os from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext from mflux.utils.image_compare import ImageCompare @@ -36,13 +35,11 @@ class ImageGeneratorKontextTestHelper: image = flux.generate_image( seed=seed, prompt=prompt, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - guidance=guidance, - image_path=kontext_image_path, - ), + num_inference_steps=steps, + height=height, + width=width, + guidance=guidance, + image_path=kontext_image_path, ) image.save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_redux_test_helper.py b/tests/image_generation/helpers/image_generation_redux_test_helper.py index 5bd789c..3458e61 100644 --- a/tests/image_generation/helpers/image_generation_redux_test_helper.py +++ b/tests/image_generation/helpers/image_generation_redux_test_helper.py @@ -1,8 +1,7 @@ import os from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.redux.flux_redux import Flux1Redux from mflux.utils.image_compare import ImageCompare @@ -12,13 +11,13 @@ class ImageGeneratorReduxTestHelper: def assert_matches_reference_image( reference_image_path: str, output_image_path: str, - model_config: ModelConfig, steps: int, seed: int, height: int, width: int, prompt: str, redux_image_path: str, + model_config: ModelConfig = ModelConfig.dev(), ): # resolve paths reference_image_path = ImageGeneratorReduxTestHelper.resolve_path(reference_image_path) @@ -36,12 +35,10 @@ class ImageGeneratorReduxTestHelper: image = flux.generate_image( seed=seed, prompt=prompt, - config=Config( - num_inference_steps=steps, - height=height, - width=width, - redux_image_paths=[redux_image_path], - ), + redux_image_paths=[redux_image_path], + num_inference_steps=steps, + height=height, + width=width, ) image.save(path=output_image_path, overwrite=True) diff --git a/tests/image_generation/helpers/image_generation_test_helper.py b/tests/image_generation/helpers/image_generation_test_helper.py index e0820df..bb9e695 100644 --- a/tests/image_generation/helpers/image_generation_test_helper.py +++ b/tests/image_generation/helpers/image_generation_test_helper.py @@ -2,8 +2,7 @@ import os from pathlib import Path from typing import Type, Union -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.txt2img.flux import Flux1 from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage from mflux.utils.image_compare import ImageCompare @@ -26,8 +25,6 @@ class ImageGeneratorTestHelper: image_strength: float | None = None, lora_paths: list[str] | None = None, lora_scales: list[float] | None = None, - lora_names: list[str] | None = None, - lora_repo_id: str | None = None, negative_prompt: str | None = None, guidance: float | None = None, mismatch_threshold: float | None = None, @@ -46,31 +43,22 @@ class ImageGeneratorTestHelper: "lora_scales": lora_scales, } - # Add HuggingFace LoRA parameters if provided - if lora_names is not None: - model_kwargs["lora_names"] = lora_names - if lora_repo_id is not None: - model_kwargs["lora_repo_id"] = lora_repo_id - model = model_class(**model_kwargs) - config_kwargs = { - "num_inference_steps": steps, - "image_path": ImageGeneratorTestHelper.resolve_path(image_path), - "image_strength": image_strength, - "height": height, - "width": width, - } - - # Add guidance if provided - if guidance is not None: - config_kwargs["guidance"] = guidance generate_kwargs = { "seed": seed, "prompt": prompt, - "config": Config(**config_kwargs), + "num_inference_steps": steps, + "image_path": ImageGeneratorTestHelper.resolve_path(image_path), + "image_strength": image_strength, + "height": height or 1024, + "width": width or 1024, } + # Add guidance if provided + if guidance is not None: + generate_kwargs["guidance"] = guidance + # Add negative_prompt for Qwen models if model_class == QwenImage and negative_prompt is not None: generate_kwargs["negative_prompt"] = negative_prompt diff --git a/tests/image_generation/helpers/image_generation_z_image_test_helper.py b/tests/image_generation/helpers/image_generation_z_image_test_helper.py new file mode 100644 index 0000000..0b307ce --- /dev/null +++ b/tests/image_generation/helpers/image_generation_z_image_test_helper.py @@ -0,0 +1,80 @@ +import os +import shutil +from pathlib import Path + +from mflux.cli.defaults.defaults import MFLUX_LORA_CACHE_DIR +from mflux.models.z_image import ZImageTurbo +from mflux.utils.image_compare import ImageCompare + + +class ImageGeneratorZImageTestHelper: + @staticmethod + def assert_matches_reference_image( + reference_image_path: str, + output_image_path: str, + prompt: str, + steps: int, + seed: int, + height: int, + width: int, + quantize: int | None = None, + lora_paths: list[str] | None = None, + lora_scales: list[float] | None = None, + mismatch_threshold: float | None = None, + clear_lora_cache_pattern: str | None = None, + ): + reference_image_path = ImageGeneratorZImageTestHelper.resolve_path(reference_image_path) + output_image_path = ImageGeneratorZImageTestHelper.resolve_path(output_image_path) + + # Clear cached LoRA to test download functionality + if clear_lora_cache_pattern: + ImageGeneratorZImageTestHelper.clear_cached_lora(clear_lora_cache_pattern) + + try: + model = ZImageTurbo( + quantize=quantize, + lora_paths=lora_paths, + lora_scales=lora_scales, + ) + + image = model.generate_image( + seed=seed, + prompt=prompt, + num_inference_steps=steps, + height=height, + width=width, + ) + + image.save(output_image_path) + + ImageCompare.check_images_close_enough( + output_image_path, + reference_image_path, + "Generated image doesn't match reference image.", + mismatch_threshold=mismatch_threshold, + ) + finally: + if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ: + os.remove(output_image_path) + + @staticmethod + def clear_cached_lora(pattern: str) -> None: + cache_dir = MFLUX_LORA_CACHE_DIR + if not cache_dir.exists(): + return + + # Look for directories matching the pattern (HuggingFace style: models--repo--name) + for item in cache_dir.iterdir(): + if pattern.lower() in item.name.lower(): + if item.is_dir(): + print(f"๐Ÿ—‘๏ธ Clearing cached LoRA directory: {item}") + shutil.rmtree(item) + elif item.is_file() or item.is_symlink(): + print(f"๐Ÿ—‘๏ธ Clearing cached LoRA file: {item}") + item.unlink() + + @staticmethod + def resolve_path(path) -> Path | None: + if path is None: + return None + return Path(__file__).parent.parent.parent / "resources" / path diff --git a/tests/image_generation/test_generate_concept.py b/tests/image_generation/test_generate_concept.py index e7d668e..23a6e09 100644 --- a/tests/image_generation/test_generate_concept.py +++ b/tests/image_generation/test_generate_concept.py @@ -1,8 +1,11 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from tests.image_generation.helpers.image_generation_concept_test_helper import ImageGenerationConceptTestHelper class TestImageGeneratorConcept: + @pytest.mark.slow def test_concept_attention_generation(self): ImageGenerationConceptTestHelper.assert_matches_reference_image_concept( reference_heatmap_path="reference_concept_schnell_heatmap.png", @@ -18,6 +21,7 @@ class TestImageGeneratorConcept: heatmap_timesteps=[0, 1, 2, 3], ) + @pytest.mark.slow def test_concept_attention_from_image(self): ImageGenerationConceptTestHelper.assert_matches_reference_image_concept_from_image( reference_heatmap_path="reference_concept_from_image_schnell_heatmap.png", diff --git a/tests/image_generation/test_generate_image.py b/tests/image_generation/test_generate_image.py index 80fe106..8329ae1 100644 --- a/tests/image_generation/test_generate_image.py +++ b/tests/image_generation/test_generate_image.py @@ -1,9 +1,12 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.txt2img.flux import Flux1 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper class TestImageGenerator: + @pytest.mark.slow def test_image_generation_schnell(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_schnell.png", @@ -17,6 +20,7 @@ class TestImageGenerator: prompt="Luxury food photograph", ) + @pytest.mark.slow def test_image_generation_dev(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev.png", @@ -30,6 +34,7 @@ class TestImageGenerator: prompt="Luxury food photograph", ) + @pytest.mark.slow def test_image_generation_dev_lora(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_lora.png", @@ -45,6 +50,7 @@ class TestImageGenerator: lora_scales=[1.0], ) + @pytest.mark.slow def test_image_generation_dev_multiple_loras(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_lora_multiple.png", @@ -60,6 +66,7 @@ class TestImageGenerator: lora_scales=[0.4, 0.6], ) + @pytest.mark.slow def test_image_generation_dev_image_to_image(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_image_to_image.png", diff --git a/tests/image_generation/test_generate_image_controlnet.py b/tests/image_generation/test_generate_image_controlnet.py index 6732205..ebfe83a 100644 --- a/tests/image_generation/test_generate_image_controlnet.py +++ b/tests/image_generation/test_generate_image_controlnet.py @@ -1,8 +1,11 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from tests.image_generation.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper class TestImageGeneratorControlnet: + @pytest.mark.slow def test_image_generation_schnell_controlnet(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_controlnet_schnell.png", @@ -17,6 +20,7 @@ class TestImageGeneratorControlnet: controlnet_strength=0.4, ) + @pytest.mark.slow def test_image_generation_dev_controlnet(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_controlnet_dev.png", @@ -31,6 +35,7 @@ class TestImageGeneratorControlnet: controlnet_strength=0.4, ) + @pytest.mark.slow def test_image_generation_dev_lora_controlnet(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_controlnet_dev_lora.png", @@ -47,6 +52,7 @@ class TestImageGeneratorControlnet: controlnet_strength=0.4, ) + @pytest.mark.slow def test_image_upscaling(self): ImageGeneratorControlnetTestHelper.assert_matches_reference_image( reference_image_path="reference_upscaled.png", diff --git a/tests/image_generation/test_generate_image_depth.py b/tests/image_generation/test_generate_image_depth.py index f8d4068..8730c77 100644 --- a/tests/image_generation/test_generate_image_depth.py +++ b/tests/image_generation/test_generate_image_depth.py @@ -1,8 +1,11 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from tests.image_generation.helpers.image_generation_depth_test_helper import ImageGeneratorDepthTestHelper class TestImageGeneratorDepth: + @pytest.mark.slow def test_image_generation_with_reference_image(self): ImageGeneratorDepthTestHelper.assert_matches_reference_image( reference_image_path="reference_depth_dev_from_image.png", diff --git a/tests/image_generation/test_generate_image_fibo.py b/tests/image_generation/test_generate_image_fibo.py index 3fe1683..1620606 100644 --- a/tests/image_generation/test_generate_image_fibo.py +++ b/tests/image_generation/test_generate_image_fibo.py @@ -1,3 +1,5 @@ +import pytest + from tests.image_generation.helpers.image_generation_fibo_test_helper import ImageGeneratorFiboTestHelper OWL_PROMPT = """ @@ -11,33 +13,33 @@ OWL_PROMPT = """ "relative_size": "large within frame", "shape_and_color": "Round head, large eyes, bulky body, predominantly brown and grey with silver accents.", "texture": "Extremely soft, fluffy, and detailed feathers, giving a plush toy-like appearance.", - "appearance_details": "The eyes are wide, dark, and reflective, conveying a sense of wonder and curiosity. The beak is small and light-colored, almost hidden by the feathers. Subtle silver highlights catch the moonlight on its feathers.", + "appearance_details": "The eyes are wide, round, and have a glossy, reflective quality, suggesting innocence and curiosity. The beak is small and light-colored, almost blending with the feathers.", "orientation": "upright, facing forward" } ], - "background_setting": "A dark, nocturnal forest setting with blurred trees and foliage, illuminated by a soft, cool moonlight. The background is out of focus, emphasizing the owl.", + "background_setting": "A dark, blurred forest at night, with hints of tree trunks and foliage visible in the background. The darkness emphasizes the owl as the central focus.", "lighting": { "conditions": "moonlight", - "direction": "backlit and side-lit from the left", - "shadows": "soft, diffused shadows on the right side of the owl and within the background foliage, indicating a single light source." + "direction": "top-left, casting subtle highlights", + "shadows": "soft, diffused shadows, contributing to the depth without harshness" }, "aesthetics": { "composition": "centered, portrait composition", - "color_scheme": "cool blues and silvers from the moonlight contrasting with warm browns and grays of the owl and forest.", - "mood_atmosphere": "mysterious, enchanting, whimsical, and serene.", + "color_scheme": "cool blues and grays of the night sky and warm browns and grays of the owl, with silver highlights", + "mood_atmosphere": "whimsical, serene, enchanting, and slightly mysterious", "aesthetic_score": "very high", "preference_score": "very high" }, "photographic_characteristics": { "depth_of_field": "shallow", - "focus": "sharp focus on the owl's face and eyes, with a soft blur in the background.", + "focus": "sharp focus on the owl's face and eyes", "camera_angle": "eye-level", "lens_focal_length": "portrait lens (e.g., 50mm-85mm)" }, "style_medium": "digital illustration", "text_render": [], - "context": "A whimsical character illustration, possibly for a children's book, animated film, or fantasy art collection.", - "artistic_style": "fantasy, illustrative, detailed" + "context": "A whimsical character illustration, suitable for children's books, animated features, or as a charming decorative art piece.", + "artistic_style": "hyperrealistic, cute, illustrative" } """ @@ -46,44 +48,45 @@ OWL_PROMPT_REFINED = """ "short_description": "A hyper-detailed, ultra-fluffy owl sitting in the trees at night, looking directly at the camera with wide, adorable, expressive eyes. Its feathers are soft and voluminous, catching the cool moonlight with subtle silver highlights. The owl's gaze is curious and full of charm, giving it a whimsical, storybook-like personality.", "objects": [ { - "description": "An adorable, fluffy owl with large, expressive eyes and soft, voluminous feathers. Its plumage is a mix of white and subtle silver highlights from the moonlight.", + "description": "An adorable, fluffy owl with large, expressive eyes and soft, voluminous feathers. Its plumage is white with subtle silver highlights from the moonlight.", "location": "center", "relationship": "The owl is the sole subject, perched comfortably within its environment.", "relative_size": "large within frame", "shape_and_color": "Round head, large eyes, bulky body, predominantly white with silver accents.", "texture": "Extremely soft, fluffy, and detailed feathers, giving a plush toy-like appearance.", - "appearance_details": "The eyes are wide, dark, and reflective, conveying a sense of wonder and curiosity. The beak is small and light-colored, almost hidden by the feathers. Subtle silver highlights catch the moonlight on its feathers.", + "appearance_details": "The eyes are wide, round, and have a glossy, reflective quality, suggesting innocence and curiosity. The beak is small and light-colored, almost blending with the feathers.", "orientation": "upright, facing forward" } ], - "background_setting": "A dark, nocturnal forest setting with blurred trees and foliage, illuminated by a soft, cool moonlight. The background is out of focus, emphasizing the owl.", + "background_setting": "A dark, blurred forest at night, with hints of tree trunks and foliage visible in the background. The darkness emphasizes the owl as the central focus.", "lighting": { "conditions": "moonlight", - "direction": "backlit and side-lit from the left", - "shadows": "soft, diffused shadows on the right side of the owl and within the background foliage, indicating a single light source." + "direction": "top-left, casting subtle highlights", + "shadows": "soft, diffused shadows, contributing to the depth without harshness" }, "aesthetics": { "composition": "centered, portrait composition", - "color_scheme": "cool blues and silvers from the moonlight contrasting with white of the owl and forest.", - "mood_atmosphere": "mysterious, enchanting, whimsical, and serene.", + "color_scheme": "cool blues and grays of the night sky and white of the owl, with silver accents", + "mood_atmosphere": "whimsical, serene, enchanting, and slightly mysterious", "aesthetic_score": "very high", "preference_score": "very high" }, "photographic_characteristics": { "depth_of_field": "shallow", - "focus": "sharp focus on the owl's face and eyes, with a soft blur in the background.", + "focus": "sharp focus on the owl's face and eyes", "camera_angle": "eye-level", "lens_focal_length": "portrait lens (e.g., 50mm-85mm)" }, "style_medium": "digital illustration", "text_render": [], - "context": "A whimsical character illustration, possibly for a children's book, animated film, or fantasy art collection.", - "artistic_style": "fantasy, illustrative, detailed" + "context": "A whimsical character illustration, suitable for children's books, animated features, or as a charming decorative art piece.", + "artistic_style": "hyperrealistic, cute, illustrative" } """ class TestImageGeneratorFibo: + @pytest.mark.slow def test_image_generation_fibo(self): ImageGeneratorFiboTestHelper.assert_matches_reference_image( reference_image_path="reference_fibo.png", @@ -91,12 +94,13 @@ class TestImageGeneratorFibo: prompt=OWL_PROMPT, # Assume this has been generated by the VLM, actual VLM tests are separate steps=20, seed=42, - height=176, - width=320, + height=352, + width=640, guidance=4.0, quantize=8, ) + @pytest.mark.slow def test_image_generation_fibo_refined_white_owl(self): ImageGeneratorFiboTestHelper.assert_matches_reference_image( reference_image_path="reference_fibo_white_owl.png", @@ -104,8 +108,8 @@ class TestImageGeneratorFibo: prompt=OWL_PROMPT_REFINED, # Assume this has been refined by the VLM, actual VLM tests are separate steps=20, seed=42, - height=176, - width=320, + height=352, + width=640, guidance=4.0, quantize=8, ) diff --git a/tests/image_generation/test_generate_image_fill.py b/tests/image_generation/test_generate_image_fill.py index d4fb718..83b78b9 100644 --- a/tests/image_generation/test_generate_image_fill.py +++ b/tests/image_generation/test_generate_image_fill.py @@ -1,8 +1,11 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from tests.image_generation.helpers.image_generation_fill_test_helper import ImageGeneratorFillTestHelper class TestImageGeneratorFill: + @pytest.mark.slow def test_fill(self): ImageGeneratorFillTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_image_to_image_fill.png", diff --git a/tests/image_generation/test_generate_image_in_context.py b/tests/image_generation/test_generate_image_in_context.py index 44912ef..5798f0e 100644 --- a/tests/image_generation/test_generate_image_in_context.py +++ b/tests/image_generation/test_generate_image_in_context.py @@ -1,9 +1,12 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config 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 class TestImageGeneratorInContext: + @pytest.mark.slow def test_in_context_lora_identity(self): ImageGeneratorInContextTestHelper.assert_matches_reference_image( reference_image_path="reference_ic_lora_identity.png", @@ -18,6 +21,7 @@ class TestImageGeneratorInContext: lora_style="identity", ) + @pytest.mark.slow def test_ic_edit_with_instruction(self): ImageGeneratorICEditTestHelper.assert_matches_reference_image( reference_image_path="reference_ic_edit.png", diff --git a/tests/image_generation/test_generate_image_kontext.py b/tests/image_generation/test_generate_image_kontext.py index 71556a1..b75227b 100644 --- a/tests/image_generation/test_generate_image_kontext.py +++ b/tests/image_generation/test_generate_image_kontext.py @@ -1,8 +1,11 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from tests.image_generation.helpers.image_generation_kontext_test_helper import ImageGeneratorKontextTestHelper class TestImageGeneratorKontext: + @pytest.mark.slow def test_image_generation_kontext(self): ImageGeneratorKontextTestHelper.assert_matches_reference_image( reference_image_path="reference_dev_kontext.png", diff --git a/tests/image_generation/test_generate_image_qwen_image.py b/tests/image_generation/test_generate_image_qwen_image.py index 628b00d..e9b2d71 100644 --- a/tests/image_generation/test_generate_image_qwen_image.py +++ b/tests/image_generation/test_generate_image_qwen_image.py @@ -1,9 +1,12 @@ -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper class TestImageGeneratorQwenImage: + @pytest.mark.slow def test_qwen_image_generation_text_to_image(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_qwen_txt2img.png", @@ -20,6 +23,7 @@ class TestImageGeneratorQwenImage: mismatch_threshold=0.35, # Qwen models produce visually similar images with minor pixel differences ) + @pytest.mark.slow def test_qwen_image_generation_image_to_image(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_qwen_img2img.png", @@ -38,6 +42,7 @@ class TestImageGeneratorQwenImage: mismatch_threshold=0.35, # Qwen models produce visually similar images with minor pixel differences ) + @pytest.mark.slow def test_qwen_image_generation_lora(self): ImageGeneratorTestHelper.assert_matches_reference_image( reference_image_path="reference_qwen_lora.png", @@ -52,7 +57,7 @@ class TestImageGeneratorQwenImage: width=768, prompt="Luxury food photograph", negative_prompt="ugly, blurry, low quality", - lora_repo_id="lightx2v/Qwen-Image-Lightning", - lora_names=["Qwen-Image-Lightning-4steps-V2.0.safetensors"], + lora_paths=["lightx2v/Qwen-Image-Lightning:Qwen-Image-Lightning-4steps-V2.0.safetensors"], + lora_scales=[1.0], mismatch_threshold=0.65, # LoRA tests have higher variance due to model updates ) diff --git a/tests/image_generation/test_generate_image_qwen_image_edit.py b/tests/image_generation/test_generate_image_qwen_image_edit.py index 7dd0663..327ae1a 100644 --- a/tests/image_generation/test_generate_image_qwen_image_edit.py +++ b/tests/image_generation/test_generate_image_qwen_image_edit.py @@ -1,11 +1,12 @@ import pytest -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.qwen.variants.edit import QwenImageEdit from tests.image_generation.helpers.image_generation_edit_test_helper import ImageGeneratorEditTestHelper class TestImageGeneratorQwenImageEdit: + @pytest.mark.slow def test_image_generation_qwen_edit(self): ImageGeneratorEditTestHelper.assert_matches_reference_image( reference_image_path="reference_qwen_edit.png", @@ -23,6 +24,7 @@ class TestImageGeneratorQwenImageEdit: mismatch_threshold=0.25, ) + @pytest.mark.slow @pytest.mark.high_memory_requirement def test_image_generation_qwen_edit_multiple_images(self): ImageGeneratorEditTestHelper.assert_matches_reference_image( diff --git a/tests/image_generation/test_generate_image_redux.py b/tests/image_generation/test_generate_image_redux.py index a4d9008..5ff3b85 100644 --- a/tests/image_generation/test_generate_image_redux.py +++ b/tests/image_generation/test_generate_image_redux.py @@ -1,13 +1,14 @@ -from mflux.config.model_config import ModelConfig +import pytest + from tests.image_generation.helpers.image_generation_redux_test_helper import ImageGeneratorReduxTestHelper class TestImageGeneratorRedux: + @pytest.mark.slow def test_image_generation_redux(self): ImageGeneratorReduxTestHelper.assert_matches_reference_image( reference_image_path="reference_redux_dev.png", output_image_path="output_redux_dev.png", - model_config=ModelConfig.dev_redux(), steps=15, seed=42, height=341, diff --git a/tests/image_generation/test_generate_image_z_image.py b/tests/image_generation/test_generate_image_z_image.py new file mode 100644 index 0000000..6edb80d --- /dev/null +++ b/tests/image_generation/test_generate_image_z_image.py @@ -0,0 +1,49 @@ +import pytest + +from tests.image_generation.helpers.image_generation_z_image_test_helper import ImageGeneratorZImageTestHelper + +ASTRONAUT_CAT_PROMPT = ( + "A fluffy orange tabby cat wearing tiny astronaut helmet, " + "floating in zero gravity inside a space station. " + "Earth visible through the window behind. " + "Photorealistic, cinematic lighting, 8k detail." +) + +FROG_LORA_PROMPT = ( + "t3chnic4lly vibrant 1960s close-up of a woman sitting under a tree in a blue skirt and white blouse, " + "she has blonde wavy short hair and a smile with green eyes lake scene by a garden with flowers in the foreground " + "1960s style film She's holding her hand out there is a small smooth frog in her palm, " + "she's making eye contact with the toad." +) + + +class TestImageGeneratorZImage: + @pytest.mark.slow + def test_image_generation_z_image_turbo(self): + ImageGeneratorZImageTestHelper.assert_matches_reference_image( + reference_image_path="reference_z_image_turbo.png", + output_image_path="output_z_image_turbo.png", + prompt=ASTRONAUT_CAT_PROMPT, + steps=9, + seed=42, + height=368, + width=640, + quantize=8, + ) + + @pytest.mark.slow + def test_image_generation_z_image_turbo_lora(self): + ImageGeneratorZImageTestHelper.assert_matches_reference_image( + reference_image_path="reference_z_image_turbo_lora.png", + output_image_path="output_z_image_turbo_lora.png", + prompt=FROG_LORA_PROMPT, + steps=9, + seed=42, + height=368, + width=640, + quantize=8, + lora_paths=["renderartist/Technically-Color-Z-Image-Turbo"], + lora_scales=[0.5], + clear_lora_cache_pattern="Technically-Color", # Test fresh LoRA download works + mismatch_threshold=0.35, # LoRA tests have higher variance + ) diff --git a/tests/image_generation/test_image_util.py b/tests/image_generation/test_image_util.py index 1ba0714..dec8c14 100644 --- a/tests/image_generation/test_image_util.py +++ b/tests/image_generation/test_image_util.py @@ -12,6 +12,7 @@ def test_image(): return PIL.Image.new("RGB", (100, 80), color="blue") +@pytest.mark.fast def test_expand_image_with_pixels(test_image): # Test expanding with pixel values expanded = ImageUtil.expand_image( @@ -37,6 +38,7 @@ def test_expand_image_with_pixels(test_image): assert expanded.getpixel((15, 25)) == (0, 0, 255) # blue +@pytest.mark.fast def test_expand_image_with_percentages(test_image): # Test expanding with percentage values expanded = ImageUtil.expand_image( @@ -65,6 +67,7 @@ def test_expand_image_with_percentages(test_image): assert expanded.getpixel((0, 0)) == (0, 255, 0) +@pytest.mark.fast def test_expand_image_with_invalid_values(test_image): # Test with invalid percentage string (a string that does not end with %" with pytest.raises(ValueError): @@ -80,6 +83,7 @@ def test_expand_image_with_invalid_values(test_image): ImageUtil.expand_image(test_image, top="25", right="30", bottom="50#", left="10*") +@pytest.mark.fast def test_create_outpaint_mask_image(): # Test with various padding values orig_width = 100 @@ -136,6 +140,7 @@ def test_create_outpaint_mask_image(): assert mask.getpixel((left_padding, top_padding + 10)) == (0, 0, 0) +@pytest.mark.fast def test_binarize(): # Create test data using numpy arrays and convert to mx arrays # Create a gradient array from 0 to 1 @@ -160,6 +165,7 @@ def test_binarize(): assert np.all(result_np[:, :, :, 5:] == 1) +@pytest.mark.fast def test_to_array_with_mask(): # Create a test image with gradient colors from PIL import Image, ImageDraw diff --git a/tests/image_generation/test_stdin_prompt_integration.py b/tests/image_generation/test_stdin_prompt_integration.py index 1867fc7..6ef9b49 100644 --- a/tests/image_generation/test_stdin_prompt_integration.py +++ b/tests/image_generation/test_stdin_prompt_integration.py @@ -11,17 +11,17 @@ def temp_output_dir(tmp_path) -> Path: return tmp_path +@pytest.mark.slow def test_stdin_prompt_with_actual_generation(temp_output_dir): - """Test that stdin prompt is correctly used in actual image generation and saved in metadata.""" stdin_prompt = "A beautiful mountain landscape with snow" output_image = temp_output_dir / "test_stdin.png" metadata_file = temp_output_dir / "test_stdin.json" - # Run the actual mflux.generate module with stdin + # Run the actual flux_generate module with stdin cmd = [ sys.executable, "-m", - "mflux.generate", + "mflux.models.flux.cli.flux_generate", "--prompt", "-", "--model", @@ -64,8 +64,8 @@ def test_stdin_prompt_with_actual_generation(temp_output_dir): assert metadata["prompt"] == stdin_prompt +@pytest.mark.slow def test_stdin_prompt_multiline_with_actual_generation(temp_output_dir): - """Test that multiline stdin prompt is correctly preserved in metadata.""" stdin_prompt = """A fantasy scene with: - Dragons flying in the sky - A castle on a mountain @@ -77,7 +77,7 @@ def test_stdin_prompt_multiline_with_actual_generation(temp_output_dir): cmd = [ sys.executable, "-m", - "mflux.generate", + "mflux.models.flux.cli.flux_generate", "--prompt", "-", "--model", @@ -109,14 +109,14 @@ def test_stdin_prompt_multiline_with_actual_generation(temp_output_dir): assert metadata["prompt"] == stdin_prompt.strip() +@pytest.mark.slow def test_empty_stdin_fails_generation(temp_output_dir): - """Test that empty stdin causes generation to fail with appropriate error.""" output_image = temp_output_dir / "test_empty.png" cmd = [ sys.executable, "-m", - "mflux.generate", + "mflux.models.flux.cli.flux_generate", "--prompt", "-", "--model", @@ -146,14 +146,14 @@ def test_empty_stdin_fails_generation(temp_output_dir): assert "No prompt provided via stdin" in process.stdout +@pytest.mark.slow def test_pipe_from_echo_command(temp_output_dir): - """Test using echo command to pipe prompt, simulating real usage.""" prompt = "A serene lake at sunset" output_image = temp_output_dir / "test_echo.png" metadata_file = temp_output_dir / "test_echo.json" # Simulate: echo "prompt" | mflux-generate --prompt - ... - echo_cmd = f'echo "{prompt}" | {sys.executable} -m mflux.generate --prompt - --model dev --steps 1 --height 256 --width 256 -q 4 --output {output_image} --metadata' + echo_cmd = f'echo "{prompt}" | {sys.executable} -m mflux.models.flux.cli.flux_generate --prompt - --model dev --steps 1 --height 256 --width 256 -q 4 --output {output_image} --metadata' process = subprocess.run(echo_cmd, shell=True, capture_output=True, text=True) diff --git a/tests/image_generation/test_upscale_dimensions.py b/tests/image_generation/test_upscale_dimensions.py index 7fae51e..14ff42c 100644 --- a/tests/image_generation/test_upscale_dimensions.py +++ b/tests/image_generation/test_upscale_dimensions.py @@ -1,10 +1,11 @@ -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import pytest -from mflux.ui.scale_factor import ScaleFactor +from mflux.utils.scale_factor import ScaleFactor +@pytest.mark.fast @pytest.mark.parametrize( "args_height,args_width,orig_height,orig_width,expected_height,expected_width", [ @@ -21,22 +22,23 @@ from mflux.ui.scale_factor import ScaleFactor def test_upscale_passes_correct_dimensions_to_generate_image( args_height, args_width, orig_height, orig_width, expected_height, expected_width ): - """Test that upscale.py passes the correct dimensions to generate_image""" - # Mock the image that will be opened - mock_image = Mock() + # Mock the image that will be opened - needs to support context manager protocol + mock_image = MagicMock() mock_image.size = (orig_width, orig_height) mock_image.height = orig_height mock_image.width = orig_width + mock_image.__enter__ = Mock(return_value=mock_image) + mock_image.__exit__ = Mock(return_value=False) # Mock the flux object mock_flux = Mock() # Import and patch the actual upscale module with patch("PIL.Image.open", return_value=mock_image): - with patch("mflux.upscale.Flux1Controlnet", return_value=mock_flux): - with patch("mflux.upscale.ModelConfig"): - with patch("mflux.upscale.CallbackManager"): - from mflux.upscale import main + with patch("mflux.models.flux.cli.flux_upscale.Flux1Controlnet", return_value=mock_flux): + with patch("mflux.models.flux.cli.flux_upscale.ModelConfig"): + with patch("mflux.models.flux.cli.flux_upscale.CallbackManager"): + from mflux.models.flux.cli.flux_upscale import main # Mock command line args mock_args = Mock() @@ -52,19 +54,20 @@ def test_upscale_passes_correct_dimensions_to_generate_image( mock_args.lora_paths = None mock_args.lora_scales = None - with patch("mflux.upscale.CommandLineParser") as mock_parser_class: + with patch("mflux.models.flux.cli.flux_upscale.CommandLineParser") as mock_parser_class: mock_parser = Mock() mock_parser.parse_args.return_value = mock_args mock_parser_class.return_value = mock_parser - with patch("mflux.upscale.PromptUtils.get_effective_prompt", return_value="test prompt"): + with patch( + "mflux.models.flux.cli.flux_upscale.PromptUtil.read_prompt", return_value="test prompt" + ): # Call the main function main() # Verify generate_image was called with correct dimensions mock_flux.generate_image.assert_called() call_args = mock_flux.generate_image.call_args - config = call_args.kwargs["config"] - assert config.height == expected_height - assert config.width == expected_width + assert call_args.kwargs["height"] == expected_height + assert call_args.kwargs["width"] == expected_width diff --git a/tests/metadata/test_metadata.py b/tests/metadata/test_metadata.py index eec3d53..c58cbc9 100644 --- a/tests/metadata/test_metadata.py +++ b/tests/metadata/test_metadata.py @@ -1,34 +1,19 @@ -""" -Test metadata embedding and reading functionality. - -This test generates a single small image using the schnell model and verifies -that all metadata fields are correctly embedded and can be read back. -""" - import json import os import tempfile from datetime import datetime from pathlib import Path -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +import pytest + +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.txt2img.flux import Flux1 from mflux.utils.metadata_reader import MetadataReader class TestMetadata: - """Test suite for image metadata functionality.""" - + @pytest.mark.slow def test_metadata_complete(self): - """ - Comprehensive test that generates one image using img2img and verifies: - - EXIF metadata (all fields including img2img-specific ones) - - mflux-info command output formatting - - Creation timestamp validity - - Optional fields handling (both None and populated) - - Metadata reader handles missing files gracefully - """ # Use a temporary file - don't keep it open to avoid conflicts fd, temp_path = tempfile.mkstemp(suffix=".png") os.close(fd) # Close the file descriptor immediately @@ -44,7 +29,9 @@ class TestMetadata: quantize=8, ) - config = Config( + image = flux.generate_image( + seed=42, + prompt="A simple test image", num_inference_steps=2, height=256, width=256, @@ -52,12 +39,6 @@ class TestMetadata: image_strength=0.3, ) - image = flux.generate_image( - seed=42, - prompt="A simple test image", - config=config, - ) - # Save with metadata (overwrite=True since mkstemp creates an empty file) image.save(path=output_path, overwrite=True) @@ -159,9 +140,9 @@ class TestMetadata: # ================================================================= # Test 11: mflux-info command output # ================================================================= - from mflux.info import format_metadata + from mflux.utils.info_util import InfoUtil - output = format_metadata(metadata) + output = InfoUtil.format_metadata(metadata) assert "A simple test image" in output, "Prompt should be in output" assert "MFLUX" in output, "MFLUX should be mentioned" assert "42" in output, "Seed should be in output" diff --git a/tests/model_config/test_model_config.py b/tests/model_config/test_model_config.py deleted file mode 100644 index 3b0aebc..0000000 --- a/tests/model_config/test_model_config.py +++ /dev/null @@ -1,253 +0,0 @@ -import pytest - -from mflux.config.model_config import ModelConfig -from mflux.utils.exceptions import InvalidBaseModel, ModelConfigError - - -def test_bfl_dev(): - model = ModelConfig.from_name("dev") - assert model.model_name == "black-forest-labs/FLUX.1-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_full_name(): - model = ModelConfig.from_name("black-forest-labs/FLUX.1-dev") - assert model.model_name == "black-forest-labs/FLUX.1-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_krea_dev(): - model = ModelConfig.from_name("krea-dev") - assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_krea_alias(): - model = ModelConfig.from_name("dev-krea") - assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_krea_dev_full_name(): - model = ModelConfig.from_name("black-forest-labs/FLUX.1-Krea-dev") - assert model.model_name == "black-forest-labs/FLUX.1-Krea-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_schnell(): - model = ModelConfig.from_name("schnell") - assert model.model_name == "black-forest-labs/FLUX.1-schnell" - assert model.max_sequence_length == 256 - assert model.num_train_steps == 1000 - assert model.supports_guidance is False - assert model.requires_sigma_shift is False - - -def test_bfl_schnell_full_name(): - model = ModelConfig.from_name("black-forest-labs/FLUX.1-schnell") - assert model.model_name == "black-forest-labs/FLUX.1-schnell" - assert model.max_sequence_length == 256 - assert model.num_train_steps == 1000 - assert model.supports_guidance is False - assert model.requires_sigma_shift is False - - -def test_bfl_dev_fill(): - model = ModelConfig.from_name("dev-fill") - assert model.model_name == "black-forest-labs/FLUX.1-Fill-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_fill_full_name(): - model = ModelConfig.from_name("black-forest-labs/FLUX.1-Fill-dev") - assert model.model_name == "black-forest-labs/FLUX.1-Fill-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_depth(): - model = ModelConfig.from_name("dev-depth") - assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_depth_full_name(): - model = ModelConfig.from_name("black-forest-labs/FLUX.1-Depth-dev") - assert model.model_name == "black-forest-labs/FLUX.1-Depth-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_redux(): - model = ModelConfig.from_name("dev-redux") - assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_redux_full_name(): - model = ModelConfig.from_name("black-forest-labs/FLUX.1-Redux-dev") - assert model.model_name == "black-forest-labs/FLUX.1-Redux-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_bfl_dev_controlnet_canny(): - model = ModelConfig.from_name("dev-controlnet-canny") - assert model.model_name == "black-forest-labs/FLUX.1-dev" - assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - assert model.is_canny() is True - - -def test_bfl_schnell_controlnet_canny(): - model = ModelConfig.from_name("schnell-controlnet-canny") - assert model.model_name == "black-forest-labs/FLUX.1-schnell" - assert model.controlnet_model == "InstantX/FLUX.1-dev-Controlnet-Canny" - assert model.max_sequence_length == 256 - assert model.num_train_steps == 1000 - assert model.supports_guidance is False - assert model.requires_sigma_shift is False - assert model.is_canny() is True - - -def test_bfl_dev_controlnet_upscaler(): - model = ModelConfig.from_name("dev-controlnet-upscaler") - assert model.model_name == "black-forest-labs/FLUX.1-dev" - assert model.controlnet_model == "jasperai/Flux.1-dev-Controlnet-Upscaler" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is False - assert model.requires_sigma_shift is False - assert model.is_canny() is False - - -def test_community_dev_fill_implicit_base_model(): - model = ModelConfig.from_name("acme-lab/some-dev-fill-model") - assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_community_dev_fill_explicit_base_model(): - model = ModelConfig.from_name("acme-lab/some-model", base_model="dev-fill") - assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_implicit_base_model_prefers_dev_fill_over_dev(): - model = ModelConfig.from_name("acme-lab/dev-fill-based-model") - assert model.base_model == "black-forest-labs/FLUX.1-Fill-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.requires_sigma_shift is True - - -def test_community_dev_implicit_base_model(): - model = ModelConfig.from_name("acme-lab/some-awesome-dev-model") - assert model.base_model == "black-forest-labs/FLUX.1-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_community_schnell_implicit_base_model(): - model = ModelConfig.from_name("acme-lab/some-quick-schnell-model") - assert model.base_model == "black-forest-labs/FLUX.1-schnell" - assert model.max_sequence_length == 256 - assert model.num_train_steps == 1000 - assert model.supports_guidance is False - assert model.requires_sigma_shift is False - - -def test_community_dev_explicit_base_model(): - model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="dev") - assert model.base_model == "black-forest-labs/FLUX.1-dev" - assert model.max_sequence_length == 512 - assert model.num_train_steps == 1000 - assert model.supports_guidance is True - assert model.requires_sigma_shift is True - - -def test_community_schnell_explicit_base_model(): - model = ModelConfig.from_name("acme-lab/some-awesome-model", base_model="schnell") - assert model.base_model == "black-forest-labs/FLUX.1-schnell" - assert model.max_sequence_length == 256 - assert model.num_train_steps == 1000 - assert model.supports_guidance is False - assert model.requires_sigma_shift is False - - -def test_model_config_error(): - with pytest.raises(ModelConfigError): - ModelConfig.from_name("acme-lab/some-model-who-knows-what-its-based-on") - - -def test_invalid_base_model_error(): - with pytest.raises(InvalidBaseModel): - ModelConfig.from_name("acme-lab/some-model-who-knows-what-its-based-on", base_model="something_unknown") - - -def test_qwen_image(): - model = ModelConfig.from_name("qwen") - assert model.model_name == "Qwen/Qwen-Image" - assert model.max_sequence_length is None - assert model.num_train_steps is None - assert model.supports_guidance is None - assert model.requires_sigma_shift is None - - -def test_qwen_image_alias(): - model = ModelConfig.from_name("qwen-image") - assert model.model_name == "Qwen/Qwen-Image" - assert model.max_sequence_length is None - assert model.num_train_steps is None - assert model.supports_guidance is None - assert model.requires_sigma_shift is None - - -def test_qwen_image_full_name(): - model = ModelConfig.from_name("Qwen/Qwen-Image") - assert model.model_name == "Qwen/Qwen-Image" - assert model.max_sequence_length is None - assert model.num_train_steps is None - assert model.supports_guidance is None - assert model.requires_sigma_shift is None diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py index accf641..dd30449 100644 --- a/tests/model_saving/test_model_saving.py +++ b/tests/model_saving/test_model_saving.py @@ -3,17 +3,18 @@ import shutil from pathlib import Path import numpy as np +import pytest -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig +from mflux.models.common.weights.loading.weight_loader import WeightLoader from mflux.models.flux.variants.txt2img.flux import Flux1 -from mflux.models.flux.weights.weight_handler import WeightHandler from mflux.utils.version_util import VersionUtil PATH = "tests/4bit/" class TestModelSaving: + @pytest.mark.slow def test_save_and_load_4bit_model(self): # Clean up any existing temporary directories from previous test runs TestModelSaving.delete_folder_if_exists(PATH) @@ -27,35 +28,31 @@ class TestModelSaving: image1 = fluxA.generate_image( seed=42, prompt="Luxury food photograph", - config=Config( - num_inference_steps=15, - height=341, - width=768, - ), + num_inference_steps=15, + height=341, + width=768, ) fluxA.save_model(PATH) del fluxA # Verify that the mflux version is correctly saved in the model's metadata - _, quantization_level, mflux_version = WeightHandler._load_vae(root_path=Path(PATH)) + _, quantization_level, mflux_version = WeightLoader._try_load_mflux_format(Path(PATH) / "vae") assert mflux_version == VersionUtil.get_mflux_version(), "mflux version not correctly saved in metadata" # fmt: off - assert quantization_level == "4", "quantization level not correctly saved in metadata" # fmt: off + assert quantization_level == 4, "quantization level not correctly saved in metadata" # fmt: off # when loading the quantized model (also without specifying bits) fluxB = Flux1( model_config=ModelConfig.dev(), - local_path=PATH, + model_path=PATH, ) # then we can load the model and generate the identical image image2 = fluxB.generate_image( seed=42, prompt="Luxury food photograph", - config=Config( - num_inference_steps=15, - height=341, - width=768, - ), + num_inference_steps=15, + height=341, + width=768, ) np.testing.assert_array_equal( np.array(image1.image), diff --git a/tests/model_saving/test_model_saving_lora.py b/tests/model_saving/test_model_saving_lora.py index f8e8e7b..7fb932e 100644 --- a/tests/model_saving/test_model_saving_lora.py +++ b/tests/model_saving/test_model_saving_lora.py @@ -3,15 +3,16 @@ import shutil from pathlib import Path import numpy as np +import pytest -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig +from mflux.models.common.config import ModelConfig from mflux.models.flux.variants.txt2img.flux import Flux1 PATH = "tests/4bit/" class TestModelSavingLora: + @pytest.mark.slow def test_save_and_load_4bit_model_with_lora(self): # Clean up any existing temporary directories from previous test runs TestModelSavingLora.delete_folder_if_exists(PATH) @@ -35,18 +36,16 @@ class TestModelSavingLora: image1 = fluxB.generate_image( seed=44, prompt="mkym this is made of wool, pizza", - config=Config( - num_inference_steps=2, - height=341, - width=768, - ), + num_inference_steps=2, + height=341, + width=768, ) del fluxB # when loading the quantized model from a local path (also without specifying bits) with a LoRA... fluxC = Flux1( model_config=ModelConfig.schnell(), - local_path=PATH, + model_path=PATH, lora_paths=TestModelSavingLora.get_lora_path(), lora_scales=[1.0], ) @@ -55,11 +54,9 @@ class TestModelSavingLora: image2 = fluxC.generate_image( seed=44, prompt="mkym this is made of wool, pizza", - config=Config( - num_inference_steps=2, - height=341, - width=768, - ), + num_inference_steps=2, + height=341, + width=768, ) # then we confirm that we get the exact *identical* image in both cases diff --git a/tests/resolution/__init__.py b/tests/resolution/__init__.py new file mode 100644 index 0000000..5c8162d --- /dev/null +++ b/tests/resolution/__init__.py @@ -0,0 +1,2 @@ +# Tests for model loading policies + diff --git a/tests/resolution/test_config_resolution.py b/tests/resolution/test_config_resolution.py new file mode 100644 index 0000000..5f7bbf9 --- /dev/null +++ b/tests/resolution/test_config_resolution.py @@ -0,0 +1,109 @@ +import pytest + +from mflux.models.common.resolution.config_resolution import ConfigResolution +from mflux.utils.exceptions import InvalidBaseModel, ModelConfigError + + +class TestConfigResolutionExactMatch: + @pytest.mark.fast + def test_exact_alias_match(self): + config = ConfigResolution.resolve(model_name="schnell") + + assert config.model_name == "black-forest-labs/FLUX.1-schnell" + assert "schnell" in config.aliases + + @pytest.mark.fast + def test_exact_alias_match_dev(self): + config = ConfigResolution.resolve(model_name="dev") + + assert config.model_name == "black-forest-labs/FLUX.1-dev" + + @pytest.mark.fast + def test_exact_alias_match_fibo(self): + config = ConfigResolution.resolve(model_name="fibo") + + assert config.model_name == "briaai/FIBO" + + @pytest.mark.fast + def test_exact_hf_name_match(self): + config = ConfigResolution.resolve(model_name="black-forest-labs/FLUX.1-schnell") + + assert config.model_name == "black-forest-labs/FLUX.1-schnell" + + +class TestConfigResolutionExplicitBase: + @pytest.mark.fast + def test_explicit_base_model(self): + config = ConfigResolution.resolve(model_name="my-custom-model", base_model="schnell") + + assert config.model_name == "my-custom-model" + assert config.base_model == "black-forest-labs/FLUX.1-schnell" + assert config.max_sequence_length == 256 # schnell's value + + @pytest.mark.fast + def test_explicit_base_model_dev(self): + config = ConfigResolution.resolve(model_name="org/my-finetune", base_model="dev") + + assert config.model_name == "org/my-finetune" + assert config.base_model == "black-forest-labs/FLUX.1-dev" + assert config.supports_guidance is True # dev's value + + @pytest.mark.fast + def test_invalid_base_model_raises(self): + with pytest.raises(InvalidBaseModel): + ConfigResolution.resolve(model_name="whatever", base_model="invalid-base") + + +class TestConfigResolutionInferSubstring: + @pytest.mark.fast + def test_infer_from_schnell_substring(self): + config = ConfigResolution.resolve(model_name="my-schnell-finetune") + + assert config.model_name == "my-schnell-finetune" + assert config.base_model == "black-forest-labs/FLUX.1-schnell" + + @pytest.mark.fast + def test_infer_from_dev_substring(self): + config = ConfigResolution.resolve(model_name="dev-lora-something") + + assert config.model_name == "dev-lora-something" + assert config.base_model == "black-forest-labs/FLUX.1-dev" + + @pytest.mark.fast + def test_infer_case_insensitive(self): + config = ConfigResolution.resolve(model_name="MY-SCHNELL-MODEL") + + assert config.base_model == "black-forest-labs/FLUX.1-schnell" + + @pytest.mark.fast + def test_longer_alias_preferred(self): + # "dev-kontext" is longer than "dev", should match dev-kontext if present + config = ConfigResolution.resolve(model_name="my-dev-kontext-model") + + assert config.base_model == "black-forest-labs/FLUX.1-Kontext-dev" + + +class TestConfigResolutionError: + @pytest.mark.fast + def test_unknown_model_without_base_raises(self): + with pytest.raises(ModelConfigError) as exc_info: + ConfigResolution.resolve(model_name="totally-unknown-model") + + assert "Cannot infer" in str(exc_info.value) + + +class TestConfigResolutionRules: + @pytest.mark.fast + def test_exact_match_takes_priority(self): + # "schnell" is both an exact alias AND would match substring + config = ConfigResolution.resolve(model_name="schnell") + + # Should return the exact config, not create a new one + assert config.model_name == "black-forest-labs/FLUX.1-schnell" + + @pytest.mark.fast + def test_explicit_base_overrides_inference(self): + # Model name contains "schnell" but explicit base is "dev" + config = ConfigResolution.resolve(model_name="schnell-style-dev", base_model="dev") + + assert config.base_model == "black-forest-labs/FLUX.1-dev" diff --git a/tests/resolution/test_lora_resolution.py b/tests/resolution/test_lora_resolution.py new file mode 100644 index 0000000..a0bc6d3 --- /dev/null +++ b/tests/resolution/test_lora_resolution.py @@ -0,0 +1,268 @@ +from unittest.mock import patch + +import pytest +from huggingface_hub.utils import LocalEntryNotFoundError + +from mflux.models.common.resolution.lora_resolution import LoraResolution + + +class TestLoraResolutionLocal: + @pytest.mark.fast + def test_existing_local_file(self, tmp_path): + lora_file = tmp_path / "my-lora.safetensors" + lora_file.touch() + + result = LoraResolution.resolve(path=str(lora_file)) + + assert result == str(lora_file) + + +class TestLoraResolutionRegistry: + @pytest.mark.fast + def test_registry_lookup(self, tmp_path): + lora_file = tmp_path / "registered-lora.safetensors" + lora_file.touch() + + original_registry = LoraResolution._registry.copy() + try: + LoraResolution._registry["my-alias"] = lora_file + + result = LoraResolution.resolve(path="my-alias") + + assert result == str(lora_file) + finally: + LoraResolution._registry = original_registry + + +class TestLoraResolutionHuggingFace: + @pytest.mark.fast + @patch("mflux.models.common.resolution.lora_resolution.snapshot_download") + def test_huggingface_repo_downloads_when_not_cached(self, mock_download, tmp_path): + lora_file = tmp_path / "lora.safetensors" + lora_file.touch() + # First call (cache check) raises, second call (download) succeeds + mock_download.side_effect = [ + LocalEntryNotFoundError("Not cached"), + str(tmp_path), + ] + + result = LoraResolution.resolve(path="org/lora-repo") + + assert mock_download.call_count == 2 + # First call should have local_files_only=True + first_call = mock_download.call_args_list[0] + assert first_call[1].get("local_files_only") is True + assert result == str(lora_file) + + @pytest.mark.fast + @patch("mflux.models.common.resolution.lora_resolution.snapshot_download") + def test_huggingface_repo_uses_cache_when_available(self, mock_download, tmp_path): + lora_file = tmp_path / "lora.safetensors" + lora_file.touch() + # Cache check succeeds - no network download needed + mock_download.return_value = str(tmp_path) + + result = LoraResolution.resolve(path="org/lora-repo") + + # Called twice with local_files_only=True (once in check, once in execute) + assert mock_download.call_count == 2 + for call in mock_download.call_args_list: + assert call[1].get("local_files_only") is True + assert result == str(lora_file) + + @pytest.mark.fast + def test_collection_format_is_detected(self): + # Collection format: "repo:filename" with a "/" in repo + assert LoraResolution._is_collection_format("org/repo:file.safetensors") + assert not LoraResolution._is_collection_format("org/repo") + assert not LoraResolution._is_collection_format("local:file") + + @pytest.mark.fast + @patch("mflux.models.common.resolution.lora_resolution.snapshot_download") + def test_multiple_safetensors_in_repo_raises_error(self, mock_download, tmp_path): + # Create multiple .safetensors files in the directory + (tmp_path / "lora_v1.safetensors").touch() + (tmp_path / "lora_v2.safetensors").touch() + + # First call (cache check) raises, second call (download) succeeds + mock_download.side_effect = [ + LocalEntryNotFoundError("Not cached"), + str(tmp_path), + ] + + with pytest.raises(ValueError) as exc_info: + LoraResolution.resolve(path="org/multi-lora-repo") + + error_msg = str(exc_info.value) + assert "Multiple .safetensors files found" in error_msg + assert "lora_v1.safetensors" in error_msg or "lora_v2.safetensors" in error_msg + assert "collection format" in error_msg + assert "org/multi-lora-repo:" in error_msg + + @pytest.mark.fast + @patch("mflux.models.common.resolution.lora_resolution.snapshot_download") + def test_multiple_safetensors_cached_raises_error(self, mock_download, tmp_path): + # Create multiple .safetensors files in the directory + (tmp_path / "lora_a.safetensors").touch() + (tmp_path / "lora_b.safetensors").touch() + + # Cache check succeeds (both calls return same path) + mock_download.return_value = str(tmp_path) + + with pytest.raises(ValueError) as exc_info: + LoraResolution.resolve(path="org/cached-multi-lora") + + assert "Multiple .safetensors files found" in str(exc_info.value) + + @pytest.mark.fast + @patch("mflux.models.common.resolution.lora_resolution.snapshot_download") + def test_single_safetensor_in_repo_succeeds(self, mock_download, tmp_path): + # Create only one .safetensors file + lora_file = tmp_path / "single-lora.safetensors" + lora_file.touch() + + mock_download.side_effect = [ + LocalEntryNotFoundError("Not cached"), + str(tmp_path), + ] + + result = LoraResolution.resolve(path="org/single-lora-repo") + + assert result == str(lora_file) + + +class TestLoraResolutionError: + @pytest.mark.fast + def test_nonexistent_file_raises(self): + with pytest.raises(FileNotFoundError) as exc_info: + LoraResolution.resolve(path="nonexistent-lora") + + assert "LoRA file not found" in str(exc_info.value) + + +class TestLoraResolutionPaths: + @pytest.mark.fast + def test_resolve_paths_empty(self): + result = LoraResolution.resolve_paths(paths=None) + assert result == [] + + @pytest.mark.fast + def test_resolve_paths_filters_invalid(self, tmp_path): + valid_lora = tmp_path / "valid.safetensors" + valid_lora.touch() + + result = LoraResolution.resolve_paths(paths=[str(valid_lora), "invalid-path"]) + + assert len(result) == 1 + assert result[0] == str(valid_lora) + + +class TestLoraResolutionScales: + @pytest.mark.fast + def test_resolve_scales_none(self): + result = LoraResolution.resolve_scales(scales=None, num_paths=3) + assert result == [1.0, 1.0, 1.0] + + @pytest.mark.fast + def test_resolve_scales_provided(self): + result = LoraResolution.resolve_scales(scales=[0.5, 0.8], num_paths=2) + assert result == [0.5, 0.8] + + @pytest.mark.fast + def test_resolve_scales_too_few_pads_with_default(self, capsys): + # 2 scales provided but 3 paths - should pad with 1.0 + result = LoraResolution.resolve_scales(scales=[0.5, 0.8], num_paths=3) + + assert result == [0.5, 0.8, 1.0] + captured = capsys.readouterr() + assert "doesn't match" in captured.out + + @pytest.mark.fast + def test_resolve_scales_too_many_truncates(self, capsys): + # 3 scales provided but only 2 paths - should truncate + result = LoraResolution.resolve_scales(scales=[0.5, 0.8, 0.9], num_paths=2) + + assert result == [0.5, 0.8] + captured = capsys.readouterr() + assert "doesn't match" in captured.out + + @pytest.mark.fast + def test_resolve_scales_matching_no_warning(self, capsys): + # Matching counts should not warn + result = LoraResolution.resolve_scales(scales=[0.5, 0.8], num_paths=2) + + assert result == [0.5, 0.8] + captured = capsys.readouterr() + assert "doesn't match" not in captured.out + + +class TestLoraResolutionRules: + @pytest.mark.fast + def test_local_takes_priority_over_registry(self, tmp_path): + lora_file = tmp_path / "lora.safetensors" + lora_file.touch() + + original_registry = LoraResolution._registry.copy() + try: + LoraResolution._registry[str(lora_file)] = tmp_path / "other.safetensors" + + result = LoraResolution.resolve(path=str(lora_file)) + + assert result == str(lora_file) + finally: + LoraResolution._registry = original_registry + + +class TestLoraResolutionTildeExpansion: + @pytest.mark.fast + def test_expands_home_directory(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + lora_dir = tmp_path / "loras" + lora_dir.mkdir() + lora_file = lora_dir / "my-style.safetensors" + lora_file.touch() + + result = LoraResolution.resolve(path="~/loras/my-style.safetensors") + + assert result == str(lora_file) + + @pytest.mark.fast + def test_tilde_path_not_treated_as_huggingface(self, tmp_path, monkeypatch): + # ~/loras/file has exactly one slash after tilde expansion path check + # It should NOT be treated as HuggingFace format + monkeypatch.setenv("HOME", str(tmp_path)) + lora_file = tmp_path / "org" / "lora.safetensors" + lora_file.parent.mkdir(parents=True) + lora_file.touch() + + # This path looks like "~/org/lora.safetensors" which could be confused with HF format + result = LoraResolution.resolve(path="~/org/lora.safetensors") + + assert result == str(lora_file) + + +class TestLoraResolutionRelativePaths: + @pytest.mark.fast + def test_relative_path_not_treated_as_huggingface(self): + # ./org/model should NOT be treated as HuggingFace even if it doesn't exist + # It should fail as a local path, not try to download from HF + assert not LoraResolution._is_hf_format("./org/model") + assert not LoraResolution._is_hf_format("../org/model") + + @pytest.mark.fast + def test_dotslash_single_segment_not_huggingface(self): + # ./lora has exactly one slash - should NOT match HuggingFace format + assert not LoraResolution._is_hf_format("./lora") + assert not LoraResolution._is_hf_format("../lora") + + @pytest.mark.fast + def test_relative_path_existing_file_resolves_locally(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + lora_dir = tmp_path / "subdir" + lora_dir.mkdir() + lora_file = lora_dir / "style.safetensors" + lora_file.touch() + + result = LoraResolution.resolve(path="./subdir/style.safetensors") + + assert "style.safetensors" in result diff --git a/tests/resolution/test_path_resolution.py b/tests/resolution/test_path_resolution.py new file mode 100644 index 0000000..93da4b8 --- /dev/null +++ b/tests/resolution/test_path_resolution.py @@ -0,0 +1,207 @@ +from unittest.mock import patch + +import pytest + +from mflux.models.common.resolution.path_resolution import PathResolution + + +class TestPathResolutionNone: + @pytest.mark.fast + def test_none_returns_none(self): + result = PathResolution.resolve(path=None) + assert result is None + + +class TestPathResolutionLocal: + @pytest.mark.fast + def test_existing_local_path(self, tmp_path): + model_dir = tmp_path / "my-model" + model_dir.mkdir() + + result = PathResolution.resolve(path=str(model_dir)) + + assert result == model_dir + + @pytest.mark.fast + def test_expands_home_directory(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + model_dir = tmp_path / "models" / "test" + model_dir.mkdir(parents=True) + + result = PathResolution.resolve(path="~/models/test") + + assert result == model_dir + + @pytest.mark.fast + def test_local_path_with_slash_preferred_over_huggingface(self, tmp_path): + # Create a local path that looks like org/model + org_dir = tmp_path / "org" + org_dir.mkdir() + model_dir = org_dir / "model" + model_dir.mkdir() + + result = PathResolution.resolve(path=str(model_dir)) + + assert result == model_dir + + @pytest.mark.fast + def test_relative_path_not_treated_as_huggingface(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + model_dir = tmp_path / "org" / "model" + model_dir.mkdir(parents=True) + + result = PathResolution.resolve(path="./org/model") + + assert result.name == "model" + + @pytest.mark.fast + def test_tilde_path_not_treated_as_huggingface(self): + # ~/org/model should NOT be treated as HuggingFace even if it doesn't exist + # It should fail as a local path, not try to download from HF + assert not PathResolution._is_hf_format("~/org/model") + assert not PathResolution._is_hf_format("~/models/custom") + + @pytest.mark.fast + def test_nonexistent_tilde_path_raises_local_error(self): + # A tilde path that doesn't exist should raise FileNotFoundError + # NOT try to download from HuggingFace + with pytest.raises(FileNotFoundError) as exc_info: + PathResolution.resolve(path="~/nonexistent/model") + + # Error should indicate it's a local path issue, not HF + assert "Model not found" in str(exc_info.value) + + +class TestPathResolutionHuggingFace: + @pytest.mark.fast + @patch("mflux.models.common.resolution.path_resolution.snapshot_download") + def test_huggingface_format_downloads_when_not_cached(self, mock_download, tmp_path): + # No cache exists, so snapshot_download is called once to download + mock_download.return_value = str(tmp_path / "cached") + + result = PathResolution.resolve(path="org/model") + + # Called once (download only - cache check uses _find_complete_cached_snapshot) + assert mock_download.call_count == 1 + call_kwargs = mock_download.call_args[1] + assert "local_files_only" not in call_kwargs + assert result == tmp_path / "cached" + + @pytest.mark.fast + def test_huggingface_uses_cache_when_available(self, tmp_path): + # Create a fake cached snapshot structure + repo_cache = tmp_path / "models--org--model" / "snapshots" / "abc123" + repo_cache.mkdir(parents=True) + (repo_cache / "model.safetensors").touch() + + with patch("mflux.models.common.resolution.path_resolution.HF_HUB_CACHE", str(tmp_path)): + result = PathResolution.resolve(path="org/model") + + # Should return the cached path without calling snapshot_download + assert result == repo_cache + + @pytest.mark.fast + @patch("mflux.models.common.resolution.path_resolution.snapshot_download") + def test_huggingface_passes_patterns(self, mock_download, tmp_path): + mock_download.return_value = str(tmp_path / "cached") + + PathResolution.resolve(path="org/model", patterns=["*.bin", "*.json"]) + + call_kwargs = mock_download.call_args[1] + assert call_kwargs["allow_patterns"] == ["*.bin", "*.json"] + + +class TestPathResolutionError: + @pytest.mark.fast + def test_nonexistent_local_path_raises(self): + with pytest.raises(FileNotFoundError) as exc_info: + PathResolution.resolve(path="/nonexistent/path/to/model") + + assert "Model not found" in str(exc_info.value) + assert "/nonexistent/path/to/model" in str(exc_info.value) + + @pytest.mark.fast + def test_invalid_format_raises(self): + with pytest.raises(FileNotFoundError) as exc_info: + PathResolution.resolve(path="not-a-valid-path") + + assert "Model not found" in str(exc_info.value) + + @pytest.mark.fast + def test_error_message_is_helpful(self): + with pytest.raises(FileNotFoundError) as exc_info: + PathResolution.resolve(path="/bad/path") + + error_msg = str(exc_info.value) + assert "local path" in error_msg.lower() + assert "org/model" in error_msg + + +class TestPathResolutionRules: + @pytest.mark.fast + def test_rules_are_checked_in_order(self, tmp_path): + # Create a local path that also matches HuggingFace format + model_dir = tmp_path / "org" / "model" + model_dir.mkdir(parents=True) + + # Local should win because it's checked first + result = PathResolution.resolve(path=str(model_dir)) + + assert result == model_dir + + @pytest.mark.fast + def test_relative_paths_are_local_not_huggingface(self): + # ./org/model should NOT be treated as HuggingFace + # It should fail because the local path doesn't exist + with pytest.raises(FileNotFoundError): + PathResolution.resolve(path="./org/model") + + @pytest.mark.fast + def test_parent_relative_paths_are_local_not_huggingface(self): + # ../org/model should NOT be treated as HuggingFace + with pytest.raises(FileNotFoundError): + PathResolution.resolve(path="../org/model") + + +class TestPathResolutionEmptyDirectory: + @pytest.mark.fast + def test_empty_directory_prints_warning(self, tmp_path, capsys): + # Create an empty directory + empty_model_dir = tmp_path / "empty-model" + empty_model_dir.mkdir() + + result = PathResolution.resolve(path=str(empty_model_dir)) + + # Should resolve (directory exists) but warn about missing files + assert result == empty_model_dir + captured = capsys.readouterr() + assert "contains no files matching" in captured.out + + @pytest.mark.fast + def test_directory_with_matching_files_no_warning(self, tmp_path, capsys): + # Create a directory with matching files + model_dir = tmp_path / "model-with-weights" + model_dir.mkdir() + (model_dir / "weights.safetensors").touch() + + result = PathResolution.resolve(path=str(model_dir)) + + # Should resolve without warning + assert result == model_dir + captured = capsys.readouterr() + assert "contains no files matching" not in captured.out + + @pytest.mark.fast + def test_empty_directory_with_custom_patterns(self, tmp_path, capsys): + # Create a directory with .bin file but looking for .json + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "weights.bin").touch() + + result = PathResolution.resolve(path=str(model_dir), patterns=["*.json"]) + + # Should warn because no .json files found + assert result == model_dir + captured = capsys.readouterr() + assert "contains no files matching" in captured.out + assert "*.json" in captured.out diff --git a/tests/resolution/test_quantization_resolution.py b/tests/resolution/test_quantization_resolution.py new file mode 100644 index 0000000..9425d5a --- /dev/null +++ b/tests/resolution/test_quantization_resolution.py @@ -0,0 +1,70 @@ +import pytest + +from mflux.models.common.resolution.quantization_resolution import QuantizationResolution + + +class TestDecideQuantization: + @pytest.mark.fast + def test_no_quantization_when_neither_specified(self): + bits, warning = QuantizationResolution.resolve(stored=None, requested=None) + assert bits is None + assert warning is None + + @pytest.mark.fast + @pytest.mark.parametrize("requested_bits", [3, 4, 5, 6, 8]) + def test_on_the_fly_quantization(self, requested_bits): + bits, warning = QuantizationResolution.resolve(stored=None, requested=requested_bits) + assert bits == requested_bits + assert warning is None + + @pytest.mark.fast + @pytest.mark.parametrize("stored_bits", [3, 4, 5, 6, 8]) + def test_prequantized_no_request(self, stored_bits): + bits, warning = QuantizationResolution.resolve(stored=stored_bits, requested=None) + assert bits == stored_bits + assert warning is None + + @pytest.mark.fast + @pytest.mark.parametrize("bits_value", [3, 4, 5, 6, 8]) + def test_prequantized_matching_request(self, bits_value): + bits, warning = QuantizationResolution.resolve(stored=bits_value, requested=bits_value) + assert bits == bits_value + assert warning is None + + @pytest.mark.fast + @pytest.mark.parametrize( + "stored_bits,requested_bits", + [(4, 8), (8, 4), (4, 3), (3, 8), (6, 4)], + ) + def test_prequantized_conflicting_request_uses_stored(self, stored_bits, requested_bits): + bits, warning = QuantizationResolution.resolve(stored=stored_bits, requested=requested_bits) + assert bits == stored_bits + + @pytest.mark.fast + @pytest.mark.parametrize( + "stored_bits,requested_bits", + [(4, 8), (8, 4), (4, 3)], + ) + def test_prequantized_conflicting_request_warns(self, stored_bits, requested_bits): + bits, warning = QuantizationResolution.resolve(stored=stored_bits, requested=requested_bits) + assert warning is not None + assert f"{stored_bits}-bit" in warning + assert f"-q {requested_bits}" in warning + + +class TestQuantizationPolicyCompleteness: + VALID_QUANT_LEVELS = [None, 3, 4, 5, 6, 8] + + @pytest.mark.fast + def test_all_combinations_handled(self): + for stored in self.VALID_QUANT_LEVELS: + for requested in self.VALID_QUANT_LEVELS: + bits, warning = QuantizationResolution.resolve(stored=stored, requested=requested) + assert bits is None or bits in self.VALID_QUANT_LEVELS + + @pytest.mark.fast + def test_result_is_always_stored_or_requested_or_none(self): + for stored in self.VALID_QUANT_LEVELS: + for requested in self.VALID_QUANT_LEVELS: + bits, _ = QuantizationResolution.resolve(stored=stored, requested=requested) + assert bits in (stored, requested, None) diff --git a/tests/resources/reference_controlnet_dev_lora_depth.png b/tests/resources/reference_controlnet_dev_lora_depth.png index 0cd2da4..eebb7f2 100644 Binary files a/tests/resources/reference_controlnet_dev_lora_depth.png and b/tests/resources/reference_controlnet_dev_lora_depth.png differ diff --git a/tests/resources/reference_fibo.png b/tests/resources/reference_fibo.png index 55ade30..ffd5f14 100644 Binary files a/tests/resources/reference_fibo.png and b/tests/resources/reference_fibo.png differ diff --git a/tests/resources/reference_fibo_white_owl.png b/tests/resources/reference_fibo_white_owl.png index 994fb3d..dee1030 100644 Binary files a/tests/resources/reference_fibo_white_owl.png and b/tests/resources/reference_fibo_white_owl.png differ diff --git a/tests/resources/reference_qwen_txt2img.png b/tests/resources/reference_qwen_txt2img.png index b0f6139..ff14fd8 100644 Binary files a/tests/resources/reference_qwen_txt2img.png and b/tests/resources/reference_qwen_txt2img.png differ diff --git a/tests/resources/reference_z_image_turbo.png b/tests/resources/reference_z_image_turbo.png new file mode 100644 index 0000000..25c1775 Binary files /dev/null and b/tests/resources/reference_z_image_turbo.png differ diff --git a/tests/resources/reference_z_image_turbo_lora.png b/tests/resources/reference_z_image_turbo_lora.png new file mode 100644 index 0000000..e865409 Binary files /dev/null and b/tests/resources/reference_z_image_turbo_lora.png differ diff --git a/tests/schedulers/test_base_scheduler.py b/tests/schedulers/test_base_scheduler.py index c9f08a0..9eec8b1 100644 --- a/tests/schedulers/test_base_scheduler.py +++ b/tests/schedulers/test_base_scheduler.py @@ -3,9 +3,7 @@ import pytest from mflux.models.common.schedulers.base_scheduler import BaseScheduler +@pytest.mark.fast def test_base_scheduler_is_abstract(): - """ - Test that BaseScheduler cannot be instantiated directly. - """ with pytest.raises(TypeError): BaseScheduler() diff --git a/tests/schedulers/test_linear_scheduler.py b/tests/schedulers/test_linear_scheduler.py index f71467e..5e2167a 100644 --- a/tests/schedulers/test_linear_scheduler.py +++ b/tests/schedulers/test_linear_scheduler.py @@ -4,13 +4,13 @@ import mlx.core as mx import numpy as np import pytest -from mflux.config.config import Config -from mflux.config.model_config import ModelConfig -from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.config.config import Config +from mflux.models.common.config.model_config import ModelConfig from mflux.models.common.schedulers import try_import_external_scheduler from mflux.models.common.schedulers.linear_scheduler import LinearScheduler +@pytest.mark.fast def test_linear_scheduler_import_by_name(): assert ( try_import_external_scheduler("mflux.models.common.schedulers.linear_scheduler.LinearScheduler") @@ -19,34 +19,27 @@ def test_linear_scheduler_import_by_name(): @pytest.fixture -def test_runtime_config(): - return RuntimeConfig( - Config( - # these are the only attributes relevant to schedulers - num_inference_steps=14, - width=1024, - height=1024, - scheduler="linear", - ), - ModelConfig.dev(), # requires_sigma_shift=True +def test_config(): + return Config( + model_config=ModelConfig.dev(), # requires_sigma_shift=True + num_inference_steps=14, + width=1024, + height=1024, + scheduler="linear", ) -def test_linear_scheduler_initialization(test_runtime_config): - """ - Test the initialization of the LinearScheduler. - """ - scheduler = LinearScheduler(runtime_config=test_runtime_config) +@pytest.mark.fast +def test_linear_scheduler_initialization(test_config): + scheduler = LinearScheduler(config=test_config) assert scheduler.sigmas is not None assert isinstance(scheduler.sigmas, mx.array) assert len(scheduler.sigmas) > 0 -def test_linear_scheduler_sigmas_property_no_shift(test_runtime_config): - """ - Test the sigmas property of the LinearScheduler without sigma shift. - """ - scheduler = LinearScheduler(runtime_config=test_runtime_config) +@pytest.mark.fast +def test_linear_scheduler_sigmas_property_no_shift(test_config): + scheduler = LinearScheduler(config=test_config) expected_sigmas_from_mflux_0_9_0 = mx.array( np.load( io.BytesIO( @@ -59,15 +52,13 @@ def test_linear_scheduler_sigmas_property_no_shift(test_runtime_config): ) assert mx.allclose(scheduler.sigmas, expected_sigmas_from_mflux_0_9_0) - assert scheduler.sigmas.shape == (test_runtime_config.num_inference_steps + 1,) + assert scheduler.sigmas.shape == (test_config.num_inference_steps + 1,) -def test_linear_scheduler_sigmas_property_with_shift(test_runtime_config): - """ - Test the sigmas property of the LinearScheduler with sigma shift. - """ - test_runtime_config.model_config = ModelConfig.schnell() # requires_sigma_shift=True - scheduler = LinearScheduler(runtime_config=test_runtime_config) +@pytest.mark.fast +def test_linear_scheduler_sigmas_property_with_shift(test_config): + test_config.model_config = ModelConfig.schnell() # requires_sigma_shift=True + scheduler = LinearScheduler(config=test_config) expected_sigmas_from_mflux_0_9_0 = mx.array( np.load( io.BytesIO( @@ -79,4 +70,4 @@ def test_linear_scheduler_sigmas_property_with_shift(test_runtime_config): ) ) assert mx.allclose(scheduler.sigmas, expected_sigmas_from_mflux_0_9_0) - assert scheduler.sigmas.shape == (test_runtime_config.num_inference_steps + 1,) + assert scheduler.sigmas.shape == (test_config.num_inference_steps + 1,) diff --git a/tests/schedulers/test_scheduler_lookup.py b/tests/schedulers/test_scheduler_lookup.py index b369ac2..9372d32 100644 --- a/tests/schedulers/test_scheduler_lookup.py +++ b/tests/schedulers/test_scheduler_lookup.py @@ -3,15 +3,18 @@ import pytest import mflux.models.common.schedulers as schedulers +@pytest.mark.fast def test_scheduler_by_path(): schedulers.try_import_external_scheduler("mflux.models.common.schedulers.linear_scheduler.LinearScheduler") +@pytest.mark.fast def test_scheduler_bad_module(): with pytest.raises(schedulers.SchedulerModuleNotFound): schedulers.try_import_external_scheduler("someone.other.project.BarScheduler") +@pytest.mark.fast def test_scheduler_bad_classname(): with pytest.raises(schedulers.SchedulerClassNotFound): schedulers.try_import_external_scheduler("mflux.models.common.schedulers.linear_scheduler.FooBarScheduler") diff --git a/tests/ui/test_scale_factor.py b/tests/ui/test_scale_factor.py deleted file mode 100644 index 0964671..0000000 --- a/tests/ui/test_scale_factor.py +++ /dev/null @@ -1,119 +0,0 @@ -import pytest - -from mflux.ui.scale_factor import ScaleFactor, parse_scale_factor - - -def test_scale_factor_init(): - """Test ScaleFactor initialization and validation""" - # Valid integer scale factor - sf = ScaleFactor(value=2) - assert sf.value == 2 - - # Valid float scale factor - sf = ScaleFactor(value=1.5) - assert sf.value == 1.5 - - # Zero should raise ValueError - with pytest.raises(ValueError, match="Scale factor must be positive"): - ScaleFactor(value=0) - - # Negative should raise ValueError - with pytest.raises(ValueError, match="Scale factor must be positive"): - ScaleFactor(value=-1) - - -def test_scale_factor_get_scaled_value(): - """Test ScaleFactor.get_scaled_value method""" - # Test with non-perfect multiple (needs rounding down) - sf = ScaleFactor(value=1.5) - assert sf.get_scaled_value(100) == 144 # 1.5 * 100 - (1.5 * 100) % 16 = 150 - 6 = 144 - - # Test with scale factor that creates remainder - sf = ScaleFactor(value=1.2) - assert sf.get_scaled_value(100) == 112 # 1.2 * 100 - (1.2 * 100) % 16 = 120 - 8 = 112 - - # Test with larger remainder - sf = ScaleFactor(value=1.1) - assert sf.get_scaled_value(200) == 208 # 1.1 * 200 - (1.1 * 200) % 16 = 220 - 12 = 208 - - # Test with custom pixel_steps - sf = ScaleFactor(value=1.3) - assert sf.get_scaled_value(100, pixel_steps=32) == 128 # 1.3 * 100 - (1.3 * 100) % 32 = 130 - 2 = 128 - - # Test edge case where result would be less than pixel_steps - sf = ScaleFactor(value=0.1) - assert sf.get_scaled_value(100) == 0 # 0.1 * 100 - (0.1 * 100) % 16 = 10 - 10 = 0 - - -def test_parse_scale_factor_valid(): - """Test parsing valid scale factor strings""" - # Integer scale factors - sf = parse_scale_factor("1x") - assert sf.value == 1 - - sf = parse_scale_factor("2x") - assert sf.value == 2 - - sf = parse_scale_factor("10x") - assert sf.value == 10 - - # Float scale factors - sf = parse_scale_factor("1.5x") - assert sf.value == 1.5 - - sf = parse_scale_factor("2.75X") - assert sf.value == 2.75 - - # With whitespace - sf = parse_scale_factor(" 2x ") - assert sf.value == 2 - - -def test_parse_scale_factor_invalid(): - """Test parsing invalid scale factor strings""" - # Missing 'x' - with pytest.raises(ValueError, match="Invalid scale factor format"): - parse_scale_factor("2") - - # Multiple 'x' - with pytest.raises(ValueError, match="Invalid scale factor format"): - parse_scale_factor("2xx") - - # Non-numeric value - with pytest.raises(ValueError, match="Invalid scale factor format"): - parse_scale_factor("abcx") - - # Empty before 'x' - with pytest.raises(ValueError, match="Invalid scale factor format"): - parse_scale_factor("x") - - # Invalid format - with pytest.raises(ValueError, match="Invalid scale factor format"): - parse_scale_factor("2.5.5x") - - # Negative values should fail at parsing - with pytest.raises(ValueError, match="Invalid scale factor format"): - parse_scale_factor("-1x") - - # Zero should parse but fail in ScaleFactor init - with pytest.raises(ValueError, match="Scale factor must be positive"): - parse_scale_factor("0x") - - -def test_scale_factor_realistic_dimensions(): - """Test scale factor with realistic image dimensions""" - # 2x upscale of 512x512 image - sf = ScaleFactor(value=2) - assert sf.get_scaled_value(512) == 1024 # 2 * 512 - (2 * 512) % 16 = 1024 - 0 = 1024 - - # 1.5x upscale of 768x768 image - sf = ScaleFactor(value=1.5) - assert sf.get_scaled_value(768) == 1152 # 1.5 * 768 - (1.5 * 768) % 16 = 1152 - 0 = 1152 - - # 3x upscale of 256x256 image - sf = ScaleFactor(value=3) - assert sf.get_scaled_value(256) == 768 # 3 * 256 - (3 * 256) % 16 = 768 - 0 = 768 - - # 0.5x downscale of 1024x1024 image - sf = ScaleFactor(value=0.5) - assert sf.get_scaled_value(1024) == 512 # 0.5 * 1024 - (0.5 * 1024) % 16 = 512 - 0 = 512 diff --git a/tests/vlm/test_fibo_vlm.py b/tests/vlm/test_fibo_vlm.py index 9f6e8c4..e3d6755 100644 --- a/tests/vlm/test_fibo_vlm.py +++ b/tests/vlm/test_fibo_vlm.py @@ -9,68 +9,78 @@ from tests.image_generation.test_generate_image_fibo import OWL_PROMPT, OWL_PROM INSPIRE_PROMPT = """ { - "short_description": "A charming, stylized illustration of a young owl perched on a mossy tree stump in a dimly lit forest. The owl has large, expressive eyes and soft, feathered textures. The background is a blur of dark trees and foliage, creating a serene and mysterious atmosphere. The overall impression is one of innocence and nature's quiet beauty.", + "short_description": "A charming, stylized illustration of a young owl sitting on a mossy forest floor. The owl is facing forward with large, expressive eyes and fluffy ear tufts. It is surrounded by a softly blurred forest environment with tall trees and dappled moonlight. The overall impression is one of innocence and nocturnal wonder, rendered in a gentle, painterly style.", "objects": [ { - "description": "A young owl with large, round, golden eyes and a small yellow beak. Its plumage is a mix of soft browns, creams, and subtle blues, with distinct feather patterns. It has prominent ear tufts and a fluffy chest.", + "description": "A young, anthropomorphic owl with large, round eyes and prominent ear tufts. Its plumage is a soft, mottled grey and beige, with distinct feather patterns.", "location": "center", - "relationship": "perched on a tree stump", - "relative_size": "medium within frame", - "shape_and_color": "Rounded body shape, predominantly brown and cream with blueish undertones. Large, round eyes with dark pupils.", - "texture": "soft, feathery, detailed", - "appearance_details": "Its eyes have a glossy sheen and appear to be looking directly at the viewer. The feathers have a layered, textured look.", + "relationship": "The owl is the sole prominent subject, resting on the ground.", + "relative_size": "medium-to-large within frame", + "shape_and_color": "Round body, large circular eyes, conical beak. Predominantly grey and beige.", + "texture": "soft, feathery", + "appearance_details": "Large, dark pupils in its eyes, a small yellow beak, and delicate talons.", "number_of_objects": 1, - "pose": "Sitting upright, with its body facing forward and its head slightly tilted.", + "pose": "Sitting upright, facing forward.", "expression": "curious and gentle", - "action": "perching", - "gender": "unknown", + "action": "resting", + "gender": "neutral", + "skin_tone_and_texture": "N/A (owl)", "orientation": "upright" }, { - "description": "A weathered tree stump covered in lush green moss. It provides a natural perch for the owl.", - "location": "bottom-center foreground", - "relationship": "supports the owl", + "description": "A patch of vibrant green moss covering the ground in the foreground and around the owl.", + "location": "bottom foreground and midground", + "relationship": "The owl is perched on this mossy ground.", "relative_size": "medium", - "shape_and_color": "Irregular, rounded shape, dark brown wood with vibrant green moss.", - "texture": "rough wood, soft mossy texture", - "appearance_details": "The moss is thick and detailed, with small blades of grass growing around the base.", - "number_of_objects": 1, - "orientation": "lying on its side" + "shape_and_color": "Irregular, organic shapes, vibrant green.", + "texture": "soft, velvety, slightly uneven", + "appearance_details": "Appears lush and damp, with small variations in shade.", + "orientation": "horizontal" }, { - "description": "Several slender tree trunks and branches, rendered in dark, muted tones.", - "location": "background and midground", - "relationship": "forms the forest environment", - "relative_size": "large, forming the environment", - "shape_and_color": "Vertical, elongated shapes in shades of dark brown and grey.", - "texture": "smooth, bark-like texture", - "appearance_details": "Some branches have sparse green leaves. The trees are out of focus, creating depth.", + "description": "Several tall, slender tree trunks forming the background.", + "location": "background", + "relationship": "They create the forest environment behind the owl.", + "relative_size": "large", + "shape_and_color": "Vertical, cylindrical shapes. Dark brown and grey.", + "texture": "rough bark", + "appearance_details": "Some trunks are covered in moss. They are out of focus.", "orientation": "vertical" + }, + { + "description": "A few delicate, glowing orbs scattered in the background, resembling fireflies or distant lights.", + "location": "background, scattered", + "relationship": "They add a magical element to the forest scene.", + "relative_size": "small", + "shape_and_color": "Small, spherical, glowing yellow-white.", + "texture": "N/A", + "appearance_details": "Softly blurred, suggesting distance.", + "orientation": "N/A" } ], - "background_setting": "A dense, dark forest with tall trees and undergrowth. The atmosphere is slightly misty or foggy, softening the distant elements.", + "background_setting": "A mystical forest at night, with tall, ancient trees, a soft undergrowth of moss, and a hint of moonlight filtering through the canopy. The atmosphere is serene and enchanting.", "lighting": { - "conditions": "dim, atmospheric lighting", - "direction": "soft, diffused light from the front-left", - "shadows": "soft, subtle shadows that enhance the depth and mood" + "conditions": "soft, ambient moonlight", + "direction": "diffused from above and behind", + "shadows": "soft, elongated shadows cast by the trees, minimal on the owl" }, "aesthetics": { - "composition": "centered composition with the owl as the primary focal point", - "color_scheme": "muted earth tones with pops of green and soft blues/greys", - "mood_atmosphere": "serene, mysterious, innocent", - "aesthetic_score": "high", - "preference_score": "high" + "composition": "centered framing with the owl as the main focal point", + "color_scheme": "cool, muted tones of blues, greys, and greens, with warm accents from the owl's eyes and beak", + "mood_atmosphere": "magical, serene, innocent", + "aesthetic_score": "very high", + "preference_score": "very high" }, "photographic_characteristics": { - "depth_of_field": "shallow, with a blurred background", + "depth_of_field": "shallow, with a strong bokeh effect in the background", "focus": "sharp focus on the owl", "camera_angle": "eye-level", "lens_focal_length": "standard lens (e.g., 50mm)" }, "style_medium": "digital illustration", "text_render": [], - "context": "This image is suitable for children's book illustrations, nature-themed art, or decorative prints.", - "artistic_style": "stylized, painterly, cute" + "context": "This image is suitable for children's book illustrations, fantasy-themed art, or decorative prints.", + "artistic_style": "stylized, painterly, whimsical" } """ @@ -79,44 +89,43 @@ SKYSCRAPERS_INSPIRE_PROMPT = """ "short_description": "A dramatic, low-angle shot of several modern skyscrapers in black and white, emphasizing their towering height and geometric forms against a bright, overcast sky. The buildings feature repetitive window patterns and varying architectural details, creating a sense of urban grandeur and scale.", "objects": [ { - "description": "A tall skyscraper with a facade of numerous rectangular windows, arranged in a grid pattern. The building has a slightly curved or angled design, giving it a dynamic appearance.", + "description": "A tall skyscraper with a facade of numerous rectangular windows, arranged in a grid pattern. The building has a slightly curved or angled top section.", "location": "center-right foreground", - "relationship": "This is the most prominent building, dominating the right side of the frame and serving as a primary focal point.", + "relationship": "It is the most prominent building, dominating the right side of the frame and appearing to recede into the background.", "relative_size": "large within frame", - "shape_and_color": "Rectangular and angular forms, dark grey to black with bright white window reflections.", - "texture": "Smooth concrete or glass, with visible mullions creating a textured grid.", - "appearance_details": "The windows reflect the bright sky, appearing as glowing rectangles. The building's structure is robust and imposing.", + "shape_and_color": "Rectangular with a complex, multi-faceted top; shades of gray and black.", + "texture": "Smooth concrete and glass, with visible window frames.", + "appearance_details": "The windows reflect the bright sky, appearing light gray or white. Some sections of the facade have decorative elements or protrusions.", "orientation": "vertical" }, { - "description": "A large, flat-roofed skyscraper with a simpler, more rectilinear design compared to the central-right building. Its facade also features windows, but they are less prominent due to the building's angle and distance.", - "location": "left foreground", - "relationship": "It stands to the left of the central skyscraper, partially obscuring other buildings behind it and contributing to the layered effect.", - "relative_size": "large within frame", - "shape_and_color": "Large, flat, angular forms, dark grey to black.", - "texture": "Rough concrete texture on the upper sections, smoother on the windowed parts.", - "appearance_details": "The building has a distinct overhang or cantilevered section at the top-left, adding architectural interest.", - "orientation": "diagonal, leaning towards the right" - }, - { - "description": "A cluster of mid-rise buildings, less distinct than the foreground structures, with visible window patterns.", - "location": "midground, behind the foreground buildings", - "relationship": "These buildings provide depth and context, suggesting a dense urban environment behind the prominent foreground structures.", + "description": "A long, rectangular skyscraper with a facade of regularly spaced, horizontal windows.", + "location": "left midground", + "relationship": "It stands to the left of the central skyscraper, appearing slightly further back and narrower.", "relative_size": "medium", - "shape_and_color": "Rectangular forms, dark grey with hints of white from windows.", - "texture": "Less defined, appearing smoother due to distance.", - "appearance_details": "Their details are softened by distance and atmospheric haze.", - "number_of_objects": 3, + "shape_and_color": "Tall, slender rectangle; shades of gray and black.", + "texture": "Smooth concrete and glass.", + "appearance_details": "The horizontal lines of windows create a strong sense of rhythm. Its top is flat and extends horizontally.", "orientation": "vertical" }, { - "description": "A tall, slender skyscraper with a facade composed of many small, closely spaced rectangular windows, creating a dense, textured pattern.", - "location": "bottom-right midground", - "relationship": "It stands slightly behind and to the right of the main central skyscraper, adding to the overall height of the urban landscape.", + "description": "A skyscraper with a facade featuring larger, rectangular windows that are not perfectly uniform, giving a slightly more textured appearance.", + "location": "center-left midground", + "relationship": "It is positioned behind the left midground skyscraper and to the left of the central skyscraper, partially obscured by the foreground elements.", "relative_size": "medium", - "shape_and_color": "Tall, slender, rectangular, dark grey with bright white window reflections.", - "texture": "Densely packed, creating a fine-grained texture from the windows.", - "appearance_details": "Its vertical lines are emphasized by the window arrangement.", + "shape_and_color": "Tall, slender rectangle; shades of gray and black.", + "texture": "Smooth concrete and glass, with subtle variations in window size.", + "appearance_details": "Its upper section has a distinct, stepped design.", + "orientation": "vertical" + }, + { + "description": "A building with a facade characterized by smaller, square-like windows, creating a more intricate pattern.", + "location": "bottom-right foreground", + "relationship": "It is partially visible at the bottom right, appearing to be closer to the viewer than the central skyscraper.", + "relative_size": "small", + "shape_and_color": "Rectangular base with a repeating pattern of small squares; shades of gray and black.", + "texture": "Smooth concrete and glass.", + "appearance_details": "The windows are arranged in a dense, grid-like pattern, almost appearing pixelated.", "orientation": "vertical" } ], @@ -124,34 +133,35 @@ SKYSCRAPERS_INSPIRE_PROMPT = """ "lighting": { "conditions": "overcast daylight", "direction": "diffused from above", - "shadows": "soft, subtle shadows on the buildings, mainly emphasizing their geometric forms rather than sharp lines" + "shadows": "soft, subtle shadows on the buildings, emphasizing their forms without harsh lines" }, "aesthetics": { - "composition": "dynamic and asymmetrical, with buildings angled and overlapping, creating a sense of depth and scale. The low angle emphasizes the height of the structures.", - "color_scheme": "monochromatic (black and white) with a wide range of greys, emphasizing form and texture.", - "mood_atmosphere": "grand, imposing, architectural, and slightly dramatic.", + "composition": "dynamic low-angle composition, with leading lines from the buildings converging towards the top-center, creating a sense of immense height and perspective.", + "color_scheme": "monochromatic (black and white) with a wide range of grays, emphasizing texture and form.", + "mood_atmosphere": "grand, imposing, architectural, and slightly mysterious.", "aesthetic_score": "very high", "preference_score": "very high" }, "photographic_characteristics": { "depth_of_field": "deep", - "focus": "sharp focus on the foreground and midground buildings, with slight softening towards the background.", + "focus": "sharp focus on all visible skyscrapers, maintaining clarity across the architectural details.", "camera_angle": "very low angle, looking upwards from ground level.", - "lens_focal_length": "wide-angle" + "lens_focal_length": "wide-angle lens, to capture the full height and breadth of the buildings." }, "style_medium": "photograph", "text_render": [], - "context": "An architectural photograph, likely intended for art prints, urban exploration photography, or a design magazine, focusing on the abstract beauty of modern cityscapes.", - "artistic_style": "minimalist, high-contrast, geometric" + "context": "An architectural photograph, likely intended for art prints, urban exploration blogs, or promotional material for city tourism, focusing on the abstract beauty of modern cityscapes.", + "artistic_style": "minimalist, abstract, architectural" } """ @pytest.fixture def vlm(): - return FiboVLM(model_id="briaai/FIBO-vlm") + return FiboVLM(model_id="briaai/FIBO-vlm", quantize=8) +@pytest.mark.slow def test_vlm_generate_json(vlm): # given input_prompt = "A hyper-detailed, ultra-fluffy owl sitting in the trees at night, looking directly at the camera with wide, adorable, expressive eyes. Its feathers are soft and voluminous, catching the cool moonlight with subtle silver highlights. The owl's gaze is curious and full of charm, giving it a whimsical, storybook-like personality." @@ -164,6 +174,7 @@ def test_vlm_generate_json(vlm): assert isinstance(json.loads(json_output), dict), "Output should be a valid JSON object" +@pytest.mark.slow def test_vlm_refine_json_text_only(vlm): # given editing_instructions = "make the owl white color but keep everything else exactly the same" @@ -180,6 +191,7 @@ def test_vlm_refine_json_text_only(vlm): assert isinstance(json.loads(json_output), dict), "Output should be a valid JSON object" +@pytest.mark.slow def test_vlm_inspire_json_from_image(vlm): # given resource_dir = Path(__file__).parent.parent / "resources" @@ -194,6 +206,7 @@ def test_vlm_inspire_json_from_image(vlm): assert isinstance(json.loads(json_output), dict), "Output should be a valid JSON object" +@pytest.mark.slow def test_vlm_inspire_json_from_image_with_prompt(vlm): # given resource_dir = Path(__file__).parent.parent / "resources" diff --git a/tests/weights/test_lora_library.py b/tests/weights/test_lora_library.py deleted file mode 100644 index d78ded5..0000000 --- a/tests/weights/test_lora_library.py +++ /dev/null @@ -1,199 +0,0 @@ -import os -import tempfile -from pathlib import Path -from unittest import mock - -import pytest - -from mflux.models.common.lora.download import lora_library - - -@pytest.fixture(autouse=True) -def reset_lora_registry(): - """Reset the global registry before and after each test.""" - original = lora_library._LORA_REGISTRY.copy() - lora_library._LORA_REGISTRY.clear() - yield - lora_library._LORA_REGISTRY.clear() - lora_library._LORA_REGISTRY.update(original) - - -@pytest.fixture -def temp_lora_dirs(): - """Create temporary directories with fake .safetensors files for testing.""" - with tempfile.TemporaryDirectory() as temp_root: - # Create two library directories - lib1 = Path(temp_root) / "library1" - lib2 = Path(temp_root) / "library2" - lib1.mkdir() - lib2.mkdir() - - # Create some .safetensors files - (lib1 / "style_a.safetensors").touch() - (lib1 / "style_b.safetensors").touch() - (lib1 / "subdir").mkdir() - (lib1 / "subdir" / "style_c.safetensors").touch() - - # Create transformer directory with digit-named files (should be ignored) - (lib1 / "transformer").mkdir() - (lib1 / "transformer" / "0.safetensors").touch() - (lib1 / "transformer" / "1.safetensors").touch() - (lib1 / "transformer" / "style_valid.safetensors").touch() # Should not be ignored - - (lib2 / "style_b.safetensors").touch() # Duplicate name - (lib2 / "style_d.safetensors").touch() - - yield lib1, lib2 - - -def test_discover_lora_files_single_path(temp_lora_dirs): - """Test discovering files in a single library path.""" - lib1, lib2 = temp_lora_dirs - - result = lora_library._discover_lora_files([lib1]) - - assert len(result) == 4 # Should not include digit-named files in transformer/ - assert "style_a" in result - assert "style_b" in result - assert "style_c" in result - assert "style_valid" in result - assert "0" not in result # Digit-named files in transformer/ should be excluded - assert "1" not in result - assert result["style_a"] == (lib1 / "style_a.safetensors").resolve() - assert result["style_c"] == (lib1 / "subdir" / "style_c.safetensors").resolve() - assert result["style_valid"] == (lib1 / "transformer" / "style_valid.safetensors").resolve() - - -def test_discover_lora_files_multiple_paths_precedence(temp_lora_dirs): - """Test that earlier paths take precedence for duplicate names.""" - lib1, lib2 = temp_lora_dirs - - # lib1 should take precedence over lib2 - result = lora_library._discover_lora_files([lib1, lib2]) - - assert len(result) == 5 # style_a, style_b, style_c, style_valid, style_d - assert "style_a" in result - assert "style_b" in result - assert "style_c" in result - assert "style_valid" in result - assert "style_d" in result - assert "0" not in result # Digit-named files should still be excluded - assert "1" not in result - - # style_b should come from lib1 (first in list) - assert result["style_b"] == (lib1 / "style_b.safetensors").resolve() - assert result["style_d"] == (lib2 / "style_d.safetensors").resolve() - - -def test_discover_lora_files_nonexistent_path(): - """Test handling of nonexistent paths.""" - result = lora_library._discover_lora_files([Path("/nonexistent/path")]) - assert result == {} - - -def test_get_lora_path_existing_file(temp_lora_dirs): - """Test that existing file paths are returned as-is.""" - lib1, _ = temp_lora_dirs - existing_file = lib1 / "style_a.safetensors" - - result = lora_library.get_lora_path(str(existing_file)) - assert result == str(existing_file) - - -def test_get_lora_path_from_registry(temp_lora_dirs): - """Test resolving paths from the registry.""" - lib1, lib2 = temp_lora_dirs - - # Mock the registry - with mock.patch.object( - lora_library, - "_LORA_REGISTRY", - { - "style_a": lib1 / "style_a.safetensors", - "style_b": lib1 / "style_b.safetensors", - }, - ): - # Should resolve from registry - assert lora_library.get_lora_path("style_a") == str(lib1 / "style_a.safetensors") - assert lora_library.get_lora_path("style_b") == str(lib1 / "style_b.safetensors") - - # Should raise FileNotFoundError if not in registry - with pytest.raises(FileNotFoundError, match="LoRA file not found: 'unknown_style'"): - lora_library.get_lora_path("unknown_style") - - -def test_initialize_registry_from_env(temp_lora_dirs): - """Test that registry is initialized from LORA_LIBRARY_PATH env var.""" - lib1, lib2 = temp_lora_dirs - - # Set up environment variable with colon-delimited paths - env_path = f"{lib1}:{lib2}" - - with mock.patch.dict(os.environ, {"LORA_LIBRARY_PATH": env_path}): - # Re-initialize the registry - lora_library._initialize_registry() - registry = lora_library.get_registry() - - assert len(registry) == 5 # Includes style_valid from transformer/ - assert "style_a" in registry - assert "style_b" in registry - assert "style_c" in registry - assert "style_valid" in registry - assert "style_d" in registry - assert "0" not in registry # Digit files should be excluded - assert "1" not in registry - - # Check precedence - style_b should be from lib1 - assert registry["style_b"] == (lib1 / "style_b.safetensors").resolve() - - -def test_initialize_registry_empty_env(): - """Test that registry is empty when env var is not set.""" - # Save the original registry - original_registry = lora_library._LORA_REGISTRY.copy() - try: - with mock.patch.dict(os.environ, {}, clear=True): - lora_library._initialize_registry() - registry = lora_library.get_registry() - assert registry == {} - finally: - # Restore the original registry - lora_library._LORA_REGISTRY = original_registry - - -def test_get_registry_returns_copy(): - """Test that get_registry returns a copy, not the original.""" - with mock.patch.object(lora_library, "_LORA_REGISTRY", {"test": Path("/test")}): - registry1 = lora_library.get_registry() - registry2 = lora_library.get_registry() - - # Should be equal but different objects - assert registry1 == registry2 - assert registry1 is not registry2 - - # Modifying the copy shouldn't affect the original - registry1["new"] = Path("/new") - assert "new" not in lora_library._LORA_REGISTRY - - -def test_transformer_digit_filtering(temp_lora_dirs): - """Test that digit-named files in transformer directories are filtered out.""" - lib1, _ = temp_lora_dirs - - # Create additional test cases - (lib1 / "9.safetensors").touch() # Should be included (not in transformer/) - (lib1 / "other_dir").mkdir() - (lib1 / "other_dir" / "5.safetensors").touch() # Should be included (not in transformer/) - - result = lora_library._discover_lora_files([lib1]) - - # Digit files NOT in transformer/ should be included - assert "9" in result - assert "5" in result - - # Digit files in transformer/ should be excluded - assert "0" not in result - assert "1" not in result - - # Non-digit files in transformer/ should be included - assert "style_valid" in result diff --git a/tools/create_outpaint_image_canvas_and_mask.py b/tools/create_outpaint_image_canvas_and_mask.py index b04ffa9..b60a30a 100755 --- a/tools/create_outpaint_image_canvas_and_mask.py +++ b/tools/create_outpaint_image_canvas_and_mask.py @@ -1,8 +1,8 @@ import sys from pathlib import Path -from mflux.ui.box_values import AbsoluteBoxValues, BoxValues -from mflux.ui.cli.parsers import CommandLineParser +from mflux.cli.parser.parsers import CommandLineParser +from mflux.utils.box_values import AbsoluteBoxValues, BoxValues from mflux.utils.image_util import ImageUtil diff --git a/uv.lock b/uv.lock index bc704c2..b176b1f 100644 --- a/uv.lock +++ b/uv.lock @@ -51,8 +51,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", version = "2.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "torch" }, ] sdist = { url = "https://files.pythonhosted.org/packages/f7/66/be171836d86dc5b8698b3a9bf4b9eb10cb53369729939f88bf650167588b/accelerate-1.10.0.tar.gz", hash = "sha256:8270568fda9036b5cccdc09703fef47872abccd56eb5f6d53b54ea5fb7581496", size = 392261 } wheels = [ @@ -769,7 +768,7 @@ wheels = [ [[package]] name = "mflux" -version = "0.12.1" +version = "0.13.0.dev0" source = { editable = "." } dependencies = [ { name = "accelerate" }, @@ -791,8 +790,7 @@ dependencies = [ { name = "sentencepiece", version = "0.2.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, { name = "tokenizers", marker = "python_full_version >= '3.13'" }, { name = "toml" }, - { name = "torch", version = "2.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "torch" }, { name = "tqdm" }, { name = "transformers" }, { name = "twine" }, @@ -1095,371 +1093,114 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/48/6b/1c6b515a83d5564b1698a61efa245727c8feecf308f4091f565988519d20/numpy-2.3.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:e610832418a2bc09d974cc9fecebfa51e9532d6190223bc5ef6a7402ebf3b5cb", size = 12927246 }, ] -[[package]] -name = "nvidia-cublas-cu12" -version = "12.6.4.1" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/af/eb/ff4b8c503fa1f1796679dce648854d58751982426e4e4b37d6fce49d259c/nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08ed2686e9875d01b58e3cb379c6896df8e76c75e0d4a7f7dace3d7b6d9ef8eb", size = 393138322 }, - { url = "https://files.pythonhosted.org/packages/97/0d/f1f0cadbf69d5b9ef2e4f744c9466cb0a850741d08350736dfdb4aa89569/nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:235f728d6e2a409eddf1df58d5b0921cf80cfa9e72b9f2775ccb7b4a87984668", size = 390794615 }, -] - [[package]] name = "nvidia-cublas-cu12" version = "12.8.4.1" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/29/99/db44d685f0e257ff0e213ade1964fc459b4a690a73293220e98feb3307cf/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b86f6dd8935884615a0683b663891d43781b819ac4f2ba2b0c9604676af346d0", size = 590537124 }, { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921 }, ] -[[package]] -name = "nvidia-cuda-cupti-cu12" -version = "12.6.80" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/8b/2f6230cb715646c3a9425636e513227ce5c93c4d65823a734f4bb86d43c3/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:166ee35a3ff1587f2490364f90eeeb8da06cd867bd5b701bf7f9a02b78bc63fc", size = 8236764 }, - { url = "https://files.pythonhosted.org/packages/25/0f/acb326ac8fd26e13c799e0b4f3b2751543e1834f04d62e729485872198d4/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_aarch64.whl", hash = "sha256:358b4a1d35370353d52e12f0a7d1769fc01ff74a191689d3870b2123156184c4", size = 8236756 }, - { url = "https://files.pythonhosted.org/packages/49/60/7b6497946d74bcf1de852a21824d63baad12cd417db4195fc1bfe59db953/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6768bad6cab4f19e8292125e5f1ac8aa7d1718704012a0e3272a6f61c4bce132", size = 8917980 }, - { url = "https://files.pythonhosted.org/packages/a5/24/120ee57b218d9952c379d1e026c4479c9ece9997a4fb46303611ee48f038/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73", size = 8917972 }, -] - [[package]] name = "nvidia-cuda-cupti-cu12" version = "12.8.90" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/d5/1f/b3bd73445e5cb342727fd24fe1f7b748f690b460acadc27ea22f904502c8/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4412396548808ddfed3f17a467b104ba7751e6b58678a4b840675c56d21cf7ed", size = 9533318 }, { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621 }, ] -[[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.6.77" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/f4/2f/72df534873235983cc0a5371c3661bebef7c4682760c275590b972c7b0f9/nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5847f1d6e5b757f1d2b3991a01082a44aad6f10ab3c5c0213fa3e25bddc25a13", size = 23162955 }, - { url = "https://files.pythonhosted.org/packages/75/2e/46030320b5a80661e88039f59060d1790298b4718944a65a7f2aeda3d9e9/nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:35b0cc6ee3a9636d5409133e79273ce1f3fd087abb0532d2d2e8fff1fe9efc53", size = 23650380 }, -] - [[package]] name = "nvidia-cuda-nvrtc-cu12" version = "12.8.93" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029 }, { url = "https://files.pythonhosted.org/packages/eb/d1/e50d0acaab360482034b84b6e27ee83c6738f7d32182b987f9c7a4e32962/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc1fec1e1637854b4c0a65fb9a8346b51dd9ee69e61ebaccc82058441f15bce8", size = 43106076 }, ] -[[package]] -name = "nvidia-cuda-runtime-cu12" -version = "12.6.77" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/ea/590b2ac00d772a8abd1c387a92b46486d2679ca6622fd25c18ff76265663/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6116fad3e049e04791c0256a9778c16237837c08b27ed8c8401e2e45de8d60cd", size = 908052 }, - { url = "https://files.pythonhosted.org/packages/b7/3d/159023799677126e20c8fd580cca09eeb28d5c5a624adc7f793b9aa8bbfa/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d461264ecb429c84c8879a7153499ddc7b19b5f8d84c204307491989a365588e", size = 908040 }, - { url = "https://files.pythonhosted.org/packages/e1/23/e717c5ac26d26cf39a27fbc076240fad2e3b817e5889d671b67f4f9f49c5/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba3b56a4f896141e25e19ab287cd71e52a6a0f4b29d0d31609f60e3b4d5219b7", size = 897690 }, - { url = "https://files.pythonhosted.org/packages/f0/62/65c05e161eeddbafeca24dc461f47de550d9fa8a7e04eb213e32b55cfd99/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8", size = 897678 }, -] - [[package]] name = "nvidia-cuda-runtime-cu12" version = "12.8.90" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/7c/75/f865a3b236e4647605ea34cc450900854ba123834a5f1598e160b9530c3a/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d", size = 965265 }, { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765 }, ] -[[package]] -name = "nvidia-cudnn-cu12" -version = "9.5.1.17" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/99/93/a201a12d3ec1caa8c6ac34c1c2f9eeb696b886f0c36ff23c638b46603bd0/nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:9fd4584468533c61873e5fda8ca41bac3a38bcb2d12350830c69b0a96a7e4def", size = 570523509 }, - { url = "https://files.pythonhosted.org/packages/2a/78/4535c9c7f859a64781e43c969a3a7e84c54634e319a996d43ef32ce46f83/nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2", size = 570988386 }, -] - [[package]] name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, + { name = "nvidia-cublas-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/fa/41/e79269ce215c857c935fd86bcfe91a451a584dfc27f1e068f568b9ad1ab7/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c9132cc3f8958447b4910a1720036d9eff5928cc3179b0a51fb6d167c6cc87d8", size = 705026878 }, { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467 }, ] -[[package]] -name = "nvidia-cufft-cu12" -version = "11.3.0.4" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/37/c50d2b2f2c07e146776389e3080f4faf70bcc4fa6e19d65bb54ca174ebc3/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d16079550df460376455cba121db6564089176d9bac9e4f360493ca4741b22a6", size = 200164144 }, - { url = "https://files.pythonhosted.org/packages/ce/f5/188566814b7339e893f8d210d3a5332352b1409815908dad6a363dcceac1/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8510990de9f96c803a051822618d42bf6cb8f069ff3f48d93a8486efdacb48fb", size = 200164135 }, - { url = "https://files.pythonhosted.org/packages/8f/16/73727675941ab8e6ffd86ca3a4b7b47065edcca7a997920b831f8147c99d/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5", size = 200221632 }, - { url = "https://files.pythonhosted.org/packages/60/de/99ec247a07ea40c969d904fc14f3a356b3e2a704121675b75c366b694ee1/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.whl", hash = "sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca", size = 200221622 }, -] - [[package]] name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/60/bc/7771846d3a0272026c416fbb7e5f4c1f146d6d80704534d0b187dd6f4800/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:848ef7224d6305cdb2a4df928759dca7b1201874787083b6e7550dd6765ce69a", size = 193109211 }, { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695 }, ] -[[package]] -name = "nvidia-cufile-cu12" -version = "1.11.1.6" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/b2/66/cc9876340ac68ae71b15c743ddb13f8b30d5244af344ec8322b449e35426/nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc23469d1c7e52ce6c1d55253273d32c565dd22068647f3aa59b3c6b005bf159", size = 1142103 }, - { url = "https://files.pythonhosted.org/packages/17/bf/cc834147263b929229ce4aadd62869f0b195e98569d4c28b23edc72b85d9/nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:8f57a0051dcf2543f6dc2b98a98cb2719c37d3cee1baba8965d57f3bbc90d4db", size = 1066155 }, -] - [[package]] name = "nvidia-cufile-cu12" version = "1.13.1.3" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834 }, { url = "https://files.pythonhosted.org/packages/1e/f5/5607710447a6fe9fd9b3283956fceeee8a06cda1d2f56ce31371f595db2a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:4beb6d4cce47c1a0f1013d72e02b0994730359e17801d395bdcbf20cfb3bb00a", size = 1120705 }, ] -[[package]] -name = "nvidia-curand-cu12" -version = "10.3.7.77" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/42/ac/36543605358a355632f1a6faa3e2d5dfb91eab1e4bc7d552040e0383c335/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:6e82df077060ea28e37f48a3ec442a8f47690c7499bff392a5938614b56c98d8", size = 56289881 }, - { url = "https://files.pythonhosted.org/packages/73/1b/44a01c4e70933637c93e6e1a8063d1e998b50213a6b65ac5a9169c47e98e/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a42cd1344297f70b9e39a1e4f467a4e1c10f1da54ff7a85c12197f6c652c8bdf", size = 56279010 }, - { url = "https://files.pythonhosted.org/packages/4a/aa/2c7ff0b5ee02eaef890c0ce7d4f74bc30901871c5e45dee1ae6d0083cd80/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:99f1a32f1ac2bd134897fc7a203f779303261268a65762a623bf30cc9fe79117", size = 56279000 }, - { url = "https://files.pythonhosted.org/packages/a6/02/5362a9396f23f7de1dd8a64369e87c85ffff8216fc8194ace0fa45ba27a5/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:7b2ed8e95595c3591d984ea3603dd66fe6ce6812b886d59049988a712ed06b6e", size = 56289882 }, -] - [[package]] name = "nvidia-curand-cu12" version = "10.3.9.90" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/45/5e/92aa15eca622a388b80fbf8375d4760738df6285b1e92c43d37390a33a9a/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dfab99248034673b779bc6decafdc3404a8a6f502462201f2f31f11354204acd", size = 63625754 }, { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976 }, ] -[[package]] -name = "nvidia-cusolver-cu12" -version = "11.7.1.2" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, - { name = "nvidia-cusparse-cu12", version = "12.5.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/93/17/dbe1aa865e4fdc7b6d4d0dd308fdd5aaab60f939abfc0ea1954eac4fb113/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ce237ef60acde1efc457335a2ddadfd7610b892d94efee7b776c64bb1cac9e0", size = 157833628 }, - { url = "https://files.pythonhosted.org/packages/f0/6e/c2cf12c9ff8b872e92b4a5740701e51ff17689c4d726fca91875b07f655d/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c", size = 158229790 }, - { url = "https://files.pythonhosted.org/packages/9f/81/baba53585da791d043c10084cf9553e074548408e04ae884cfe9193bd484/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6", size = 158229780 }, - { url = "https://files.pythonhosted.org/packages/7c/5f/07d0ba3b7f19be5a5ec32a8679fc9384cfd9fc6c869825e93be9f28d6690/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:dbbe4fc38ec1289c7e5230e16248365e375c3673c9c8bac5796e2e20db07f56e", size = 157833630 }, -] - [[package]] name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] dependencies = [ - { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, - { name = "nvidia-cusparse-cu12", version = "12.5.8.93", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, + { name = "nvidia-cublas-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, + { name = "nvidia-cusparse-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/32/f7cd6ce8a7690544d084ea21c26e910a97e077c9b7f07bf5de623ee19981/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:db9ed69dbef9715071232caa9b69c52ac7de3a95773c2db65bdba85916e4e5c0", size = 267229841 }, { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905 }, ] -[[package]] -name = "nvidia-cusparse-cu12" -version = "12.5.4.2" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/eb/eb/6681efd0aa7df96b4f8067b3ce7246833dd36830bb4cec8896182773db7d/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d25b62fb18751758fe3c93a4a08eff08effedfe4edf1c6bb5afd0890fe88f887", size = 216451147 }, - { url = "https://files.pythonhosted.org/packages/d3/56/3af21e43014eb40134dea004e8d0f1ef19d9596a39e4d497d5a7de01669f/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7aa32fa5470cf754f72d1116c7cbc300b4e638d3ae5304cfa4a638a5b87161b1", size = 216451135 }, - { url = "https://files.pythonhosted.org/packages/06/1e/b8b7c2f4099a37b96af5c9bb158632ea9e5d9d27d7391d7eb8fc45236674/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73", size = 216561367 }, - { url = "https://files.pythonhosted.org/packages/43/ac/64c4316ba163e8217a99680c7605f779accffc6a4bcd0c778c12948d3707/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f", size = 216561357 }, -] - [[package]] name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] dependencies = [ - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, + { name = "nvidia-nvjitlink-cu12", marker = "(platform_machine != 'aarch64' and platform_system != 'Darwin') or (platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/bc/f7/cd777c4109681367721b00a106f491e0d0d15cfa1fd59672ce580ce42a97/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9b6c161cb130be1a07a27ea6923df8141f3c295852f4b260c65f18f3e0a091dc", size = 288117129 }, { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466 }, ] -[[package]] -name = "nvidia-cusparselt-cu12" -version = "0.6.3" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/62/da/4de092c61c6dea1fc9c936e69308a02531d122e12f1f649825934ad651b5/nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8371549623ba601a06322af2133c4a44350575f5a3108fb75f3ef20b822ad5f1", size = 156402859 }, - { url = "https://files.pythonhosted.org/packages/3b/9a/72ef35b399b0e183bc2e8f6f558036922d453c4d8237dab26c666a04244b/nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e5c8a26c36445dd2e6812f1177978a24e2d37cacce7e090f297a688d1ec44f46", size = 156785796 }, -] - [[package]] name = "nvidia-cusparselt-cu12" version = "0.7.1" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/73/b9/598f6ff36faaece4b3c50d26f50e38661499ff34346f00e057760b35cc9d/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:8878dce784d0fac90131b6817b607e803c36e629ba34dc5b433471382196b6a5", size = 283835557 }, { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691 }, @@ -1467,91 +1208,35 @@ wheels = [ [[package]] name = "nvidia-nccl-cu12" -version = "2.26.2" +version = "2.27.5" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ - { url = "https://files.pythonhosted.org/packages/69/5b/ca2f213f637305633814ae8c36b153220e40a07ea001966dcd87391f3acb/nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5c196e95e832ad30fbbb50381eb3cbd1fadd5675e587a548563993609af19522", size = 291671495 }, - { url = "https://files.pythonhosted.org/packages/67/ca/f42388aed0fddd64ade7493dbba36e1f534d4e6fdbdd355c6a90030ae028/nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:694cf3879a206553cc9d7dbda76b13efaf610fdb70a50cba303de1b0d1530ac6", size = 201319755 }, -] - -[[package]] -name = "nvidia-nccl-cu12" -version = "2.27.3" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/4b/7b/8354b784cf73b0ba51e566b4baba3ddd44fe8288a3d39ef1e06cd5417226/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:9ddf1a245abc36c550870f26d537a9b6087fb2e2e3d6e0ef03374c6fd19d984f", size = 322397768 }, - { url = "https://files.pythonhosted.org/packages/5c/5b/4e4fff7bad39adf89f735f2bc87248c81db71205b62bcc0d5ca5b606b3c3/nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adf27ccf4238253e0b826bce3ff5fa532d65fc42322c8bfdfaf28024c0fbe039", size = 322364134 }, -] - -[[package]] -name = "nvidia-nvjitlink-cu12" -version = "12.6.85" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/d7/c5383e47c7e9bf1c99d5bd2a8c935af2b6d705ad831a7ec5c97db4d82f4f/nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:eedc36df9e88b682efe4309aa16b5b4e78c2407eac59e8c10a6a47535164369a", size = 19744971 }, - { url = "https://files.pythonhosted.org/packages/31/db/dc71113d441f208cdfe7ae10d4983884e13f464a6252450693365e166dcf/nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cf4eaa7d4b6b543ffd69d6abfb11efdeb2db48270d94dfd3a452c24150829e41", size = 19270338 }, + { url = "https://files.pythonhosted.org/packages/bb/1c/857979db0ef194ca5e21478a0612bcdbbe59458d7694361882279947b349/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:31432ad4d1fb1004eb0c56203dc9bc2178a1ba69d1d9e02d64a6938ab5e40e7a", size = 322400625 }, + { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229 }, ] [[package]] name = "nvidia-nvjitlink-cu12" version = "12.8.93" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836 }, { url = "https://files.pythonhosted.org/packages/2a/a2/8cee5da30d13430e87bf99bb33455d2724d0a4a9cb5d7926d80ccb96d008/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:adccd7161ace7261e01bb91e44e88da350895c270d23f744f0820c818b7229e7", size = 38386204 }, ] [[package]] -name = "nvidia-nvtx-cu12" -version = "12.6.77" +name = "nvidia-nvshmem-cu12" +version = "3.3.20" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ - { url = "https://files.pythonhosted.org/packages/b9/93/80f8a520375af9d7ee44571a6544653a176e53c2b8ccce85b97b83c2491b/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f44f8d86bb7d5629988d61c8d3ae61dddb2015dee142740536bc7481b022fe4b", size = 90549 }, - { url = "https://files.pythonhosted.org/packages/2b/53/36e2fd6c7068997169b49ffc8c12d5af5e5ff209df6e1a2c4d373b3a638f/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_aarch64.whl", hash = "sha256:adcaabb9d436c9761fca2b13959a2d237c5f9fd406c8e4b723c695409ff88059", size = 90539 }, - { url = "https://files.pythonhosted.org/packages/56/9a/fff8376f8e3d084cd1530e1ef7b879bb7d6d265620c95c1b322725c694f4/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b90bed3df379fa79afbd21be8e04a0314336b8ae16768b58f2d34cb1d04cd7d2", size = 89276 }, - { url = "https://files.pythonhosted.org/packages/9e/4e/0d0c945463719429b7bd21dece907ad0bde437a2ff12b9b12fee94722ab0/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1", size = 89265 }, + { url = "https://files.pythonhosted.org/packages/92/9d/3dd98852568fb845ec1f7902c90a22b240fe1cbabda411ccedf2fd737b7b/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0b0b960da3842212758e4fa4696b94f129090b30e5122fea3c5345916545cff0", size = 124484616 }, + { url = "https://files.pythonhosted.org/packages/3b/6c/99acb2f9eb85c29fc6f3a7ac4dccfd992e22666dd08a642b303311326a97/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d00f26d3f9b2e3c3065be895e3059d6479ea5c638a3f38c9fec49b1b9dd7c1e5", size = 124657145 }, ] [[package]] name = "nvidia-nvtx-cu12" version = "12.8.90" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] wheels = [ { url = "https://files.pythonhosted.org/packages/10/c0/1b303feea90d296f6176f32a2a70b5ef230f9bdeb3a72bddb0dc922dc137/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d7ad891da111ebafbf7e015d34879f7112832fc239ff0d7d776b6cb685274615", size = 91161 }, { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954 }, @@ -2358,144 +2043,62 @@ wheels = [ [[package]] name = "torch" -version = "2.7.1" +version = "2.9.1" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version == '3.12.*' and platform_system == 'Darwin' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'darwin'", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_system == 'Darwin' and sys_platform == 'linux'", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'linux'", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'linux'", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system == 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system == 'Darwin' and sys_platform != 'darwin' and sys_platform != 'linux')", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux'", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "python_full_version == '3.11.*' and platform_system == 'Darwin' and sys_platform == 'darwin'", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'darwin'", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_system == 'Darwin' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'linux'", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system == 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system == 'Darwin' and sys_platform != 'darwin' and sys_platform != 'linux')", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux'", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "python_full_version < '3.11' and platform_system == 'Darwin' and sys_platform == 'darwin'", - "python_full_version < '3.11' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'darwin'", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "python_full_version < '3.11' and platform_machine == 'aarch64' and platform_system == 'Darwin' and sys_platform == 'linux'", - "python_full_version < '3.11' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'linux'", - "python_full_version < '3.11' and platform_machine == 'aarch64' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'linux'", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system == 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system == 'Darwin' and sys_platform != 'darwin' and sys_platform != 'linux')", - "python_full_version < '3.11' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux'", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] dependencies = [ - { name = "filelock", marker = "python_full_version < '3.13'" }, - { name = "fsspec", marker = "python_full_version < '3.13'" }, - { name = "jinja2", marker = "python_full_version < '3.13'" }, - { name = "networkx", marker = "python_full_version < '3.13'" }, - { name = "nvidia-cublas-cu12", version = "12.6.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cuda-cupti-cu12", version = "12.6.80", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", version = "12.6.77", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cuda-runtime-cu12", version = "12.6.77", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cudnn-cu12", version = "9.5.1.17", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cufft-cu12", version = "11.3.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cufile-cu12", version = "1.11.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-curand-cu12", version = "10.3.7.77", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cusolver-cu12", version = "11.7.1.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cusparse-cu12", version = "12.5.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cusparselt-cu12", version = "0.6.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-nccl-cu12", version = "2.26.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-nvjitlink-cu12", version = "12.6.85", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-nvtx-cu12", version = "12.6.77", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "setuptools", marker = "python_full_version == '3.12.*'" }, - { name = "sympy", marker = "python_full_version < '3.13'" }, - { name = "triton", version = "3.3.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "setuptools", marker = "python_full_version >= '3.12'" }, + { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" }, + { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/6a/27/2e06cb52adf89fe6e020963529d17ed51532fc73c1e6d1b18420ef03338c/torch-2.7.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:a103b5d782af5bd119b81dbcc7ffc6fa09904c423ff8db397a1e6ea8fd71508f", size = 99089441 }, - { url = "https://files.pythonhosted.org/packages/0a/7c/0a5b3aee977596459ec45be2220370fde8e017f651fecc40522fd478cb1e/torch-2.7.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:fe955951bdf32d182ee8ead6c3186ad54781492bf03d547d31771a01b3d6fb7d", size = 821154516 }, - { url = "https://files.pythonhosted.org/packages/f9/91/3d709cfc5e15995fb3fe7a6b564ce42280d3a55676dad672205e94f34ac9/torch-2.7.1-cp310-cp310-win_amd64.whl", hash = "sha256:885453d6fba67d9991132143bf7fa06b79b24352f4506fd4d10b309f53454162", size = 216093147 }, - { url = "https://files.pythonhosted.org/packages/92/f6/5da3918414e07da9866ecb9330fe6ffdebe15cb9a4c5ada7d4b6e0a6654d/torch-2.7.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:d72acfdb86cee2a32c0ce0101606f3758f0d8bb5f8f31e7920dc2809e963aa7c", size = 68630914 }, - { url = "https://files.pythonhosted.org/packages/11/56/2eae3494e3d375533034a8e8cf0ba163363e996d85f0629441fa9d9843fe/torch-2.7.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:236f501f2e383f1cb861337bdf057712182f910f10aeaf509065d54d339e49b2", size = 99093039 }, - { url = "https://files.pythonhosted.org/packages/e5/94/34b80bd172d0072c9979708ccd279c2da2f55c3ef318eceec276ab9544a4/torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:06eea61f859436622e78dd0cdd51dbc8f8c6d76917a9cf0555a333f9eac31ec1", size = 821174704 }, - { url = "https://files.pythonhosted.org/packages/50/9e/acf04ff375b0b49a45511c55d188bcea5c942da2aaf293096676110086d1/torch-2.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:8273145a2e0a3c6f9fd2ac36762d6ee89c26d430e612b95a99885df083b04e52", size = 216095937 }, - { url = "https://files.pythonhosted.org/packages/5b/2b/d36d57c66ff031f93b4fa432e86802f84991477e522adcdffd314454326b/torch-2.7.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aea4fc1bf433d12843eb2c6b2204861f43d8364597697074c8d38ae2507f8730", size = 68640034 }, - { url = "https://files.pythonhosted.org/packages/87/93/fb505a5022a2e908d81fe9a5e0aa84c86c0d5f408173be71c6018836f34e/torch-2.7.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:27ea1e518df4c9de73af7e8a720770f3628e7f667280bce2be7a16292697e3fa", size = 98948276 }, - { url = "https://files.pythonhosted.org/packages/56/7e/67c3fe2b8c33f40af06326a3d6ae7776b3e3a01daa8f71d125d78594d874/torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c33360cfc2edd976c2633b3b66c769bdcbbf0e0b6550606d188431c81e7dd1fc", size = 821025792 }, - { url = "https://files.pythonhosted.org/packages/a1/37/a37495502bc7a23bf34f89584fa5a78e25bae7b8da513bc1b8f97afb7009/torch-2.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:d8bf6e1856ddd1807e79dc57e54d3335f2b62e6f316ed13ed3ecfe1fc1df3d8b", size = 216050349 }, - { url = "https://files.pythonhosted.org/packages/3a/60/04b77281c730bb13460628e518c52721257814ac6c298acd25757f6a175c/torch-2.7.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:787687087412c4bd68d315e39bc1223f08aae1d16a9e9771d95eabbb04ae98fb", size = 68645146 }, - { url = "https://files.pythonhosted.org/packages/66/81/e48c9edb655ee8eb8c2a6026abdb6f8d2146abd1f150979ede807bb75dcb/torch-2.7.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:03563603d931e70722dce0e11999d53aa80a375a3d78e6b39b9f6805ea0a8d28", size = 98946649 }, - { url = "https://files.pythonhosted.org/packages/3a/24/efe2f520d75274fc06b695c616415a1e8a1021d87a13c68ff9dce733d088/torch-2.7.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:d632f5417b6980f61404a125b999ca6ebd0b8b4bbdbb5fbbba44374ab619a412", size = 821033192 }, - { url = "https://files.pythonhosted.org/packages/dd/d9/9c24d230333ff4e9b6807274f6f8d52a864210b52ec794c5def7925f4495/torch-2.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:23660443e13995ee93e3d844786701ea4ca69f337027b05182f5ba053ce43b38", size = 216055668 }, - { url = "https://files.pythonhosted.org/packages/95/bf/e086ee36ddcef9299f6e708d3b6c8487c1651787bb9ee2939eb2a7f74911/torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:0da4f4dba9f65d0d203794e619fe7ca3247a55ffdcbd17ae8fb83c8b2dc9b585", size = 68925988 }, - { url = "https://files.pythonhosted.org/packages/69/6a/67090dcfe1cf9048448b31555af6efb149f7afa0a310a366adbdada32105/torch-2.7.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e08d7e6f21a617fe38eeb46dd2213ded43f27c072e9165dc27300c9ef9570934", size = 99028857 }, - { url = "https://files.pythonhosted.org/packages/90/1c/48b988870823d1cc381f15ec4e70ed3d65e043f43f919329b0045ae83529/torch-2.7.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:30207f672328a42df4f2174b8f426f354b2baa0b7cca3a0adb3d6ab5daf00dc8", size = 821098066 }, - { url = "https://files.pythonhosted.org/packages/7b/eb/10050d61c9d5140c5dc04a89ed3257ef1a6b93e49dd91b95363d757071e0/torch-2.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:79042feca1c634aaf6603fe6feea8c6b30dfa140a6bbc0b973e2260c7e79a22e", size = 216336310 }, - { url = "https://files.pythonhosted.org/packages/b1/29/beb45cdf5c4fc3ebe282bf5eafc8dfd925ead7299b3c97491900fe5ed844/torch-2.7.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:988b0cbc4333618a1056d2ebad9eb10089637b659eb645434d0809d8d937b946", size = 68645708 }, -] - -[[package]] -name = "torch" -version = "2.8.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13' and platform_system == 'Darwin' and sys_platform == 'darwin'", - "python_full_version >= '3.13' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'darwin'", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "python_full_version >= '3.13' and platform_machine == 'aarch64' and platform_system == 'Darwin' and sys_platform == 'linux'", - "python_full_version >= '3.13' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform == 'linux'", - "python_full_version >= '3.13' and platform_machine == 'aarch64' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'linux'", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system == 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system == 'Darwin' and sys_platform != 'darwin' and sys_platform != 'linux')", - "python_full_version >= '3.13' and platform_machine == 'aarch64' and platform_system == 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux'", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "filelock", marker = "python_full_version >= '3.13'" }, - { name = "fsspec", marker = "python_full_version >= '3.13'" }, - { name = "jinja2", marker = "python_full_version >= '3.13'" }, - { name = "networkx", marker = "python_full_version >= '3.13'" }, - { name = "nvidia-cublas-cu12", version = "12.8.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cuda-cupti-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cuda-runtime-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cudnn-cu12", version = "9.10.2.21", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cufft-cu12", version = "11.3.3.83", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cufile-cu12", version = "1.13.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-curand-cu12", version = "10.3.9.90", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cusolver-cu12", version = "11.7.3.90", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cusparse-cu12", version = "12.5.8.93", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-cusparselt-cu12", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-nccl-cu12", version = "2.27.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-nvjitlink-cu12", version = "12.8.93", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "nvidia-nvtx-cu12", version = "12.8.90", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "setuptools", marker = "python_full_version >= '3.13'" }, - { name = "sympy", marker = "python_full_version >= '3.13'" }, - { name = "triton", version = "3.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" }, - { name = "typing-extensions", marker = "python_full_version >= '3.13'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/63/28/110f7274254f1b8476c561dada127173f994afa2b1ffc044efb773c15650/torch-2.8.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:0be92c08b44009d4131d1ff7a8060d10bafdb7ddcb7359ef8d8c5169007ea905", size = 102052793 }, - { url = "https://files.pythonhosted.org/packages/70/1c/58da560016f81c339ae14ab16c98153d51c941544ae568da3cb5b1ceb572/torch-2.8.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:89aa9ee820bb39d4d72b794345cccef106b574508dd17dbec457949678c76011", size = 888025420 }, - { url = "https://files.pythonhosted.org/packages/70/87/f69752d0dd4ba8218c390f0438130c166fa264a33b7025adb5014b92192c/torch-2.8.0-cp310-cp310-win_amd64.whl", hash = "sha256:e8e5bf982e87e2b59d932769938b698858c64cc53753894be25629bdf5cf2f46", size = 241363614 }, - { url = "https://files.pythonhosted.org/packages/ef/d6/e6d4c57e61c2b2175d3aafbfb779926a2cfd7c32eeda7c543925dceec923/torch-2.8.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:a3f16a58a9a800f589b26d47ee15aca3acf065546137fc2af039876135f4c760", size = 73611154 }, - { url = "https://files.pythonhosted.org/packages/8f/c4/3e7a3887eba14e815e614db70b3b529112d1513d9dae6f4d43e373360b7f/torch-2.8.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:220a06fd7af8b653c35d359dfe1aaf32f65aa85befa342629f716acb134b9710", size = 102073391 }, - { url = "https://files.pythonhosted.org/packages/5a/63/4fdc45a0304536e75a5e1b1bbfb1b56dd0e2743c48ee83ca729f7ce44162/torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:c12fa219f51a933d5f80eeb3a7a5d0cbe9168c0a14bbb4055f1979431660879b", size = 888063640 }, - { url = "https://files.pythonhosted.org/packages/84/57/2f64161769610cf6b1c5ed782bd8a780e18a3c9d48931319f2887fa9d0b1/torch-2.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:8c7ef765e27551b2fbfc0f41bcf270e1292d9bf79f8e0724848b1682be6e80aa", size = 241366752 }, - { url = "https://files.pythonhosted.org/packages/a4/5e/05a5c46085d9b97e928f3f037081d3d2b87fb4b4195030fc099aaec5effc/torch-2.8.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:5ae0524688fb6707c57a530c2325e13bb0090b745ba7b4a2cd6a3ce262572916", size = 73621174 }, - { url = "https://files.pythonhosted.org/packages/49/0c/2fd4df0d83a495bb5e54dca4474c4ec5f9c62db185421563deeb5dabf609/torch-2.8.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:e2fab4153768d433f8ed9279c8133a114a034a61e77a3a104dcdf54388838705", size = 101906089 }, - { url = "https://files.pythonhosted.org/packages/99/a8/6acf48d48838fb8fe480597d98a0668c2beb02ee4755cc136de92a0a956f/torch-2.8.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:b2aca0939fb7e4d842561febbd4ffda67a8e958ff725c1c27e244e85e982173c", size = 887913624 }, - { url = "https://files.pythonhosted.org/packages/af/8a/5c87f08e3abd825c7dfecef5a0f1d9aa5df5dd0e3fd1fa2f490a8e512402/torch-2.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:2f4ac52f0130275d7517b03a33d2493bab3693c83dcfadf4f81688ea82147d2e", size = 241326087 }, - { url = "https://files.pythonhosted.org/packages/be/66/5c9a321b325aaecb92d4d1855421e3a055abd77903b7dab6575ca07796db/torch-2.8.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:619c2869db3ada2c0105487ba21b5008defcc472d23f8b80ed91ac4a380283b0", size = 73630478 }, - { url = "https://files.pythonhosted.org/packages/10/4e/469ced5a0603245d6a19a556e9053300033f9c5baccf43a3d25ba73e189e/torch-2.8.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:2b2f96814e0345f5a5aed9bf9734efa913678ed19caf6dc2cddb7930672d6128", size = 101936856 }, - { url = "https://files.pythonhosted.org/packages/16/82/3948e54c01b2109238357c6f86242e6ecbf0c63a1af46906772902f82057/torch-2.8.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:65616ca8ec6f43245e1f5f296603e33923f4c30f93d65e103d9e50c25b35150b", size = 887922844 }, - { url = "https://files.pythonhosted.org/packages/e3/54/941ea0a860f2717d86a811adf0c2cd01b3983bdd460d0803053c4e0b8649/torch-2.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:659df54119ae03e83a800addc125856effda88b016dfc54d9f65215c3975be16", size = 241330968 }, - { url = "https://files.pythonhosted.org/packages/de/69/8b7b13bba430f5e21d77708b616f767683629fc4f8037564a177d20f90ed/torch-2.8.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:1a62a1ec4b0498930e2543535cf70b1bef8c777713de7ceb84cd79115f553767", size = 73915128 }, - { url = "https://files.pythonhosted.org/packages/15/0e/8a800e093b7f7430dbaefa80075aee9158ec22e4c4fc3c1a66e4fb96cb4f/torch-2.8.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:83c13411a26fac3d101fe8035a6b0476ae606deb8688e904e796a3534c197def", size = 102020139 }, - { url = "https://files.pythonhosted.org/packages/4a/15/5e488ca0bc6162c86a33b58642bc577c84ded17c7b72d97e49b5833e2d73/torch-2.8.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8f0a9d617a66509ded240add3754e462430a6c1fc5589f86c17b433dd808f97a", size = 887990692 }, - { url = "https://files.pythonhosted.org/packages/b4/a8/6a04e4b54472fc5dba7ca2341ab219e529f3c07b6941059fbf18dccac31f/torch-2.8.0-cp313-cp313t-win_amd64.whl", hash = "sha256:a7242b86f42be98ac674b88a4988643b9bc6145437ec8f048fea23f72feb5eca", size = 241603453 }, - { url = "https://files.pythonhosted.org/packages/04/6e/650bb7f28f771af0cb791b02348db8b7f5f64f40f6829ee82aa6ce99aabe/torch-2.8.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:7b677e17f5a3e69fdef7eb3b9da72622f8d322692930297e4ccb52fefc6c8211", size = 73632395 }, + { url = "https://files.pythonhosted.org/packages/5f/56/9577683b23072075ed2e40d725c52c2019d71a972fab8e083763da8e707e/torch-2.9.1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:1cc208435f6c379f9b8fdfd5ceb5be1e3b72a6bdf1cb46c0d2812aa73472db9e", size = 104207681 }, + { url = "https://files.pythonhosted.org/packages/38/45/be5a74f221df8f4b609b78ff79dc789b0cc9017624544ac4dd1c03973150/torch-2.9.1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:9fd35c68b3679378c11f5eb73220fdcb4e6f4592295277fbb657d31fd053237c", size = 899794036 }, + { url = "https://files.pythonhosted.org/packages/67/95/a581e8a382596b69385a44bab2733f1273d45c842f5d4a504c0edc3133b6/torch-2.9.1-cp310-cp310-win_amd64.whl", hash = "sha256:2af70e3be4a13becba4655d6cc07dcfec7ae844db6ac38d6c1dafeb245d17d65", size = 110969861 }, + { url = "https://files.pythonhosted.org/packages/ad/51/1756dc128d2bf6ea4e0a915cb89ea5e730315ff33d60c1ff56fd626ba3eb/torch-2.9.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:a83b0e84cc375e3318a808d032510dde99d696a85fe9473fc8575612b63ae951", size = 74452222 }, + { url = "https://files.pythonhosted.org/packages/15/db/c064112ac0089af3d2f7a2b5bfbabf4aa407a78b74f87889e524b91c5402/torch-2.9.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:62b3fd888277946918cba4478cf849303da5359f0fb4e3bfb86b0533ba2eaf8d", size = 104220430 }, + { url = "https://files.pythonhosted.org/packages/56/be/76eaa36c9cd032d3b01b001e2c5a05943df75f26211f68fae79e62f87734/torch-2.9.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:d033ff0ac3f5400df862a51bdde9bad83561f3739ea0046e68f5401ebfa67c1b", size = 899821446 }, + { url = "https://files.pythonhosted.org/packages/47/cc/7a2949e38dfe3244c4df21f0e1c27bce8aedd6c604a587dd44fc21017cb4/torch-2.9.1-cp311-cp311-win_amd64.whl", hash = "sha256:0d06b30a9207b7c3516a9e0102114024755a07045f0c1d2f2a56b1819ac06bcb", size = 110973074 }, + { url = "https://files.pythonhosted.org/packages/1e/ce/7d251155a783fb2c1bb6837b2b7023c622a2070a0a72726ca1df47e7ea34/torch-2.9.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:52347912d868653e1528b47cafaf79b285b98be3f4f35d5955389b1b95224475", size = 74463887 }, + { url = "https://files.pythonhosted.org/packages/0f/27/07c645c7673e73e53ded71705045d6cb5bae94c4b021b03aa8d03eee90ab/torch-2.9.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:da5f6f4d7f4940a173e5572791af238cb0b9e21b1aab592bd8b26da4c99f1cd6", size = 104126592 }, + { url = "https://files.pythonhosted.org/packages/19/17/e377a460603132b00760511299fceba4102bd95db1a0ee788da21298ccff/torch-2.9.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:27331cd902fb4322252657f3902adf1c4f6acad9dcad81d8df3ae14c7c4f07c4", size = 899742281 }, + { url = "https://files.pythonhosted.org/packages/b1/1a/64f5769025db846a82567fa5b7d21dba4558a7234ee631712ee4771c436c/torch-2.9.1-cp312-cp312-win_amd64.whl", hash = "sha256:81a285002d7b8cfd3fdf1b98aa8df138d41f1a8334fd9ea37511517cedf43083", size = 110940568 }, + { url = "https://files.pythonhosted.org/packages/6e/ab/07739fd776618e5882661d04c43f5b5586323e2f6a2d7d84aac20d8f20bd/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:c0d25d1d8e531b8343bea0ed811d5d528958f1dcbd37e7245bc686273177ad7e", size = 74479191 }, + { url = "https://files.pythonhosted.org/packages/20/60/8fc5e828d050bddfab469b3fe78e5ab9a7e53dda9c3bdc6a43d17ce99e63/torch-2.9.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:c29455d2b910b98738131990394da3e50eea8291dfeb4b12de71ecf1fdeb21cb", size = 104135743 }, + { url = "https://files.pythonhosted.org/packages/f2/b7/6d3f80e6918213babddb2a37b46dbb14c15b14c5f473e347869a51f40e1f/torch-2.9.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:524de44cd13931208ba2c4bde9ec7741fd4ae6bfd06409a604fc32f6520c2bc9", size = 899749493 }, + { url = "https://files.pythonhosted.org/packages/a6/47/c7843d69d6de8938c1cbb1eba426b1d48ddf375f101473d3e31a5fc52b74/torch-2.9.1-cp313-cp313-win_amd64.whl", hash = "sha256:545844cc16b3f91e08ce3b40e9c2d77012dd33a48d505aed34b7740ed627a1b2", size = 110944162 }, + { url = "https://files.pythonhosted.org/packages/28/0e/2a37247957e72c12151b33a01e4df651d9d155dd74d8cfcbfad15a79b44a/torch-2.9.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5be4bf7496f1e3ffb1dd44b672adb1ac3f081f204c5ca81eba6442f5f634df8e", size = 74830751 }, + { url = "https://files.pythonhosted.org/packages/4b/f7/7a18745edcd7b9ca2381aa03353647bca8aace91683c4975f19ac233809d/torch-2.9.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:30a3e170a84894f3652434b56d59a64a2c11366b0ed5776fab33c2439396bf9a", size = 104142929 }, + { url = "https://files.pythonhosted.org/packages/f4/dd/f1c0d879f2863ef209e18823a988dc7a1bf40470750e3ebe927efdb9407f/torch-2.9.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8301a7b431e51764629208d0edaa4f9e4c33e6df0f2f90b90e261d623df6a4e2", size = 899748978 }, + { url = "https://files.pythonhosted.org/packages/1f/9f/6986b83a53b4d043e36f3f898b798ab51f7f20fdf1a9b01a2720f445043d/torch-2.9.1-cp313-cp313t-win_amd64.whl", hash = "sha256:2e1c42c0ae92bf803a4b2409fdfed85e30f9027a66887f5e7dcdbc014c7531db", size = 111176995 }, + { url = "https://files.pythonhosted.org/packages/40/60/71c698b466dd01e65d0e9514b5405faae200c52a76901baf6906856f17e4/torch-2.9.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:2c14b3da5df416cf9cb5efab83aa3056f5b8cd8620b8fde81b4987ecab730587", size = 74480347 }, + { url = "https://files.pythonhosted.org/packages/48/50/c4b5112546d0d13cc9eaa1c732b823d676a9f49ae8b6f97772f795874a03/torch-2.9.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1edee27a7c9897f4e0b7c14cfc2f3008c571921134522d5b9b5ec4ebbc69041a", size = 74433245 }, + { url = "https://files.pythonhosted.org/packages/81/c9/2628f408f0518b3bae49c95f5af3728b6ab498c8624ab1e03a43dd53d650/torch-2.9.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:19d144d6b3e29921f1fc70503e9f2fc572cde6a5115c0c0de2f7ca8b1483e8b6", size = 104134804 }, + { url = "https://files.pythonhosted.org/packages/28/fc/5bc91d6d831ae41bf6e9e6da6468f25330522e92347c9156eb3f1cb95956/torch-2.9.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c432d04376f6d9767a9852ea0def7b47a7bbc8e7af3b16ac9cf9ce02b12851c9", size = 899747132 }, + { url = "https://files.pythonhosted.org/packages/63/5d/e8d4e009e52b6b2cf1684bde2a6be157b96fb873732542fb2a9a99e85a83/torch-2.9.1-cp314-cp314-win_amd64.whl", hash = "sha256:d187566a2cdc726fc80138c3cdb260970fab1c27e99f85452721f7759bbd554d", size = 110934845 }, + { url = "https://files.pythonhosted.org/packages/bd/b2/2d15a52516b2ea3f414643b8de68fa4cb220d3877ac8b1028c83dc8ca1c4/torch-2.9.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cb10896a1f7fedaddbccc2017ce6ca9ecaaf990f0973bdfcf405439750118d2c", size = 74823558 }, + { url = "https://files.pythonhosted.org/packages/86/5c/5b2e5d84f5b9850cd1e71af07524d8cbb74cba19379800f1f9f7c997fc70/torch-2.9.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0a2bd769944991c74acf0c4ef23603b9c777fdf7637f115605a4b2d8023110c7", size = 104145788 }, + { url = "https://files.pythonhosted.org/packages/a9/8c/3da60787bcf70add986c4ad485993026ac0ca74f2fc21410bc4eb1bb7695/torch-2.9.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:07c8a9660bc9414c39cac530ac83b1fb1b679d7155824144a40a54f4a47bfa73", size = 899735500 }, + { url = "https://files.pythonhosted.org/packages/db/2b/f7818f6ec88758dfd21da46b6cd46af9d1b3433e53ddbb19ad1e0da17f9b/torch-2.9.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c88d3299ddeb2b35dcc31753305612db485ab6f1823e37fb29451c8b2732b87e", size = 111163659 }, ] [[package]] @@ -2512,7 +2115,7 @@ wheels = [ [[package]] name = "transformers" -version = "4.57.1" +version = "4.57.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, @@ -2527,51 +2130,30 @@ dependencies = [ { name = "tokenizers" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d6/68/a39307bcc4116a30b2106f2e689130a48de8bd8a1e635b5e1030e46fcd9e/transformers-4.57.1.tar.gz", hash = "sha256:f06c837959196c75039809636cd964b959f6604b75b8eeec6fdfc0440b89cc55", size = 10142511 } +sdist = { url = "https://files.pythonhosted.org/packages/dd/70/d42a739e8dfde3d92bb2fff5819cbf331fe9657323221e79415cd5eb65ee/transformers-4.57.3.tar.gz", hash = "sha256:df4945029aaddd7c09eec5cad851f30662f8bd1746721b34cc031d70c65afebc", size = 10139680 } wheels = [ - { url = "https://files.pythonhosted.org/packages/71/d3/c16c3b3cf7655a67db1144da94b021c200ac1303f82428f2beef6c2e72bb/transformers-4.57.1-py3-none-any.whl", hash = "sha256:b10d05da8fa67dc41644dbbf9bc45a44cb86ae33da6f9295f5fbf5b7890bd267", size = 11990925 }, + { url = "https://files.pythonhosted.org/packages/6a/6b/2f416568b3c4c91c96e5a365d164f8a4a4a88030aa8ab4644181fdadce97/transformers-4.57.3-py3-none-any.whl", hash = "sha256:c77d353a4851b1880191603d36acb313411d3577f6e2897814f333841f7003f4", size = 11993463 }, ] [[package]] name = "triton" -version = "3.3.1" +version = "3.5.1" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.12.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version == '3.11.*' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version < '3.11' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "setuptools", marker = "(python_full_version < '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version < '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, -] wheels = [ - { url = "https://files.pythonhosted.org/packages/8d/a9/549e51e9b1b2c9b854fd761a1d23df0ba2fbc60bd0c13b489ffa518cfcb7/triton-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b74db445b1c562844d3cfad6e9679c72e93fdfb1a90a24052b03bb5c49d1242e", size = 155600257 }, - { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937 }, - { url = "https://files.pythonhosted.org/packages/24/5f/950fb373bf9c01ad4eb5a8cd5eaf32cdf9e238c02f9293557a2129b9c4ac/triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43", size = 155669138 }, - { url = "https://files.pythonhosted.org/packages/74/1f/dfb531f90a2d367d914adfee771babbd3f1a5b26c3f5fbc458dee21daa78/triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240", size = 155673035 }, - { url = "https://files.pythonhosted.org/packages/28/71/bd20ffcb7a64c753dc2463489a61bf69d531f308e390ad06390268c4ea04/triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42", size = 155735832 }, -] - -[[package]] -name = "triton" -version = "3.4.0" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform == 'darwin')", - "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin' and sys_platform != 'darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'darwin' and sys_platform != 'linux')", -] -dependencies = [ - { name = "setuptools", marker = "(python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_system != 'Darwin') or (python_full_version >= '3.13' and platform_system != 'Darwin' and platform_system != 'Linux' and sys_platform != 'linux')" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/62/ee/0ee5f64a87eeda19bbad9bc54ae5ca5b98186ed00055281fd40fb4beb10e/triton-3.4.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7ff2785de9bc02f500e085420273bb5cc9c9bb767584a4aa28d6e360cec70128", size = 155430069 }, - { url = "https://files.pythonhosted.org/packages/7d/39/43325b3b651d50187e591eefa22e236b2981afcebaefd4f2fc0ea99df191/triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b70f5e6a41e52e48cfc087436c8a28c17ff98db369447bcaff3b887a3ab4467", size = 155531138 }, - { url = "https://files.pythonhosted.org/packages/d0/66/b1eb52839f563623d185f0927eb3530ee4d5ffe9d377cdaf5346b306689e/triton-3.4.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:31c1d84a5c0ec2c0f8e8a072d7fd150cab84a9c239eaddc6706c081bfae4eb04", size = 155560068 }, - { url = "https://files.pythonhosted.org/packages/30/7b/0a685684ed5322d2af0bddefed7906674f67974aa88b0fae6e82e3b766f6/triton-3.4.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00be2964616f4c619193cb0d1b29a99bd4b001d7dc333816073f92cf2a8ccdeb", size = 155569223 }, - { url = "https://files.pythonhosted.org/packages/20/63/8cb444ad5cdb25d999b7d647abac25af0ee37d292afc009940c05b82dda0/triton-3.4.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7936b18a3499ed62059414d7df563e6c163c5e16c3773678a3ee3d417865035d", size = 155659780 }, + { url = "https://files.pythonhosted.org/packages/d9/2e/f95e673222afa2c7f0c687d8913e98fcf2589ef0b1405de76894e37fe18f/triton-3.5.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f63e34dcb32d7bd3a1d0195f60f30d2aee8b08a69a0424189b71017e23dfc3d2", size = 159821655 }, + { url = "https://files.pythonhosted.org/packages/fd/6e/676ab5019b4dde8b9b7bab71245102fc02778ef3df48218b298686b9ffd6/triton-3.5.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5fc53d849f879911ea13f4a877243afc513187bc7ee92d1f2c0f1ba3169e3c94", size = 170320692 }, + { url = "https://files.pythonhosted.org/packages/dc/dc/6ce44d055f2fc2403c4ec6b3cfd3a9b25f57b7d95efadccdea91497f8e81/triton-3.5.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da47169e30a779bade679ce78df4810fca6d78a955843d2ddb11f226adc517dc", size = 159928005 }, + { url = "https://files.pythonhosted.org/packages/b0/72/ec90c3519eaf168f22cb1757ad412f3a2add4782ad3a92861c9ad135d886/triton-3.5.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:61413522a48add32302353fdbaaf92daaaab06f6b5e3229940d21b5207f47579", size = 170425802 }, + { url = "https://files.pythonhosted.org/packages/db/53/2bcc46879910991f09c063eea07627baef2bc62fe725302ba8f46a2c1ae5/triton-3.5.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:275a045b6ed670dd1bd005c3e6c2d61846c74c66f4512d6f33cc027b11de8fd4", size = 159940689 }, + { url = "https://files.pythonhosted.org/packages/f2/50/9a8358d3ef58162c0a415d173cfb45b67de60176e1024f71fbc4d24c0b6d/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d2c6b915a03888ab931a9fd3e55ba36785e1fe70cbea0b40c6ef93b20fc85232", size = 170470207 }, + { url = "https://files.pythonhosted.org/packages/f1/ba/805684a992ee32d486b7948d36aed2f5e3c643fc63883bf8bdca1c3f3980/triton-3.5.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:56765ffe12c554cd560698398b8a268db1f616c120007bfd8829d27139abd24a", size = 159955460 }, + { url = "https://files.pythonhosted.org/packages/27/46/8c3bbb5b0a19313f50edcaa363b599e5a1a5ac9683ead82b9b80fe497c8d/triton-3.5.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3f4346b6ebbd4fad18773f5ba839114f4826037c9f2f34e0148894cd5dd3dba", size = 170470410 }, + { url = "https://files.pythonhosted.org/packages/84/1e/7df59baef41931e21159371c481c31a517ff4c2517343b62503d0cd2be99/triton-3.5.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:02c770856f5e407d24d28ddc66e33cf026e6f4d360dcb8b2fabe6ea1fc758621", size = 160072799 }, + { url = "https://files.pythonhosted.org/packages/37/92/e97fcc6b2c27cdb87ce5ee063d77f8f26f19f06916aa680464c8104ef0f6/triton-3.5.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0b4d2c70127fca6a23e247f9348b8adde979d2e7a20391bfbabaac6aebc7e6a8", size = 170579924 }, + { url = "https://files.pythonhosted.org/packages/14/f9/0430e879c1e63a1016cb843261528fd3187c872c3a9539132efc39514753/triton-3.5.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f617aa7925f9ea9968ec2e1adaf93e87864ff51549c8f04ce658f29bbdb71e2d", size = 159956163 }, + { url = "https://files.pythonhosted.org/packages/a4/e6/c595c35e5c50c4bc56a7bac96493dad321e9e29b953b526bbbe20f9911d0/triton-3.5.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0637b1efb1db599a8e9dc960d53ab6e4637db7d4ab6630a0974705d77b14b60", size = 170480488 }, + { url = "https://files.pythonhosted.org/packages/41/1e/63d367c576c75919e268e4fbc33c1cb33b6dc12bb85e8bfe531c2a8bd5d3/triton-3.5.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8932391d7f93698dfe5bc9bead77c47a24f97329e9f20c10786bb230a9083f56", size = 160073620 }, + { url = "https://files.pythonhosted.org/packages/16/b5/b0d3d8b901b6a04ca38df5e24c27e53afb15b93624d7fd7d658c7cd9352a/triton-3.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bac7f7d959ad0f48c0e97d6643a1cc0fd5786fe61cb1f83b537c6b2d54776478", size = 170582192 }, ] [[package]]