From b26a1d66b39c143d7ed22b867d01447e772af37e Mon Sep 17 00:00:00 2001
From: Anthony Wu <462072+anthonywu@users.noreply.github.com>
Date: Fri, 27 Jun 2025 10:51:00 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=96=BC=EF=B8=8F=E2=86=97=EF=B8=8F=20Add?=
 =?UTF-8?q?=20scale=20factor=20support=20for=20image=20upscaling=20dimensi?=
 =?UTF-8?q?ons=20(#215)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>
Co-authored-by: filipstrand <strand.filip@gmail.com>
---
 README.md                                     | 255 +++++++++++++-----
 src/mflux/controlnet/controlnet_util.py       |  11 +-
 src/mflux/controlnet/flux_controlnet.py       |   4 +-
 src/mflux/generate_controlnet.py              |   2 +-
 src/mflux/post_processing/image_util.py       |   8 +-
 src/mflux/ui/cli/parsers.py                   |  44 ++-
 src/mflux/ui/defaults.py                      |   2 +
 src/mflux/ui/scale_factor.py                  |  48 ++++
 src/mflux/upscale.py                          |  54 +++-
 tests/arg_parser/test_cli_argparser.py        |   2 +-
 tests/arg_parser/test_upscale_argparser.py    | 194 +++++++++++++
 .../test_upscale_dimensions.py                |  70 +++++
 tests/ui/test_scale_factor.py                 | 119 ++++++++
 13 files changed, 721 insertions(+), 92 deletions(-)
 create mode 100644 src/mflux/ui/scale_factor.py
 create mode 100644 tests/arg_parser/test_upscale_argparser.py
 create mode 100644 tests/image_generation/test_upscale_dimensions.py
 create mode 100644 tests/ui/test_scale_factor.py

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)
 
 <!-- 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).
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