diff --git a/README.md b/README.md index a5d216f..b456fec 100644 --- a/README.md +++ b/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) +--- + ### 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) +
+⌨️ Shell Completions (Optional) -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). +
+ +--- + ### πŸ–ΌοΈ 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 +
+πŸ“œ Command-Line Arguments Reference + +#### General Command-Line Arguments + +
+Click to expand General arguments + - **`--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 +
+ +#### In-Context LoRA Command-Line Arguments + +
+Click to expand In-Context LoRA arguments 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 +
+ +#### CatVTON Command-Line Arguments + +
+Click to expand CatVTON arguments 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 +
+ +#### IC-Edit Command-Line Arguments + +
+Click to expand IC-Edit arguments 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 +
+ +#### Redux Tool Command-Line Arguments + +
+Click to expand Redux arguments 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 +
+ +#### Concept Attention Command-Line Arguments + +
+Click to expand Concept Attention arguments 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 +
+ +#### Fill Tool Command-Line Arguments + +
+Click to expand Fill Tool arguments 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 +
+ +#### Depth Tool Command-Line Arguments + +
+Click to expand Depth Tool arguments 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 +
+ +#### ControlNet Command-Line Arguments + +
+Click to expand ControlNet arguments 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. +
-#### Dynamic Prompts with `--prompt-file` +#### Upscale Command-Line Arguments + +
+Click to expand Upscale arguments + +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. + +
+ +#### Training Arguments + +
+Click to expand 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. + +
+ +
+ +
+πŸ“ Dynamic Prompts with `--prompt-file` 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. - +
-parameters supported by config files +βš™οΈ Parameters supported by config files #### How configs are used @@ -474,32 +583,7 @@ mflux-generate --prompt-file my_prompt.txt --auto-seeds 10 ```
-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 \ +--- + ### πŸ’½ 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: +
+πŸ“ Required directory structure + ``` . β”œβ”€β”€ 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 ``` + +
+ 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 +
+βš™οΈ Configuration details 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 +
+ +
+⚠️ Memory issues 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 +
+ +
+πŸ“ Misc 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) +
+ --- @@ -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). diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index d0d1d8f..5bd3330 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -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 diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index c3fd8d8..3a29d02 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -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 diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index 778b1f6..7c843d1 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -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() diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 9164e75..c465108 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -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( diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 4275a1f..a56a27b 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/src/mflux/ui/cli/parsers.py @@ -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") diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py index 4441d56..305e564 100644 --- a/src/mflux/ui/defaults.py +++ b/src/mflux/ui/defaults.py @@ -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 = { diff --git a/src/mflux/ui/scale_factor.py b/src/mflux/ui/scale_factor.py new file mode 100644 index 0000000..711ec22 --- /dev/null +++ b/src/mflux/ui/scale_factor.py @@ -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) diff --git a/src/mflux/upscale.py b/src/mflux/upscale.py index a730a0c..8c98033 100644 --- a/src/mflux/upscale.py +++ b/src/mflux/upscale.py @@ -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() diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index 33a3359..6a6b82f 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/tests/arg_parser/test_cli_argparser.py @@ -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 diff --git a/tests/arg_parser/test_upscale_argparser.py b/tests/arg_parser/test_upscale_argparser.py new file mode 100644 index 0000000..1877775 --- /dev/null +++ b/tests/arg_parser/test_upscale_argparser.py @@ -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" diff --git a/tests/image_generation/test_upscale_dimensions.py b/tests/image_generation/test_upscale_dimensions.py new file mode 100644 index 0000000..47424ee --- /dev/null +++ b/tests/image_generation/test_upscale_dimensions.py @@ -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 diff --git a/tests/ui/test_scale_factor.py b/tests/ui/test_scale_factor.py new file mode 100644 index 0000000..0964671 --- /dev/null +++ b/tests/ui/test_scale_factor.py @@ -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