🖼️↗️ Add scale factor support for image upscaling dimensions (#215)
Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com> Co-authored-by: filipstrand <strand.filip@gmail.com>
This commit is contained in:
Anthony Wu
GitHub
parent
0125103634
commit
b26a1d66b3
255
README.md
255
README.md
@ -12,54 +12,38 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
|
||||
|
||||
- [Philosophy](#philosophy)
|
||||
- [💿 Installation](#-installation)
|
||||
- [🚀 Shell Completions (Quick Start)](#-shell-completions-quick-start)
|
||||
- [🖼️ Generating an image](#%EF%B8%8F-generating-an-image)
|
||||
* [📜 Full list of Command-Line Arguments](#-full-list-of-command-line-arguments)
|
||||
- [⏱️ 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)
|
||||
* [📊 Size comparisons for quantized models](#-size-comparisons-for-quantized-models)
|
||||
* [💾 Saving a quantized version to disk](#-saving-a-quantized-version-to-disk)
|
||||
* [💽 Loading and running a quantized version from disk](#-loading-and-running-a-quantized-version-from-disk)
|
||||
- [💽 Running a non-quantized model directly from disk](#-running-a-non-quantized-model-directly-from-disk)
|
||||
- [🌐 Third-Party HuggingFace Model Support](#-third-party-huggingface-model-support)
|
||||
- [🎨 Image-to-Image](#-image-to-image)
|
||||
- [🔌 LoRA](#-lora)
|
||||
* [Multi-LoRA](#multi-lora)
|
||||
* [LoRA Library Path](#lora-library-path)
|
||||
* [Supported LoRA formats (updated)](#supported-lora-formats-updated)
|
||||
- [🎭 In-Context Generation](#-in-context-generation)
|
||||
* [🎨 In-Context LoRA](#-in-context-lora)
|
||||
+ [Available Styles](#available-styles)
|
||||
+ [How It Works](#how-it-works)
|
||||
+ [Tips for Best Results](#tips-for-best-results)
|
||||
* [👕 CatVTON (Virtual Try-On)](#-catvton-virtual-try-on)
|
||||
* [✏️ IC-Edit (In-Context Editing)](#%EF%B8%8F-ic-edit-in-context-editing)
|
||||
- [🛠️ Flux Tools](#%EF%B8%8F-flux-tools)
|
||||
* [🖌️ Fill](#%EF%B8%8F-fill)
|
||||
+ [Inpainting](#inpainting)
|
||||
+ [Outpainting](#outpainting)
|
||||
* [🔍 Depth](#-depth)
|
||||
* [🔄 Redux](#-redux)
|
||||
- [🕹️ Controlnet](#%EF%B8%8F-controlnet)
|
||||
- [🔎 Upscale](#-upscale)
|
||||
- [🎛️ Dreambooth fine-tuning](#-dreambooth-fine-tuning)
|
||||
* [Training configuration](#training-configuration)
|
||||
* [Training example](#training-example)
|
||||
* [Resuming a training run](#resuming-a-training-run)
|
||||
* [Configuration details](#configuration-details)
|
||||
* [Memory issues](#memory-issues)
|
||||
* [Misc](#misc)
|
||||
- [🎛️ Dreambooth fine-tuning](#%EF%B8%8F-dreambooth-fine-tuning)
|
||||
- [🧠 Concept Attention](#-concept-attention)
|
||||
- [🚧 Current limitations](#-current-limitations)
|
||||
- [💡Workflow tips](#workflow-tips)
|
||||
- [🔬 Cool research / features to support](#-cool-research--features-to-support-)
|
||||
- [🔬 Cool research](#-cool-research)
|
||||
- [🌱 Related projects](#-related-projects)
|
||||
- [🙏 Acknowledgements](#-acknowledgements)
|
||||
- [⚖️ License](#-license)
|
||||
- [⚖️ License](#%EF%B8%8F-license)
|
||||
|
||||
<!-- TOC end -->
|
||||
|
||||
---
|
||||
|
||||
### Philosophy
|
||||
|
||||
MFLUX is a line-by-line port of the FLUX implementation in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) library to [Apple MLX](https://github.com/ml-explore/mlx).
|
||||
@ -71,6 +55,7 @@ All models are implemented from scratch in MLX and only the tokenizers are used
|
||||
[Huggingface Transformers](https://github.com/huggingface/transformers) library. Other than that, there are only minimal dependencies
|
||||
like [Numpy](https://numpy.org) and [Pillow](https://pypi.org/project/pillow/) for simple image post-processing.
|
||||
|
||||
---
|
||||
|
||||
### 💿 Installation
|
||||
For users, the easiest way to install MFLUX is to use `uv tool`: If you have [installed `uv`](https://github.com/astral-sh/uv?tab=readme-ov-file#installation), simply:
|
||||
@ -140,9 +125,10 @@ pip install -U mflux
|
||||
|
||||
*If you have trouble installing MFLUX, please see the [installation related issues section](https://github.com/filipstrand/mflux/issues?q=is%3Aissue+install+).*
|
||||
|
||||
### ⌨️ Shell Completions (Quick Start)
|
||||
<details>
|
||||
<summary>⌨️ <strong>Shell Completions (Optional)</strong></summary>
|
||||
|
||||
MFLUX supports ZSH (default on macOS) shell completions for all CLI commands.
|
||||
MFLUX supports ZSH (default on macOS) shell completions for all CLI commands. This provides tab completion for all [command-line arguments](#-full-list-of-command-line-arguments) and options.
|
||||
|
||||
To enable completions:
|
||||
|
||||
@ -165,6 +151,10 @@ mflux-completions --check
|
||||
|
||||
For more details and troubleshooting, see the [completions documentation](src/mflux/completions/README.md).
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
### 🖼️ Generating an image
|
||||
|
||||
Run the command `mflux-generate` by specifying a prompt and the model and some optional arguments. For example, here we use a quantized version of the `schnell` model for 2 steps:
|
||||
@ -187,6 +177,33 @@ echo "A majestic mountain landscape" | mflux-generate --prompt - --model schnell
|
||||
|
||||
This is useful for integrating MFLUX into shell scripts or dynamically generating prompts using LLM inference tools such as [`llm`](https://llm.datasette.io/en/stable/), [`mlx-lm`](https://github.com/ml-explore/mlx-lm), [`ollama`](https://ollama.ai/), etc.
|
||||
|
||||
Alternatively, you can use MFLUX directly in Python:
|
||||
|
||||
```python
|
||||
from mflux import Flux1, Config
|
||||
|
||||
# Load the model
|
||||
flux = Flux1.from_name(
|
||||
model_name="schnell", # "schnell" or "dev"
|
||||
quantize=8, # 4 or 8
|
||||
)
|
||||
|
||||
# Generate an image
|
||||
image = flux.generate_image(
|
||||
seed=2,
|
||||
prompt="Luxury food photograph",
|
||||
config=Config(
|
||||
num_inference_steps=2, # "schnell" works well with 2-4 steps, "dev" works well with 20-25 steps
|
||||
height=1024,
|
||||
width=1024,
|
||||
)
|
||||
)
|
||||
|
||||
image.save(path="image.png")
|
||||
```
|
||||
|
||||
For more advanced Python usage and additional configuration options, you can explore the entry point files in the source code, such as [`generate.py`](src/mflux/generate.py), [`generate_controlnet.py`](src/mflux/generate_controlnet.py), [`generate_fill.py`](src/mflux/generate_fill.py), and others in the [`src/mflux/`](src/mflux/) directory. These files demonstrate how to use the Python API for various features and provide examples of advanced configurations.
|
||||
|
||||
⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#%EF%B8%8F-quantization) section for running compressed versions of the model.* ⚠️
|
||||
|
||||
*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/`).*
|
||||
@ -197,6 +214,14 @@ This is useful for integrating MFLUX into shell scripts or dynamically generatin
|
||||
|
||||
#### 📜 Full list of Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>📜 <strong>Command-Line Arguments Reference</strong></summary>
|
||||
|
||||
#### General Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand General arguments</summary>
|
||||
|
||||
- **`--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 models (`"schnell"` or `"dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`).
|
||||
@ -247,7 +272,12 @@ This is useful for integrating MFLUX into shell scripts or dynamically generatin
|
||||
|
||||
- **`--vae-tiling-split`** (optional, `str`, default: `"horizontal"`): When VAE tiling is enabled, this parameter controls the direction to split the latents. Options are `"horizontal"` (splits into top/bottom) or `"vertical"` (splits into left/right). Use this option to control where potential seams might appear in the final image.
|
||||
|
||||
#### 📜 In-Context LoRA Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### In-Context LoRA Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand In-Context LoRA arguments</summary>
|
||||
|
||||
The `mflux-generate-in-context` command supports most of the same arguments as `mflux-generate`, with these additional parameters:
|
||||
|
||||
@ -257,7 +287,12 @@ The `mflux-generate-in-context` command supports most of the same arguments as `
|
||||
|
||||
See the [In-Context Generation](#-in-context-generation) section for more details on how to use this feature effectively.
|
||||
|
||||
#### 📜 CatVTON Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### CatVTON Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand CatVTON arguments</summary>
|
||||
|
||||
The `mflux-generate-in-context-catvton` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
@ -271,7 +306,12 @@ The `mflux-generate-in-context-catvton` command supports most of the same argume
|
||||
|
||||
See the [CatVTON (Virtual Try-On)](#-catvton-virtual-try-on) section for more details on this feature.
|
||||
|
||||
#### 📜 IC-Edit Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### IC-Edit Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand IC-Edit arguments</summary>
|
||||
|
||||
The `mflux-generate-in-context-edit` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
@ -285,7 +325,12 @@ The `mflux-generate-in-context-edit` command supports most of the same arguments
|
||||
|
||||
See the [IC-Edit (In-Context Editing)](#-ic-edit-in-context-editing) section for more details on this feature.
|
||||
|
||||
#### 📜 Redux Tool Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### Redux Tool Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand Redux arguments</summary>
|
||||
|
||||
The `mflux-generate-redux` command uses most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
@ -295,7 +340,12 @@ The `mflux-generate-redux` command uses most of the same arguments as `mflux-gen
|
||||
|
||||
See the [Redux](#-redux) section for more details on this feature.
|
||||
|
||||
#### 📜 Concept Attention Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### Concept Attention Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand Concept Attention arguments</summary>
|
||||
|
||||
The `mflux-concept` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
@ -311,7 +361,12 @@ The `mflux-concept-from-image` command uses most of the same arguments as `mflux
|
||||
|
||||
See the [Concept Attention](#-concept-attention) section for more details on this feature.
|
||||
|
||||
#### 📜 Fill Tool Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### Fill Tool Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand Fill Tool arguments</summary>
|
||||
|
||||
The `mflux-generate-fill` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
@ -323,7 +378,12 @@ The `mflux-generate-fill` command supports most of the same arguments as `mflux-
|
||||
|
||||
See the [Fill](#-fill) section for more details on inpainting and outpainting.
|
||||
|
||||
#### 📜 Depth Tool Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### Depth Tool Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand Depth Tool arguments</summary>
|
||||
|
||||
The `mflux-generate-depth` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
@ -341,7 +401,12 @@ The `mflux-save-depth` command for extracting depth maps without generating imag
|
||||
|
||||
See the [Depth](#-depth) section for more details on this feature.
|
||||
|
||||
#### 📜 ControlNet Command-Line Arguments
|
||||
</details>
|
||||
|
||||
#### ControlNet Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand ControlNet arguments</summary>
|
||||
|
||||
The `mflux-generate-controlnet` command supports most of the same arguments as `mflux-generate`, with these additional parameters:
|
||||
|
||||
@ -353,8 +418,57 @@ The `mflux-generate-controlnet` command supports most of the same arguments as `
|
||||
|
||||
See the [Controlnet](#%EF%B8%8F-controlnet) section for more details on how to use this feature effectively.
|
||||
|
||||
</details>
|
||||
|
||||
#### Dynamic Prompts with `--prompt-file`
|
||||
#### Upscale Command-Line Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand Upscale arguments</summary>
|
||||
|
||||
The `mflux-upscale` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
|
||||
|
||||
- **`--height`** (optional, `int` or scale factor, default: `auto`): Image height. Can be specified as pixels (e.g., `1024`), scale factors (e.g., `2x`, `1.5x`), or `auto` to use the source image height.
|
||||
|
||||
- **`--width`** (optional, `int` or scale factor, default: `auto`): Image width. Can be specified as pixels (e.g., `1024`), scale factors (e.g., `2x`, `1.5x`), or `auto` to use the source image width.
|
||||
|
||||
- **`--controlnet-image-path`** (required, `str`): Path to the source image to upscale.
|
||||
|
||||
- **`--controlnet-strength`** (optional, `float`, default: `0.4`): Degree of influence the control image has on the output. Ranges from `0.0` (no influence) to `1.0` (full influence).
|
||||
|
||||
**Scale Factor Examples:**
|
||||
```bash
|
||||
# Scale by 2x in both dimensions
|
||||
mflux-upscale --height 2x --width 2x --controlnet-image-path source.png
|
||||
|
||||
# Scale height by 1.5x, set width to specific pixels
|
||||
mflux-upscale --height 1.5x --width 1920 --controlnet-image-path source.png
|
||||
|
||||
# Use auto (keeps original dimensions)
|
||||
mflux-upscale --height auto --width auto --controlnet-image-path source.png
|
||||
|
||||
# Mix scale factors and absolute values
|
||||
mflux-upscale --height 2x --width 1024 --controlnet-image-path source.png
|
||||
```
|
||||
|
||||
See the [Upscale](#-upscale) section for more details on how to use this feature effectively.
|
||||
|
||||
</details>
|
||||
|
||||
#### Training Arguments
|
||||
|
||||
<details>
|
||||
<summary>Click to expand Training arguments</summary>
|
||||
|
||||
- **`--train-config`** (optional, `str`): Local path of the training configuration file. This file defines all aspects of the training process including model parameters, optimizer settings, and training data. See the [Training configuration](#training-configuration) section for details on the structure of this file.
|
||||
|
||||
- **`--train-checkpoint`** (optional, `str`): Local path of the checkpoint file which specifies how to continue the training process. Used when resuming an interrupted training run.
|
||||
|
||||
</details>
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>📝 <strong>Dynamic Prompts with `--prompt-file`</strong></summary>
|
||||
|
||||
MFlux supports dynamic prompt updates through the `--prompt-file` option. Instead of providing a fixed prompt with `--prompt`, you can specify a plain text file containing your prompt. The file is re-read before each generation, allowing you to modify prompts between iterations without restarting.
|
||||
|
||||
@ -399,15 +513,10 @@ mflux-generate --prompt-file my_prompt.txt --auto-seeds 10
|
||||
- Empty prompt files or non-existent files will raise appropriate errors
|
||||
- Each generated image's metadata will contain the actual prompt used for that specific generation
|
||||
|
||||
#### 📜 Training Arguments
|
||||
|
||||
- **`--train-config`** (optional, `str`): Local path of the training configuration file. This file defines all aspects of the training process including model parameters, optimizer settings, and training data. See the [Training configuration](#training-configuration) section for details on the structure of this file.
|
||||
|
||||
- **`--train-checkpoint`** (optional, `str`): Local path of the checkpoint file which specifies how to continue the training process. Used when resuming an interrupted training run.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>parameters supported by config files</summary>
|
||||
<summary>⚙️ <strong>Parameters supported by config files</strong></summary>
|
||||
|
||||
#### How configs are used
|
||||
|
||||
@ -474,32 +583,7 @@ mflux-generate --prompt-file my_prompt.txt --auto-seeds 10
|
||||
```
|
||||
</details>
|
||||
|
||||
Or, with the correct python environment active, create and run a separate script like the following:
|
||||
|
||||
```python
|
||||
from mflux import Flux1, Config
|
||||
|
||||
# Load the model
|
||||
flux = Flux1.from_name(
|
||||
model_name="schnell", # "schnell" or "dev"
|
||||
quantize=8, # 4 or 8
|
||||
)
|
||||
|
||||
# Generate an image
|
||||
image = flux.generate_image(
|
||||
seed=2,
|
||||
prompt="Luxury food photograph",
|
||||
config=Config(
|
||||
num_inference_steps=2, # "schnell" works well with 2-4 steps, "dev" works well with 20-25 steps
|
||||
height=1024,
|
||||
width=1024,
|
||||
)
|
||||
)
|
||||
|
||||
image.save(path="image.png")
|
||||
```
|
||||
|
||||
For more options on how to configure MFLUX, please see [generate.py](src/mflux/generate.py).
|
||||
---
|
||||
|
||||
### ⏱️ Image generation speed (updated)
|
||||
|
||||
@ -542,6 +626,8 @@ If we assume that the model is already loaded, you can inspect the image metadat
|
||||
|
||||
*These benchmarks are not very scientific and is only intended to give ballpark numbers. They were performed during different times with different MFLUX and MLX-versions etc. Additional hardware information such as number of GPU cores, Mac device etc. are not always known.*
|
||||
|
||||
---
|
||||
|
||||
### ↔️ Equivalent to Diffusers implementation
|
||||
|
||||
There is only a single source of randomness when generating an image: The initial latent array.
|
||||
@ -697,6 +783,8 @@ mflux-generate \
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
### 💽 Running a non-quantized model directly from disk
|
||||
|
||||
MFLUX also supports running a non-quantized model directly from a custom location.
|
||||
@ -716,6 +804,9 @@ Note that the `--model` flag must be set when loading a model from disk.
|
||||
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:
|
||||
|
||||
<details>
|
||||
<summary>📁 <strong>Required directory structure</strong></summary>
|
||||
|
||||
```
|
||||
.
|
||||
├── text_encoder
|
||||
@ -740,6 +831,9 @@ when loading a model directly from disk, we require the downloaded models to loo
|
||||
└── vae
|
||||
└── diffusion_pytorch_model.safetensors
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
This mirrors how the resources are placed in the [HuggingFace Repo](https://huggingface.co/black-forest-labs/FLUX.1-schnell/tree/main) for FLUX.1.
|
||||
*Huggingface weights, unlike quantized ones exported directly from this project, have to be
|
||||
processed a bit differently, which is why we require this structure above.*
|
||||
@ -1437,7 +1531,8 @@ The zip file will contain configuration files which point to the original datase
|
||||
*⚠️ Note: One current limitation is that a training run can only be resumed if it has not yet been completed.
|
||||
In other words, only checkpoints that represent an interrupted training-run can be resumed and run until completion.*
|
||||
|
||||
#### Configuration details
|
||||
<details>
|
||||
<summary>⚙️ <strong>Configuration details</strong></summary>
|
||||
|
||||
Currently, MFLUX supports fine-tuning only for the transformer part of the model.
|
||||
In the training configuration, under `lora_layers`, you can specify which layers you want to train. The available ones are:
|
||||
@ -1499,7 +1594,10 @@ In other words, training later layers, such as only the `single_transformer_bloc
|
||||
|
||||
*Under the `examples` section, there is an argument called `"path"` which specifies where the images are located. This path is relative to the config file itself.*
|
||||
|
||||
#### Memory issues
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>⚠️ <strong>Memory issues</strong></summary>
|
||||
|
||||
Depending on the configuration of the training setup, fine-tuning can be quite memory intensive.
|
||||
In the worst case, if your Mac runs out of memory it might freeze completely and crash!
|
||||
@ -1518,9 +1616,12 @@ will allow a 32GB M1 Pro to perform a successful fine-tuning run.
|
||||
Note, however, that reducing the trainable parameters might lead to worse performance.
|
||||
|
||||
|
||||
*Additional techniques such as gradient checkpoint and other strategies might be implemented in the future.*
|
||||
*Additional techniques such as gradient checkpoint and other strategies might be implemented in the future.*
|
||||
|
||||
#### Misc
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary>📝 <strong>Misc</strong></summary>
|
||||
|
||||
This feature is currently v1 and can be considered a bit experimental. Interfaces might change (configuration file setup etc.)
|
||||
The aim is to also gradually expand the scope of this feature with alternative techniques, data augmentation etc.
|
||||
@ -1536,6 +1637,8 @@ The aim is to also gradually expand the scope of this feature with alternative t
|
||||
- The fine-tuning script in [mlx-examples](https://github.com/ml-explore/mlx-examples/tree/main/flux#finetuning)
|
||||
- The original fine-tuning script in [Diffusers](https://huggingface.co/docs/diffusers/v0.11.0/en/training/dreambooth)
|
||||
|
||||
</details>
|
||||
|
||||
|
||||
---
|
||||
|
||||
@ -1623,6 +1726,8 @@ This will generate the following image
|
||||
- Dreambooth training currently does not support sending in training parameters as flags.
|
||||
- In-Context Generation features currently only support a left-right image setup (reference image on left, generated image on right).
|
||||
|
||||
---
|
||||
|
||||
### Optional Tool: Batch Image Renamer
|
||||
|
||||
With a large number of generated images, some users want to automatically rename their image outputs to reflect the prompts and configs.
|
||||
@ -1642,6 +1747,8 @@ and `uv run your/path/rename_images.py`.
|
||||
This script's renaming logic can be customized to your needs.
|
||||
See `uv run tools/rename_images.py --help` for full CLI usage help.
|
||||
|
||||
---
|
||||
|
||||
### 💡Workflow Tips
|
||||
|
||||
- To hide the model fetching status progress bars, `export HF_HUB_DISABLE_PROGRESS_BARS=1`
|
||||
@ -1654,9 +1761,13 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
|
||||
- When generating multiple images with different seeds, use `--seed` with multiple values or `--auto-seeds` to automatically generate a series of random seeds
|
||||
- Use `--stepwise-image-output-dir` to save intermediate images at each denoising step, which can be useful for debugging or creating animations of the generation process
|
||||
|
||||
### 🔬 Cool research / features to support
|
||||
---
|
||||
|
||||
### 🔬 Cool research
|
||||
- [ ] [PuLID](https://github.com/ToTheBeginning/PuLID)
|
||||
|
||||
---
|
||||
|
||||
### 🌱 Related projects
|
||||
|
||||
- [Mflux-ComfyUI](https://github.com/raysers/Mflux-ComfyUI) by [@raysers](https://github.com/raysers)
|
||||
@ -1664,6 +1775,8 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
|
||||
- [mflux-fasthtml](https://github.com/anthonywu/mflux-fasthtml) by [@anthonywu](https://github.com/anthonywu)
|
||||
- [mflux-streamlit](https://github.com/elitexp/mflux-streamlit) by [@elitexp](https://github.com/elitexp)
|
||||
|
||||
---
|
||||
|
||||
### 🙏 Acknowledgements
|
||||
|
||||
MFLUX would not be possible without the great work of:
|
||||
@ -1674,6 +1787,8 @@ MFLUX would not be possible without the great work of:
|
||||
- Depth Pro authors for the [Depth Pro model](https://github.com/apple/ml-depth-pro?tab=readme-ov-file#citation)
|
||||
- The MLX community and all [contributors and testers](https://github.com/filipstrand/mflux/graphs/contributors)
|
||||
|
||||
---
|
||||
|
||||
### ⚖️ License
|
||||
|
||||
This project is licensed under the [MIT License](LICENSE).
|
||||
|
||||
@ -7,6 +7,7 @@ import PIL.Image
|
||||
|
||||
from mflux.models.vae.vae import VAE
|
||||
from mflux.post_processing.array_util import ArrayUtil
|
||||
from mflux.post_processing.image_util import StrOrBytesPath
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
@ -17,7 +18,7 @@ class ControlnetUtil:
|
||||
vae: VAE,
|
||||
height: int,
|
||||
width: int,
|
||||
controlnet_image_path: str,
|
||||
controlnet_image_path: StrOrBytesPath,
|
||||
is_canny: bool,
|
||||
) -> tuple[mx.array, PIL.Image.Image]:
|
||||
from mflux import ImageUtil
|
||||
@ -34,7 +35,7 @@ class ControlnetUtil:
|
||||
return controlnet_cond, control_image
|
||||
|
||||
@staticmethod
|
||||
def _preprocess_canny(img: PIL.Image) -> PIL.Image:
|
||||
def _preprocess_canny(img: PIL.Image.Image) -> PIL.Image.Image:
|
||||
image_to_canny = np.array(img)
|
||||
image_to_canny = cv2.Canny(image_to_canny, 100, 200)
|
||||
image_to_canny = np.array(image_to_canny[:, :, None])
|
||||
@ -42,8 +43,10 @@ class ControlnetUtil:
|
||||
return PIL.Image.fromarray(image_to_canny)
|
||||
|
||||
@staticmethod
|
||||
def _scale_image(height: int, width: int, img: PIL.Image) -> PIL.Image:
|
||||
def _scale_image(height: int, width: int, img: PIL.Image.Image) -> PIL.Image.Image:
|
||||
if height != img.height or width != img.width:
|
||||
log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}")
|
||||
log.warning(
|
||||
f"Control image {img.width}x{img.height} has different dimensions than the model requirements or requested width x height. Resizing to {width}x{height}"
|
||||
)
|
||||
img = img.resize((width, height), PIL.Image.LANCZOS)
|
||||
return img
|
||||
|
||||
@ -18,7 +18,7 @@ from mflux.models.transformer.transformer import Transformer
|
||||
from mflux.models.vae.vae import VAE
|
||||
from mflux.post_processing.array_util import ArrayUtil
|
||||
from mflux.post_processing.generated_image import GeneratedImage
|
||||
from mflux.post_processing.image_util import ImageUtil
|
||||
from mflux.post_processing.image_util import ImageUtil, StrOrBytesPath
|
||||
from mflux.weights.model_saver import ModelSaver
|
||||
|
||||
|
||||
@ -52,7 +52,7 @@ class Flux1Controlnet(nn.Module):
|
||||
self,
|
||||
seed: int,
|
||||
prompt: str,
|
||||
controlnet_image_path: str,
|
||||
controlnet_image_path: StrOrBytesPath,
|
||||
config: Config,
|
||||
) -> GeneratedImage:
|
||||
# 0. Create a new runtime config based on the model type and input parameters
|
||||
|
||||
@ -13,7 +13,7 @@ def main():
|
||||
parser.add_model_arguments(require_model_arg=True)
|
||||
parser.add_lora_arguments()
|
||||
parser.add_image_generator_arguments(supports_metadata_config=False)
|
||||
parser.add_controlnet_arguments()
|
||||
parser.add_controlnet_arguments(mode="canny")
|
||||
parser.add_output_arguments()
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@ -7,6 +7,7 @@ import numpy as np
|
||||
import piexif
|
||||
import PIL.Image
|
||||
import PIL.ImageDraw
|
||||
from PIL._typing import StrOrBytesPath
|
||||
|
||||
from mflux.community.concept_attention.attention_data import ConceptHeatmap
|
||||
from mflux.config.runtime_config import RuntimeConfig
|
||||
@ -117,8 +118,11 @@ class ImageUtil:
|
||||
return array
|
||||
|
||||
@staticmethod
|
||||
def load_image(path: str | Path) -> PIL.Image.Image:
|
||||
return PIL.Image.open(path).convert("RGB")
|
||||
def load_image(image_or_path: PIL.Image.Image | StrOrBytesPath) -> PIL.Image.Image:
|
||||
if isinstance(image_or_path, PIL.Image.Image):
|
||||
return image_or_path.convert("RGB")
|
||||
else:
|
||||
return PIL.Image.open(image_or_path).convert("RGB")
|
||||
|
||||
@staticmethod
|
||||
def expand_image(
|
||||
|
||||
@ -9,6 +9,7 @@ from mflux.community.in_context.utils.in_context_loras import LORA_NAME_MAP, LOR
|
||||
from mflux.ui import (
|
||||
box_values,
|
||||
defaults as ui_defaults,
|
||||
scale_factor,
|
||||
)
|
||||
from mflux.weights.lora_library import get_lora_path
|
||||
|
||||
@ -29,6 +30,25 @@ class ModelSpecAction(argparse.Action):
|
||||
setattr(namespace, self.dest, values)
|
||||
|
||||
|
||||
def int_or_special_value(value) -> int | scale_factor.ScaleFactor:
|
||||
if value.lower() == "auto":
|
||||
return scale_factor.ScaleFactor(value=1)
|
||||
|
||||
# Try to parse as integer first
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# If not an integer, try to parse as scale factor
|
||||
try:
|
||||
return scale_factor.parse_scale_factor(value)
|
||||
except ValueError:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"'{value}' is not a valid integer or 'auto' or a scale factor like '2x' or '3.5x'"
|
||||
)
|
||||
|
||||
|
||||
# fmt: off
|
||||
class CommandLineParser(argparse.ArgumentParser):
|
||||
|
||||
@ -37,6 +57,7 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
self.supports_metadata_config = False
|
||||
self.supports_image_generation = False
|
||||
self.supports_controlnet = False
|
||||
self.supports_dimension_scale_factor = False
|
||||
self.supports_image_to_image = False
|
||||
self.supports_image_outpaint = False
|
||||
self.supports_lora = False
|
||||
@ -67,20 +88,26 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
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) -> None:
|
||||
def _add_image_generator_common_arguments(self, supports_dimension_scale_factor=False) -> None:
|
||||
self.supports_image_generation = True
|
||||
self.add_argument("--height", type=int, default=ui_defaults.HEIGHT, help=f"Image height (Default is {ui_defaults.HEIGHT})")
|
||||
self.add_argument("--width", type=int, default=ui_defaults.WIDTH, help=f"Image width (Default is {ui_defaults.HEIGHT})")
|
||||
if supports_dimension_scale_factor:
|
||||
self.supports_dimension_scale_factor = True
|
||||
self.add_argument("--height", type=int_or_special_value, default="auto", help="Image height (Default is source image height)")
|
||||
self.add_argument("--width", type=int_or_special_value, default="auto", help="Image width (Default is source image width)")
|
||||
else:
|
||||
self.add_argument("--height", type=int, default=ui_defaults.HEIGHT, help=f"Image height (Default is {ui_defaults.HEIGHT})")
|
||||
self.add_argument("--width", type=int, default=ui_defaults.WIDTH, help=f"Image width (Default is {ui_defaults.HEIGHT})")
|
||||
|
||||
self.add_argument("--steps", type=int, default=None, help="Inference Steps")
|
||||
self.add_argument("--guidance", type=float, default=None, help=f"Guidance Scale (Default varies by tool: {ui_defaults.GUIDANCE_SCALE} for most, {ui_defaults.DEFAULT_DEV_FILL_GUIDANCE} for fill tools, {ui_defaults.DEFAULT_DEPTH_GUIDANCE} for depth)")
|
||||
|
||||
def add_image_generator_arguments(self, supports_metadata_config=False, require_prompt=True) -> None:
|
||||
def add_image_generator_arguments(self, supports_metadata_config=False, require_prompt=True, supports_dimension_scale_factor=False) -> None:
|
||||
prompt_group = self.add_mutually_exclusive_group(required=(require_prompt and not supports_metadata_config))
|
||||
prompt_group.add_argument("--prompt", type=str, help="The textual description of the image to generate.")
|
||||
prompt_group.add_argument("--prompt-file", type=Path, help="Path to a file containing the prompt text. The file will be re-read before each generation, allowing you to edit the prompt between iterations when using multiple seeds without restarting the program.")
|
||||
self.add_argument("--seed", type=int, default=None, nargs='+', help="Specify 1+ Entropy Seeds (Default is 1 time-based random-seed)")
|
||||
self.add_argument("--auto-seeds", type=int, default=-1, help="Auto generate N Entropy Seeds (random ints between 0 and 1 billion")
|
||||
self._add_image_generator_common_arguments()
|
||||
self._add_image_generator_common_arguments(supports_dimension_scale_factor=supports_dimension_scale_factor)
|
||||
if supports_metadata_config:
|
||||
self.add_metadata_config()
|
||||
self.require_prompt = require_prompt
|
||||
@ -134,11 +161,12 @@ class CommandLineParser(argparse.ArgumentParser):
|
||||
self.supports_image_outpaint = True
|
||||
self.add_argument("--image-outpaint-padding", type=str, default=None, required=required, help="For outpainting mode: CSS-style box padding values to extend the canvas of image specified by--image-path. E.g. '20', '50%%'")
|
||||
|
||||
def add_controlnet_arguments(self) -> None:
|
||||
def add_controlnet_arguments(self, mode: str | None = None, require_image=False) -> None:
|
||||
self.supports_controlnet = True
|
||||
self.add_argument("--controlnet-image-path", type=str, required=False, help="Local path of the image to use as input for controlnet.")
|
||||
self.add_argument("--controlnet-image-path", type=str, required=require_image, help="Local path of the image to use as input for controlnet.")
|
||||
self.add_argument("--controlnet-strength", type=float, default=ui_defaults.CONTROLNET_STRENGTH, help=f"Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is {ui_defaults.CONTROLNET_STRENGTH})")
|
||||
self.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.")
|
||||
if mode == 'canny':
|
||||
self.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.")
|
||||
|
||||
def add_concept_attention_arguments(self) -> None:
|
||||
concept_group = self.add_argument_group("Concept Attention configuration")
|
||||
|
||||
@ -11,8 +11,10 @@ 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
|
||||
HEIGHT, WIDTH = 1024, 1024
|
||||
MAX_PIXELS_WARNING_THRESHOLD = 2048 * 2048
|
||||
IMAGE_STRENGTH = 0.4
|
||||
MODEL_CHOICES = ["dev", "dev-fill", "schnell"]
|
||||
MODEL_INFERENCE_STEPS = {
|
||||
|
||||
48
src/mflux/ui/scale_factor.py
Normal file
48
src/mflux/ui/scale_factor.py
Normal file
@ -0,0 +1,48 @@
|
||||
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)
|
||||
|
||||
|
||||
# Regex pattern for scale factors
|
||||
SCALE_FACTOR_PATTERN = re.compile(
|
||||
r"^(\d+(?:\.\d+)?)x$", # Matches integer or decimal followed by 'x'
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
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.")
|
||||
|
||||
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)
|
||||
@ -1,8 +1,14 @@
|
||||
import sys
|
||||
|
||||
import PIL.Image
|
||||
|
||||
from mflux import Config, Flux1Controlnet, ModelConfig, StopImageGenerationException
|
||||
from mflux.callbacks.callback_manager import CallbackManager
|
||||
from mflux.error.exceptions import PromptFileReadError
|
||||
from mflux.ui import defaults as ui_defaults
|
||||
from mflux.ui.cli.parsers import CommandLineParser
|
||||
from mflux.ui.prompt_utils import get_effective_prompt
|
||||
from mflux.ui.scale_factor import ScaleFactor
|
||||
|
||||
|
||||
def main():
|
||||
@ -11,8 +17,8 @@ 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=False)
|
||||
parser.add_controlnet_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()
|
||||
|
||||
@ -29,6 +35,9 @@ def main():
|
||||
memory_saver = CallbackManager.register_callbacks(args=args, flux=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(
|
||||
@ -37,8 +46,8 @@ def main():
|
||||
controlnet_image_path=args.controlnet_image_path,
|
||||
config=Config(
|
||||
num_inference_steps=args.steps,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
height=height,
|
||||
width=width,
|
||||
controlnet_strength=args.controlnet_strength,
|
||||
),
|
||||
)
|
||||
@ -52,5 +61,42 @@ def main():
|
||||
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)
|
||||
else:
|
||||
output_height = args.height
|
||||
|
||||
if isinstance(args.width, ScaleFactor):
|
||||
output_width: int = args.width.get_scaled_value(orig_image.width)
|
||||
else:
|
||||
output_width = args.width
|
||||
|
||||
# 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()
|
||||
|
||||
@ -19,7 +19,7 @@ def _create_mflux_generate_parser(with_controlnet=False, require_model_arg=False
|
||||
parser.add_image_to_image_arguments(required=False)
|
||||
parser.add_image_outpaint_arguments()
|
||||
if with_controlnet:
|
||||
parser.add_controlnet_arguments()
|
||||
parser.add_controlnet_arguments(mode="canny")
|
||||
parser.add_output_arguments()
|
||||
return parser
|
||||
|
||||
|
||||
194
tests/arg_parser/test_upscale_argparser.py
Normal file
194
tests/arg_parser/test_upscale_argparser.py
Normal file
@ -0,0 +1,194 @@
|
||||
from pathlib import Path
|
||||
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
|
||||
|
||||
|
||||
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)
|
||||
parser.add_lora_arguments()
|
||||
|
||||
# Manually add the image generator arguments with scale factor support
|
||||
prompt_group = parser.add_mutually_exclusive_group(required=False)
|
||||
prompt_group.add_argument("--prompt", type=str, help="The textual description of the image to generate.")
|
||||
prompt_group.add_argument("--prompt-file", type=Path, help="Path to a file containing the prompt text.")
|
||||
parser.add_argument("--seed", type=int, default=None, nargs="+", help="Specify 1+ Entropy Seeds")
|
||||
parser.add_argument("--auto-seeds", type=int, default=-1, help="Auto generate N Entropy Seeds")
|
||||
|
||||
# Add height/width with scale factor support
|
||||
parser.supports_image_generation = True
|
||||
parser.supports_dimension_scale_factor = True
|
||||
parser.add_argument(
|
||||
"--height", type=int_or_special_value, default="auto", help="Image height (Default is source image height)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width", type=int_or_special_value, default="auto", help="Image width (Default is source image width)"
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=None, help="Inference Steps")
|
||||
parser.add_argument("--guidance", type=float, default=None, help="Guidance Scale")
|
||||
|
||||
parser.add_controlnet_arguments()
|
||||
parser.add_output_arguments()
|
||||
return parser
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mflux_upscale_parser() -> CommandLineParser:
|
||||
return _create_custom_upscale_parser()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mflux_upscale_minimal_argv() -> list[str]:
|
||||
return ["mflux-upscale", "--prompt", "upscaled image", "--controlnet-image-path", "image.png"]
|
||||
|
||||
|
||||
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)
|
||||
assert args.height.value == 1
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 1
|
||||
|
||||
|
||||
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()
|
||||
assert isinstance(args.height, ScaleFactor)
|
||||
assert args.height.value == 2
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 3
|
||||
|
||||
# Test float scale factor
|
||||
with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "1.5x", "--width", "2.5x"]):
|
||||
args = mflux_upscale_parser.parse_args()
|
||||
assert isinstance(args.height, ScaleFactor)
|
||||
assert args.height.value == 1.5
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 2.5
|
||||
|
||||
# Test decimal scale factor
|
||||
with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "3.14x", "--width", "0.5x"]):
|
||||
args = mflux_upscale_parser.parse_args()
|
||||
assert isinstance(args.height, ScaleFactor)
|
||||
assert args.height.value == 3.14
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 0.5
|
||||
|
||||
|
||||
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)
|
||||
assert args.height == 1024
|
||||
assert isinstance(args.width, int)
|
||||
assert args.width == 768
|
||||
|
||||
|
||||
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)
|
||||
assert args.height.value == 2
|
||||
assert isinstance(args.width, int)
|
||||
assert args.width == 1024
|
||||
|
||||
with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "768", "--width", "1.5x"]):
|
||||
args = mflux_upscale_parser.parse_args()
|
||||
assert isinstance(args.height, int)
|
||||
assert args.height == 768
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 1.5
|
||||
|
||||
|
||||
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)
|
||||
assert isinstance(args.height, ScaleFactor)
|
||||
assert args.height.value == 1
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 1
|
||||
|
||||
|
||||
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):
|
||||
mflux_upscale_parser.parse_args()
|
||||
|
||||
# Invalid format with multiple 'x'
|
||||
with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "2xx"]):
|
||||
with pytest.raises(SystemExit):
|
||||
mflux_upscale_parser.parse_args()
|
||||
|
||||
# Invalid non-numeric value
|
||||
with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "abcx"]):
|
||||
with pytest.raises(SystemExit):
|
||||
mflux_upscale_parser.parse_args()
|
||||
|
||||
# Invalid empty value before 'x'
|
||||
with patch("sys.argv", mflux_upscale_minimal_argv + ["--height", "x"]):
|
||||
with pytest.raises(SystemExit):
|
||||
mflux_upscale_parser.parse_args()
|
||||
|
||||
|
||||
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)
|
||||
assert args.height.value == 2
|
||||
assert isinstance(args.width, ScaleFactor)
|
||||
assert args.width.value == 1.5
|
||||
|
||||
|
||||
def test_upscale_with_all_arguments(mflux_upscale_parser):
|
||||
"""Test upscale parser with all arguments"""
|
||||
full_argv = [
|
||||
"mflux-upscale",
|
||||
"--prompt",
|
||||
"upscaled beautiful landscape",
|
||||
"--controlnet-image-path",
|
||||
"source.png",
|
||||
"--height",
|
||||
"2x",
|
||||
"--width",
|
||||
"1920",
|
||||
"--steps",
|
||||
"20",
|
||||
"--guidance",
|
||||
"7.5",
|
||||
"--controlnet-strength",
|
||||
"0.8",
|
||||
"--seed",
|
||||
"42",
|
||||
"--output",
|
||||
"upscaled.png",
|
||||
]
|
||||
with patch("sys.argv", full_argv):
|
||||
args = mflux_upscale_parser.parse_args()
|
||||
assert args.prompt == "upscaled beautiful landscape"
|
||||
assert args.controlnet_image_path == "source.png"
|
||||
assert isinstance(args.height, ScaleFactor)
|
||||
assert args.height.value == 2
|
||||
assert isinstance(args.width, int)
|
||||
assert args.width == 1920
|
||||
assert args.steps == 20
|
||||
assert args.guidance == 7.5
|
||||
assert args.controlnet_strength == 0.8
|
||||
assert args.seed == [42]
|
||||
assert args.output == "upscaled.png"
|
||||
70
tests/image_generation/test_upscale_dimensions.py
Normal file
70
tests/image_generation/test_upscale_dimensions.py
Normal file
@ -0,0 +1,70 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from mflux.ui.scale_factor import ScaleFactor
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args_height,args_width,orig_height,orig_width,expected_height,expected_width",
|
||||
[
|
||||
# ScaleFactor dimensions
|
||||
(ScaleFactor(value=2), ScaleFactor(value=1.5), 768, 512, 1536, 768),
|
||||
# Integer dimensions
|
||||
(1024, 768, 512, 512, 1024, 768),
|
||||
# Mixed: ScaleFactor height, integer width
|
||||
(ScaleFactor(value=2.5), 1280, 480, 640, 1200, 1280),
|
||||
# Auto (ScaleFactor with value 1)
|
||||
(ScaleFactor(value=1), ScaleFactor(value=1), 512, 1024, 512, 1024),
|
||||
],
|
||||
)
|
||||
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_image.size = (orig_width, orig_height)
|
||||
mock_image.height = orig_height
|
||||
mock_image.width = orig_width
|
||||
|
||||
# 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
|
||||
|
||||
# Mock command line args
|
||||
mock_args = Mock()
|
||||
mock_args.height = args_height
|
||||
mock_args.width = args_width
|
||||
mock_args.controlnet_image_path = "test.png"
|
||||
mock_args.seed = [42]
|
||||
mock_args.prompt = "test prompt"
|
||||
mock_args.steps = 20
|
||||
mock_args.controlnet_strength = 0.4
|
||||
mock_args.quantize = None
|
||||
mock_args.path = None
|
||||
mock_args.lora_paths = None
|
||||
mock_args.lora_scales = None
|
||||
|
||||
with patch("mflux.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.get_effective_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
|
||||
119
tests/ui/test_scale_factor.py
Normal file
119
tests/ui/test_scale_factor.py
Normal file
@ -0,0 +1,119 @@
|
||||
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
|
||||
Loading…
Reference in New Issue
Block a user