Release 0.11.1 (#278)

Co-authored-by: Filip Strand <filip@Host-022.local>
2025-11-13 01:05:29 +01:00 · 2025-11-13 01:05:29 +01:00 · 2e92984125
commit 2e92984125
parent d4005efac0
22 changed files with 1787 additions and 1624 deletions
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@ -5,6 +5,34 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

+## [0.11.1] - 2025-11-11
+
+# MFLUX v.0.11.1 Release Notes
+
+### 🎨 New Model Support
+
+- **Qwen Image Edit Support**: Added support for the Qwen Image Edit model, enabling natural language image editing capabilities
+- **New command**: `mflux-generate-qwen-edit` for image editing with text instructions
+- **Multiple image support**: Edit images using multiple reference images via `--image-paths` parameter
+- **Model**: Uses `Qwen/Qwen-Image-Edit-2509` for high-quality image editing
+- **Quantization support**: Full support for quantized models (8-bit recommended for optimal quality)
+
+### 🔧 Improvements
+
+- **Dedicated Qwen Image command**: Added `mflux-generate-qwen` as a dedicated command for Qwen Image model generation. The `mflux-generate` command now only supports Flux models.
+- **Image comparison utility refactoring**: Refactored `image_compare.py` into a cleaner class-based structure with static methods
+- **Error handling**: Moved `ReferenceVsOutputImageError` to the main exceptions module for better organization
+
+### 🔄 Breaking Changes
+
+⚠️ **Qwen Image Command Change**: The Qwen Image model now requires using the dedicated `mflux-generate-qwen` command instead of `mflux-generate --model qwen`. This provides better separation between Flux and Qwen model families and improves command clarity.
+
+### 👩‍💻 Contributors
+
+- **Filip Strand (@filipstrand)**: Qwen Image Edit model implementation, code refactoring
+
+---
+
 ## [0.11.0] - 2025-10-14

 # MFLUX v.0.11.0 Release Notes
--- a/README.md
+++ b/README.md
@ -4,7 +4,7 @@

 ### About

-Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black Forest Labs](https://blackforestlabs.ai) locally on your Mac!
+Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) and [Qwen Image](https://github.com/QwenLM/Qwen-Image) models locally on your Mac!

 ### Table of contents

@ -20,6 +20,9 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
 - [💽 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)
+- [🦙 Qwen Models](#-qwen-models)
+  * [🖼️ Qwen Image](#%EF%B8%8F-qwen-image)
+  * [✏️ Qwen Image Edit](#%EF%B8%8F-qwen-image-edit)
 - [🔌 LoRA](#-lora)
 - [🎭 In-Context Generation](#-in-context-generation)
  * [📸 Kontext](#-kontext)
@ -47,16 +50,13 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black

 ### 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).
+MFLUX is a line-by-line port of the FLUX and Qwen implementations in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) and [Huggingface Transformers](https://github.com/huggingface/transformers) libraries to [Apple MLX](https://github.com/ml-explore/mlx).
 MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
 except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models
 (thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#%EF%B8%8F-image-generation-speed-updated), [and even faster quantized](#%EF%B8%8F-quantization).

-All models are implemented from scratch in MLX and only the tokenizers are used via the
-[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.
+All models are implemented from scratch in MLX and only the tokenizers are used via the [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.

-As of v.0.11.0, MFLUX now supports the Qwen Image model.
 ---

 ### 💿 Installation
@ -156,7 +156,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2
 This example uses the `qwen` model with 20 time steps:

 ```sh
-mflux-generate --model qwen --prompt "Luxury food photograph" --steps 20 --seed 2 -q 6
+mflux-generate-qwen --prompt "Luxury food photograph" --steps 20 --seed 2 -q 6
 ```

 You can also pipe prompts from other commands using stdin:
@ -175,7 +175,7 @@ from mflux.config.config import Config

 # Load the model
 flux = Flux1.from_name(
-   model_name="schnell",  # "schnell", "dev", "krea-dev", or "qwen"
+   model_name="schnell",  # "schnell", "dev", or "krea-dev"
   quantize=8,            # 3, 4, 5, 6, or 8
 )

@ -184,7 +184,7 @@ image = flux.generate_image(
   seed=2,
   prompt="Luxury food photograph",
   config=Config(
-      num_inference_steps=2,  # "schnell" works well with 2-4 steps, "dev", "krea-dev", and "qwen" work well with 20-25 steps
+      num_inference_steps=2,  # "schnell" works well with 2-4 steps, "dev" and "krea-dev" work well with 20-25 steps
      height=1024,
      width=1024,
   )
@ -235,9 +235,9 @@ mflux-generate \

 - **`--prompt`** (required, `str`): Text description of the image to generate. Use `-` to read the prompt from stdin (e.g., `echo "A beautiful sunset" | mflux-generate --prompt -`).

- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official models (`"schnell"`, `"dev"`, `"krea-dev"`, or `"qwen"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`).
+- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official Flux models (`"schnell"`, `"dev"`, or `"krea-dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`). For Qwen models, use `mflux-generate-qwen` instead.

- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"`, `"dev"`, `"krea-dev"`, or `"qwen"`). Required when using third-party models from HuggingFace.
+- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"`, `"dev"`, or `"krea-dev"`). Required when using third-party models from HuggingFace.

 - **`--output`** (optional, `str`, default: `"image.png"`): Output image filename. If `--seed` or `--auto-seeds` establishes multiple seed values, the output filename will automatically be modified to include the seed value (e.g., `image_seed_42.png`).

@ -251,7 +251,7 @@ mflux-generate \

 - **`--steps`** (optional, `int`, default: `4`): Number of inference steps.

- **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale (only used for `"dev"`, `"krea-dev"` and `qwen` models).
+- **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale (only used for `"dev"` and `"krea-dev"` models).

 - **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.

@ -355,6 +355,40 @@ See the [IC-Edit (In-Context Editing)](#-ic-edit-in-context-editing) section for

 </details>

+#### Qwen Image Command-Line Arguments
+
+<details>
+<summary>Click to expand Qwen Image arguments</summary>
+
+The `mflux-generate-qwen` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
+
+- **`--model`** (optional, `str`, default: `"qwen-image"`): Model to use for generation. Defaults to `"qwen-image"` if not specified.
+
+- **`--guidance`** (optional, `float`, default: `3.5`): Guidance scale for the generation. Qwen Image typically works well with values around 3.5.
+
+**Note**: The Qwen Image tool automatically uses the Qwen Image model, so you typically don't need to specify `--model`.
+
+See the [Qwen Image](#-qwen-image) section for more details on this feature.
+
+</details>
+
+#### Qwen Image Edit Command-Line Arguments
+
+<details>
+<summary>Click to expand Qwen Image Edit arguments</summary>
+
+The `mflux-generate-qwen-edit` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
+
+- **`--image-paths`** (required, `[str]`): Paths to one or more reference images. For single image editing, provide one path. For multiple image editing (e.g., combining elements from multiple images), provide multiple paths.
+
+- **`--guidance`** (optional, `float`, default: `2.5`): Guidance scale for the generation. Qwen Image Edit typically works well with values around 2.5.
+
+**Note**: The Qwen Image Edit tool automatically uses the appropriate model (`qwen-image-edit`), so you don't need to specify `--model`.
+
+See the [Qwen Image Edit](#-qwen-image-edit) section for more details on this feature.
+
+</details>
+
 #### Redux Tool Command-Line Arguments

 <details>
@ -928,6 +962,131 @@ In the examples above the following LoRAs are used [Sketching](https://civitai.c

 ---

+### 🦙 Qwen Models
+
+MFLUX supports the [Qwen Image](https://github.com/QwenLM/Qwen-Image) family of vision-language models, providing both text-to-image generation and natural language image editing capabilities. Released approximately a year after FLUX, Qwen models achieve state-of-the-art performance in most areas, though they are comparatively heavier to run.
+
+#### 🖼️ Qwen Image
+
+**Qwen Image** is a powerful 20B parameter text-to-image model ([technical report](https://arxiv.org/abs/2508.02324)) that enables high-quality image generation from text prompts. It uses a vision-language architecture with a 7B text encoder (Qwen2.5-VL) to understand and generate images based on natural language descriptions.
+
+The Qwen Image model has its own dedicated command `mflux-generate-qwen`. Qwen Image excels at multilingual prompts, including Chinese characters, and can render Chinese text as part of the image content (like signs, menus, and calligraphy).
+
+![Qwen Image Examples](src/mflux/assets/qwen_image_example.jpg)
+
+**Example: Wildlife Portrait**
+
+```sh
+mflux-generate-qwen \
+  --prompt "Close-up portrait of a majestic tiger in its natural habitat, detailed fur texture, piercing eyes, natural forest background, soft natural lighting, wildlife photography, photorealistic, high detail, professional wildlife shot" \
+  --negative-prompt "blurry, low quality, distorted, deformed, ugly, bad anatomy, bad proportions, extra limbs, duplicate, watermark, signature, text, letters, cartoon, anime, painting, drawing, illustration, 3d render, cgi, zoo, cage, artificial" \
+  --width 1920 \
+  --height 816 \
+  --steps 30 \
+  --seed 42 \
+  -q 8
+```
+
+<details>
+<summary><strong>Click to expand additional example commands - These are the exact commands used to generate the images shown above</strong></summary>
+
+**Chinese Calligraphy:**
+
+```sh
+mflux-generate-qwen \
+  --prompt "Traditional Chinese calligraphy studio, ancient scrolls with beautiful Chinese characters, ink brushes, inkstone, traditional paper, warm natural lighting, peaceful atmosphere, photorealistic, high detail, cultural heritage" \
+  --negative-prompt "blurry, low quality, distorted, deformed, ugly, bad anatomy, bad proportions, extra limbs, duplicate, watermark, signature, text, letters, cartoon, anime, painting, drawing, illustration, 3d render, cgi, modern, digital" \
+  --width 1920 \
+  --height 816 \
+  --steps 30 \
+  --seed 42 \
+  -q 8
+```
+
+**Chinese Street Signs:**
+
+```sh
+mflux-generate-qwen \
+  --prompt "Traditional Chinese street scene, old neighborhood with shop signs displaying Chinese characters (店铺, 餐厅, 书店), red lanterns, narrow alleys, traditional architecture, bustling street life, natural lighting, photorealistic, high detail, street photography" \
+  --negative-prompt "blurry, low quality, distorted, deformed, ugly, bad anatomy, bad proportions, extra limbs, duplicate, watermark, signature, cartoon, anime, painting, drawing, illustration, 3d render, cgi, modern signs, English text only" \
+  --width 1920 \
+  --height 816 \
+  --steps 30 \
+  --seed 42 \
+  -q 8
+```
+
+**Food Photography:**
+
+```sh
+mflux-generate-qwen \
+  --prompt "Professional food photography, gourmet Chinese cuisine, steamed dumplings, colorful vegetables, traditional table setting, restaurant lighting, shallow depth of field, photorealistic, high detail, magazine quality" \
+  --negative-prompt "blurry, low quality, distorted, deformed, ugly, bad anatomy, bad proportions, extra limbs, duplicate, watermark, signature, text, letters, cartoon, anime, painting, drawing, illustration, 3d render, cgi, fast food, unappetizing" \
+  --width 1920 \
+  --height 816 \
+  --steps 30 \
+  --seed 42 \
+  -q 8
+```
+
+</details>
+
+⚠️ *Note: The Qwen Image model requires downloading the `Qwen/Qwen-Image` model weights (~58GB for the full model, or use quantization for smaller sizes).*
+
+#### ✏️ Qwen Image Edit
+
+**Qwen Image Edit** enables precise natural language image editing, allowing you to modify images using text instructions while maintaining their original structure and context. The model uses a vision-language encoder to understand both the input image and your editing instructions, making it ideal for tasks like changing specific elements, adjusting poses, modifying clothing, or altering backgrounds while preserving the overall composition.
+
+Qwen Image Edit supports natural language editing with descriptive text instructions, maintains original poses and body positions when requested, supports multiple images for complex compositions, and works seamlessly with LoRA adapters for specialized transformations like camera angles and styles. The model uses `Qwen/Qwen-Image-Edit-2509`.
+
+![Qwen Image Edit Examples](src/mflux/assets/qwen_edit_example.jpg)
+*Examples showing dog replacement with two-image input and monkey camera angle transformations with LoRAs. Source images: [Golden Retriever](https://images.unsplash.com/photo-1552053831-71594a27632d), [Grey Dog](https://images.unsplash.com/photo-1566710582818-d673dc761201), and [Monkey](https://images.unsplash.com/photo-1578948610588-ffe24448f5ed).*
+
+**Example 1: Two-Image Transformation (Dog Replacement)**
+
+Qwen Image Edit excels at complex transformations using multiple reference images. This example replaces a golden retriever with a grey dog and changes the rose color from white to red:
+
+```sh
+mflux-generate-qwen-edit \
+  --image-paths "dog1.png" "dog2.png" \
+  --prompt "Replace the golden retriever (standing outside, holding white rose) in Image 1 with the grey dog from Image 2 (which is standing inside in a studio). The grey dog should hold a red rose in its mouth and stand outside in the same position as the golden retriever. Maintain the outside environment, background, lighting, and all surroundings completely unchanged." \
+  --steps 30 \
+  --guidance 2.5 \
+  --width 624 \
+  --height 1024
+```
+
+**Example 2: Single Image with LoRAs (Camera Angle Transformations)**
+
+Qwen Image Edit works seamlessly with LoRA adapters for specialized transformations. This example uses two LoRAs to transform camera angles on a single image:
+
+```sh
+mflux-generate-qwen-edit \
+  --image-paths "monkey.png" \
+  --prompt "将镜头极度拉近，使用超长焦镜头进行极端特写拍摄，主体占据画面的大部分空间，背景完全虚化，营造出强烈的视觉冲击力和亲密感。Extreme zoom in with a super telephoto lens, creating an intense close-up where the subject dominates most of the frame, with the background completely blurred, creating a strong visual impact and sense of intimacy." \
+  --steps 8 \
+  --guidance 2.5 \
+  --width 1024 \
+  --height 1024 \
+  --lora-paths "/path/to/Qwen-Image-Lightning-4steps-V2.0.safetensors" "/path/to/镜头转换.safetensors" \
+  --lora-scales 0.5 1.0
+```
+
+*Uses [Qwen Image Lightning LoRA](https://huggingface.co/lightx2v/Qwen-Image-Lightning) for fast generation and [Camera Angle LoRA](https://huggingface.co/dx8152/Qwen-Edit-2509-Multiple-angles) for precise camera control.*
+
+**Tips for Qwen Image Edit:**
+1. **Detailed Prompts**: The model works best with detailed, specific editing instructions
+2. **Pose Maintenance**: Explicitly mention maintaining poses, body positions, or overall stance when you want to preserve the original structure
+3. **Single Focus**: Focus on one or a few related edits at a time for more predictable results
+4. **LoRA Combinations**: Combine multiple LoRAs for complex effects (e.g., fast generation + camera control)
+5. **Quantization**: 6-bit or below can degrade the image a lot more compared to Flux, use with caution
+6. **Seed Variation**: Qwen models typically do not vary much with seed changes. If you want more variation, vary the prompt instead
+7. **Image Quality**: Qwen images come out quite soft compared to Flux models
+
+⚠️ *Note: The Qwen Image Edit model requires downloading the `Qwen/Qwen-Image-Edit-2509` model weights (~58GB for the full model, or use quantization for smaller sizes).*
+
+---
+
 ### 🔌 LoRA

 MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).
@ -1873,7 +2032,7 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
 - Set up shell aliases for required args examples:
  - shortcut for dev model: `alias mflux-dev='mflux-generate --model dev'`
  - shortcut for schnell model *and* always save metadata: `alias mflux-schnell='mflux-generate --model schnell --metadata'`
-  - shortcut for qwen model: `alias mflux-qwen='mflux-generate --model qwen'`
+  - shortcut for qwen model: `alias mflux-qwen='mflux-generate-qwen'`
 - For systems with limited memory, use the `--low-ram` flag to reduce memory usage by constraining the MLX cache size and releasing components after use
 - On battery-powered Macs, use `--battery-percentage-stop-limit` (or `-B`) to prevent your laptop from shutting down during long generation sessions
 - When generating multiple images with different seeds, use `--seed` with multiple values or `--auto-seeds` to automatically generate a series of random seeds
--- a/pyproject.toml
+++ b/pyproject.toml
@ -17,7 +17,7 @@ source-exclude = [

 [project]
 name = "mflux"
-version = "0.11.0"
+version = "0.11.1"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]
@ -84,6 +84,7 @@ mflux-generate-fill = "mflux.generate_fill:main"
 mflux-generate-depth = "mflux.generate_depth:main"
 mflux-generate-redux = "mflux.generate_redux:main"
 mflux-generate-kontext = "mflux.generate_kontext:main"
+mflux-generate-qwen = "mflux.generate_qwen:main"
 mflux-generate-qwen-edit = "mflux.generate_qwen_edit:main"
 mflux-concept = "mflux.concept:main"
 mflux-concept-from-image = "mflux.concept_from_image:main"
@ -146,7 +147,10 @@ docstring-code-line-length = "dynamic"
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 python_files = "test_*.py"
-addopts = "-v --failed-first --showlocals --tb=long --full-trace"
+addopts = '-v --failed-first --showlocals --tb=long --full-trace -m "not high_memory_requirement"'
+markers = [
+    "high_memory_requirement: marks tests that require high memory (deselect with '-m \"not high_memory_requirement\"')",
+]

 # https://docs.astral.sh/ruff/settings/#lintisort
 [tool.ruff.lint.isort]
--- a/src/mflux/assets/qwen_edit_example.jpg
+++ b/src/mflux/assets/qwen_edit_example.jpg
--- a/src/mflux/assets/qwen_image_example.jpg
+++ b/src/mflux/assets/qwen_image_example.jpg
--- a/src/mflux/generate.py
+++ b/src/mflux/generate.py
@ -3,7 +3,6 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
 from mflux.models.flux.variants.txt2img.flux import Flux1
-from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
 from mflux.ui import defaults as ui_defaults
 from mflux.ui.cli.parsers import CommandLineParser
 from mflux.ui.prompt_utils import get_effective_negative_prompt, get_effective_prompt
@ -25,8 +24,7 @@ def main():
        args.guidance = ui_defaults.GUIDANCE_SCALE

    # 1. Load the model
-    model_class = QwenImage if "qwen" in args.model.lower() else Flux1
-    model = model_class(
+    model = Flux1(
        model_config=ModelConfig.from_name(model_name=args.model, base_model=args.base_model),
        quantize=args.quantize,
        local_path=args.path,
--- a/src/mflux/generate_qwen.py
+++ b/src/mflux/generate_qwen.py
@ -0,0 +1,65 @@
+from mflux.callbacks.callback_manager import CallbackManager
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+from mflux.ui import defaults as ui_defaults
+from mflux.ui.cli.parsers import CommandLineParser
+from mflux.ui.prompt_utils import get_effective_negative_prompt, get_effective_prompt
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Generate an image using Qwen Image model.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=True)
+    parser.add_image_to_image_arguments(required=False)
+    parser.add_output_arguments()
+    args = parser.parse_args()
+
+    # 0. Set default guidance value if not provided by user
+    if args.guidance is None:
+        args.guidance = ui_defaults.GUIDANCE_SCALE
+
+    # 1. Load the model
+    qwen = QwenImage(
+        model_config=ModelConfig.from_name(model_name=args.model or "qwen-image", base_model=args.base_model),
+        quantize=args.quantize,
+        local_path=args.path,
+        lora_paths=args.lora_paths,
+        lora_scales=args.lora_scales,
+    )
+
+    # 2. Register callbacks
+    memory_saver = CallbackManager.register_callbacks(args=args, model=qwen)
+
+    try:
+        for seed in args.seed:
+            # 3. Generate an image for each seed value
+            image = qwen.generate_image(
+                seed=seed,
+                prompt=get_effective_prompt(args),
+                negative_prompt=get_effective_negative_prompt(args),
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    image_path=args.image_path,
+                    image_strength=args.image_strength,
+                    scheduler=args.scheduler,
+                ),
+            )
+            # 4. Save the image
+            image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
+    except (StopImageGenerationException, PromptFileReadError) as exc:
+        print(exc)
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
+
+
+if __name__ == "__main__":
+    main()
--- a/src/mflux/generate_qwen_edit.py
+++ b/src/mflux/generate_qwen_edit.py
@ -16,13 +16,12 @@ def main():
    parser.add_model_arguments(require_model_arg=False)
    parser.add_lora_arguments()
    parser.add_image_generator_arguments(supports_metadata_config=True)
-    parser.add_image_to_image_arguments(required=True)
    parser.add_argument(
        "--image-paths",
        type=Path,
        nargs="+",
-        default=None,
-        help="Local paths to multiple init images. If not provided, uses --image-path as single image.",
+        required=True,
+        help="Local paths to one or more init images. For single image editing, provide one path. For multiple image editing, provide multiple paths.",
    )
    parser.add_output_arguments()
    args = parser.parse_args()
@ -44,12 +43,12 @@ def main():

    try:
        for seed in args.seed:
-            # 3. Prepare image paths: use --image-paths if provided, otherwise fallback to --image-path
-            image_paths = None
-            if args.image_paths:
+            # 3. Prepare image paths
            image_paths = [str(p) for p in args.image_paths]
-            elif args.image_path:
-                image_paths = [str(args.image_path)]
+
+            # Use first image path for config.image_path (for backward compatibility with Config)
+            # All image paths are passed to generate_image and stored in metadata
+            config_image_path = image_paths[0]

            # 4. Generate an image for each seed value
            image = qwen.generate_image(
@ -60,7 +59,7 @@ def main():
                    height=args.height,
                    width=args.width,
                    guidance=args.guidance,
-                    image_path=args.image_path,
+                    image_path=config_image_path,
                ),
                negative_prompt=get_effective_negative_prompt(args),
                image_paths=image_paths,
--- a/src/mflux/models/flux/weights/weight_handler.py
+++ b/src/mflux/models/flux/weights/weight_handler.py
@ -217,9 +217,13 @@ class WeightHandler:
                repo_id=repo_id,
                allow_patterns=[
                    "text_encoder/*.safetensors",
+                    "text_encoder/*.json",
                    "text_encoder_2/*.safetensors",
+                    "text_encoder_2/*.json",
                    "transformer/*.safetensors",
+                    "transformer/*.json",
                    "vae/*.safetensors",
+                    "vae/*.json",
                ],
            )
        )
--- a/src/mflux/models/qwen/variants/edit/qwen_image_edit.py
+++ b/src/mflux/models/qwen/variants/edit/qwen_image_edit.py
@ -182,6 +182,7 @@ class QwenImageEdit(nn.Module):
            lora_paths=self.lora_paths,
            lora_scales=self.lora_scales,
            image_path=runtime_config.image_path,
+            image_paths=image_paths,
            generation_time=time_steps.format_dict["elapsed"],
            negative_prompt=negative_prompt,
        )
--- a/src/mflux/models/qwen/weights/qwen_lora_mapping.py
+++ b/src/mflux/models/qwen/weights/qwen_lora_mapping.py
@ -4,7 +4,6 @@ from mflux.models.common.lora.mapping.lora_mapping import LoRAMapping, LoRATarge


 class QwenLoRAMapping(LoRAMapping):
-
    @staticmethod
    def get_mapping() -> List[LoRATarget]:
        return [
@ -14,18 +13,20 @@ class QwenLoRAMapping(LoRAMapping):
                    "transformer_blocks.{block}.attn.to_q.lora_up.weight",
                    "transformer.transformer_blocks.{block}.attn.to_q.lora.up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_q.lora_B.weight",
+                    "transformer_blocks.{block}.attn.to_q.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_q.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.to_q.lora_down.weight",
                    "transformer.transformer_blocks.{block}.attn.to_q.lora.down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_q.lora_A.weight",
+                    "transformer_blocks.{block}.attn.to_q.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_q.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.to_q.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_to_q.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.to_k",
@ -33,18 +34,20 @@ class QwenLoRAMapping(LoRAMapping):
                    "transformer_blocks.{block}.attn.to_k.lora_up.weight",
                    "transformer.transformer_blocks.{block}.attn.to_k.lora.up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_k.lora_B.weight",
+                    "transformer_blocks.{block}.attn.to_k.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_k.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.to_k.lora_down.weight",
                    "transformer.transformer_blocks.{block}.attn.to_k.lora.down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_k.lora_A.weight",
+                    "transformer_blocks.{block}.attn.to_k.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_k.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.to_k.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_to_k.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.to_v",
@ -52,18 +55,20 @@ class QwenLoRAMapping(LoRAMapping):
                    "transformer_blocks.{block}.attn.to_v.lora_up.weight",
                    "transformer.transformer_blocks.{block}.attn.to_v.lora.up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_v.lora_B.weight",
+                    "transformer_blocks.{block}.attn.to_v.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_v.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.to_v.lora_down.weight",
                    "transformer.transformer_blocks.{block}.attn.to_v.lora.down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_v.lora_A.weight",
+                    "transformer_blocks.{block}.attn.to_v.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_v.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.to_v.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_to_v.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.attn_to_out.0",
@ -71,154 +76,171 @@ class QwenLoRAMapping(LoRAMapping):
                    "transformer_blocks.{block}.attn.to_out.0.lora_up.weight",
                    "transformer.transformer_blocks.{block}.attn.to_out.0.lora.up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_out.0.lora_B.weight",
+                    "transformer_blocks.{block}.attn.to_out.0.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_out_0.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.to_out.0.lora_down.weight",
                    "transformer.transformer_blocks.{block}.attn.to_out.0.lora.down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_out.0.lora_A.weight",
+                    "transformer_blocks.{block}.attn.to_out.0.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_out_0.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.to_out.0.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_to_out_0.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.add_q_proj",
                possible_up_patterns=[
                    "transformer_blocks.{block}.attn.add_q_proj.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.add_q_proj.lora_B.weight",
+                    "transformer_blocks.{block}.attn.add_q_proj.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_add_q_proj.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.add_q_proj.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.add_q_proj.lora_A.weight",
+                    "transformer_blocks.{block}.attn.add_q_proj.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_add_q_proj.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.add_q_proj.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_add_q_proj.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.add_k_proj",
                possible_up_patterns=[
                    "transformer_blocks.{block}.attn.add_k_proj.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.add_k_proj.lora_B.weight",
+                    "transformer_blocks.{block}.attn.add_k_proj.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_add_k_proj.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.add_k_proj.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.add_k_proj.lora_A.weight",
+                    "transformer_blocks.{block}.attn.add_k_proj.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_add_k_proj.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.add_k_proj.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_add_k_proj.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.add_v_proj",
                possible_up_patterns=[
                    "transformer_blocks.{block}.attn.add_v_proj.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.add_v_proj.lora_B.weight",
+                    "transformer_blocks.{block}.attn.add_v_proj.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_add_v_proj.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.add_v_proj.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.add_v_proj.lora_A.weight",
+                    "transformer_blocks.{block}.attn.add_v_proj.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_add_v_proj.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.add_v_proj.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_add_v_proj.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.attn.to_add_out",
                possible_up_patterns=[
                    "transformer_blocks.{block}.attn.to_add_out.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_add_out.lora_B.weight",
+                    "transformer_blocks.{block}.attn.to_add_out.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_add_out.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.attn.to_add_out.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.attn.to_add_out.lora_A.weight",
+                    "transformer_blocks.{block}.attn.to_add_out.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_attn_to_add_out.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.attn.to_add_out.alpha",
                    "lora_unet_transformer_blocks_{block}_attn_to_add_out.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.img_ff.mlp_in",
                possible_up_patterns=[
                    "transformer_blocks.{block}.img_mlp.net.0.proj.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.img_mlp.net.0.proj.lora_B.weight",
+                    "transformer_blocks.{block}.img_mlp.net.0.proj.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_img_mlp_net_0_proj.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.img_mlp.net.0.proj.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.img_mlp.net.0.proj.lora_A.weight",
+                    "transformer_blocks.{block}.img_mlp.net.0.proj.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_img_mlp_net_0_proj.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.img_mlp.net.0.proj.alpha",
                    "lora_unet_transformer_blocks_{block}_img_mlp_net_0_proj.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.img_ff.mlp_out",
                possible_up_patterns=[
                    "transformer_blocks.{block}.img_mlp.net.2.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.img_mlp.net.2.lora_B.weight",
+                    "transformer_blocks.{block}.img_mlp.net.2.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_img_mlp_net_2.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.img_mlp.net.2.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.img_mlp.net.2.lora_A.weight",
+                    "transformer_blocks.{block}.img_mlp.net.2.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_img_mlp_net_2.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.img_mlp.net.2.alpha",
                    "lora_unet_transformer_blocks_{block}_img_mlp_net_2.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.txt_ff.mlp_in",
                possible_up_patterns=[
                    "transformer_blocks.{block}.txt_mlp.net.0.proj.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.txt_mlp.net.0.proj.lora_B.weight",
+                    "transformer_blocks.{block}.txt_mlp.net.0.proj.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_txt_mlp_net_0_proj.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.txt_mlp.net.0.proj.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.txt_mlp.net.0.proj.lora_A.weight",
+                    "transformer_blocks.{block}.txt_mlp.net.0.proj.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_txt_mlp_net_0_proj.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.txt_mlp.net.0.proj.alpha",
                    "lora_unet_transformer_blocks_{block}_txt_mlp_net_0_proj.alpha",
-                ]
+                ],
            ),
            LoRATarget(
                model_path="transformer_blocks.{block}.txt_ff.mlp_out",
                possible_up_patterns=[
                    "transformer_blocks.{block}.txt_mlp.net.2.lora_up.weight",
                    "diffusion_model.transformer_blocks.{block}.txt_mlp.net.2.lora_B.weight",
+                    "transformer_blocks.{block}.txt_mlp.net.2.lora_B.default.weight",
                    "lora_unet_transformer_blocks_{block}_txt_mlp_net_2.lora_up.weight",
                ],
                possible_down_patterns=[
                    "transformer_blocks.{block}.txt_mlp.net.2.lora_down.weight",
                    "diffusion_model.transformer_blocks.{block}.txt_mlp.net.2.lora_A.weight",
+                    "transformer_blocks.{block}.txt_mlp.net.2.lora_A.default.weight",
                    "lora_unet_transformer_blocks_{block}_txt_mlp_net_2.lora_down.weight",
                ],
                possible_alpha_patterns=[
                    "transformer_blocks.{block}.txt_mlp.net.2.alpha",
                    "lora_unet_transformer_blocks_{block}_txt_mlp_net_2.alpha",
-                ]
+                ],
            ),
        ]
-
--- a/src/mflux/models/qwen/weights/qwen_weight_handler.py
+++ b/src/mflux/models/qwen/weights/qwen_weight_handler.py
@ -110,7 +110,7 @@ class QwenWeightHandler:

        if loading_mode == "single":
            # VAE style: Single file loading
-            safetensors_files = list(path.glob("*.safetensors"))
+            safetensors_files = [f for f in path.glob("*.safetensors") if not f.name.startswith("._")]
            if not safetensors_files:
                raise FileNotFoundError(f"No safetensors files found in {path}")

@ -180,7 +180,7 @@ class QwenWeightHandler:

    @staticmethod
    def _detect_metadata(path: Path) -> tuple[int | None, str | None]:
-        file_glob = sorted(path.glob("*.safetensors"))
+        file_glob = sorted([f for f in path.glob("*.safetensors") if not f.name.startswith("._")])
        if file_glob:
            data = mx.load(str(file_glob[0]), return_metadata=True)
            if len(data) > 1:
--- a/src/mflux/post_processing/generated_image.py
+++ b/src/mflux/post_processing/generated_image.py
@ -28,6 +28,7 @@ class GeneratedImage:
        controlnet_image_path: str | Path | None = None,
        controlnet_strength: float | None = None,
        image_path: str | Path | None = None,
+        image_paths: list[str] | list[Path] | None = None,
        image_strength: float | None = None,
        masked_image_path: str | Path | None = None,
        depth_image_path: str | Path | None = None,
@ -52,6 +53,7 @@ class GeneratedImage:
        self.controlnet_image_path = controlnet_image_path
        self.controlnet_strength = controlnet_strength
        self.image_path = image_path
+        self.image_paths = image_paths
        self.image_strength = image_strength
        self.masked_image_path = masked_image_path
        self.depth_image_path = depth_image_path
@ -155,7 +157,8 @@ class GeneratedImage:
            "lora_paths": [str(p) for p in self.lora_paths] if self.lora_paths else None,
            "lora_scales": [round(scale, 2) for scale in self.lora_scales] if self.lora_scales else None,
            "image_path": str(self.image_path) if self.image_path else None,
-            "image_strength": self.image_strength if self.image_path else None,
+            "image_paths": [str(p) for p in self.image_paths] if self.image_paths else None,
+            "image_strength": self.image_strength if (self.image_path or self.image_paths) else None,
            "controlnet_image_path": str(self.controlnet_image_path) if self.controlnet_image_path else None,
            "controlnet_strength": round(self.controlnet_strength, 2) if self.controlnet_strength else None,
            "masked_image_path": str(self.masked_image_path) if self.masked_image_path else None,
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@ -31,6 +31,7 @@ class ImageUtil:
        lora_scales: list[float],
        controlnet_image_path: str | Path | None = None,
        image_path: str | Path | None = None,
+        image_paths: list[str] | list[Path] | None = None,
        redux_image_paths: list[str] | list[Path] | None = None,
        redux_image_strengths: list[float] | None = None,
        image_strength: float | None = None,
@ -57,6 +58,7 @@ class ImageUtil:
            height=config.height,
            width=config.width,
            image_path=image_path,
+            image_paths=image_paths,
            image_strength=image_strength,
            controlnet_image_path=controlnet_image_path,
            controlnet_strength=config.controlnet_strength,
--- a/tests/image_generation/helpers/image_generation_edit_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_edit_test_helper.py
@ -26,6 +26,9 @@ class ImageGeneratorEditTestHelper:
        negative_prompt: str | None = None,
        quantize: int = 8,
        image_paths: list[str] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+        mismatch_threshold: float | None = None,
    ):
        # resolve paths
        reference_image_path = ImageGeneratorEditTestHelper.resolve_path(reference_image_path)
@ -46,6 +49,13 @@ class ImageGeneratorEditTestHelper:
            model_kwargs = {
                "quantize": quantize,
            }
+
+            # Add HuggingFace LoRA parameters if provided
+            if lora_names is not None:
+                model_kwargs["lora_names"] = lora_names
+            if lora_repo_id is not None:
+                model_kwargs["lora_repo_id"] = lora_repo_id
+
            model = model_class(**model_kwargs)

            # when
@ -90,6 +100,7 @@ class ImageGeneratorEditTestHelper:
                output_image_path,
                reference_image_path,
                f"Generated {model_name} image doesn't match reference image.",
+                mismatch_threshold=mismatch_threshold,
            )
        finally:
            # cleanup
--- a/tests/image_generation/helpers/image_generation_kontext_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_kontext_test_helper.py
@ -0,0 +1,64 @@
+import os
+from pathlib import Path
+
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
+from mflux.utils.image_compare import ImageCompare
+
+
+class ImageGeneratorKontextTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_config: ModelConfig,
+        steps: int,
+        seed: int,
+        height: int,
+        width: int,
+        prompt: str,
+        kontext_image_path: str,
+        guidance: float = 2.5,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorKontextTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorKontextTestHelper.resolve_path(output_image_path)
+        kontext_image_path = ImageGeneratorKontextTestHelper.resolve_path(kontext_image_path)
+
+        try:
+            # given
+            flux = Flux1Kontext(
+                quantize=8,
+            )
+
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    height=height,
+                    width=width,
+                    guidance=guidance,
+                    image_path=kontext_image_path,
+                ),
+            )
+            image.save(path=output_image_path, overwrite=True)
+
+            # then
+            ImageCompare.check_images_close_enough(
+                output_image_path,
+                reference_image_path,
+                "Generated kontext image doesn't match reference image.",
+            )
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path) and "MFLUX_PRESERVE_TEST_OUTPUT" not in os.environ:
+                os.remove(output_image_path)
+
+    @staticmethod
+    def resolve_path(path) -> Path | None:
+        if path is None:
+            return None
+        return Path(__file__).parent.parent.parent / "resources" / path
--- a/tests/image_generation/helpers/image_generation_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_test_helper.py
@ -30,6 +30,7 @@ class ImageGeneratorTestHelper:
        lora_repo_id: str | None = None,
        negative_prompt: str | None = None,
        guidance: float | None = None,
+        mismatch_threshold: float | None = None,
    ):
        # resolve paths
        reference_image_path = ImageGeneratorTestHelper.resolve_path(reference_image_path)
@ -83,6 +84,7 @@ class ImageGeneratorTestHelper:
                output_image_path,
                reference_image_path,
                "Generated image doesn't match reference image.",
+                mismatch_threshold=mismatch_threshold,
            )
        finally:
            # cleanup
--- a/tests/image_generation/test_generate_image.py
+++ b/tests/image_generation/test_generate_image.py
@ -1,6 +1,5 @@
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.txt2img.flux import Flux1
-from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper


@ -75,53 +74,3 @@ class TestImageGenerator:
            image_path="reference_dev_lora.png",
            prompt="Luxury food photograph of a burger",
        )
-
-    def test_qwen_image_generation_text_to_image(self):
-        ImageGeneratorTestHelper.assert_matches_reference_image(
-            reference_image_path="reference_qwen_txt2img.png",
-            output_image_path="output_qwen_txt2img.png",
-            model_class=QwenImage,
-            model_config=ModelConfig.qwen_image(),
-            quantize=6,
-            steps=20,
-            seed=42,
-            height=341,
-            width=768,
-            prompt="Luxury food photograph",
-            negative_prompt="ugly, blurry, low quality",
-        )
-
-    def test_qwen_image_generation_image_to_image(self):
-        ImageGeneratorTestHelper.assert_matches_reference_image(
-            reference_image_path="reference_qwen_img2img.png",
-            output_image_path="output_qwen_img2img.png",
-            model_class=QwenImage,
-            model_config=ModelConfig.qwen_image(),
-            quantize=6,
-            steps=20,
-            seed=44,
-            height=341,
-            width=768,
-            image_path="reference_dev_image_to_image.png",
-            image_strength=0.4,
-            prompt="Luxury food photograph of a burger",
-            negative_prompt="ugly, blurry, low quality",
-        )
-
-    def test_qwen_image_generation_lora(self):
-        ImageGeneratorTestHelper.assert_matches_reference_image(
-            reference_image_path="reference_qwen_lora.png",
-            output_image_path="output_qwen_lora.png",
-            model_class=QwenImage,
-            model_config=ModelConfig.qwen_image(),
-            guidance=1.0,
-            quantize=6,
-            steps=4,
-            seed=42,
-            height=341,
-            width=768,
-            prompt="Luxury food photograph",
-            negative_prompt="ugly, blurry, low quality",
-            lora_repo_id="lightx2v/Qwen-Image-Lightning",
-            lora_names=["Qwen-Image-Lightning-4steps-V2.0.safetensors"],
-        )
--- a/tests/image_generation/test_generate_image_kontext.py
+++ b/tests/image_generation/test_generate_image_kontext.py
@ -0,0 +1,18 @@
+from mflux.config.model_config import ModelConfig
+from tests.image_generation.helpers.image_generation_kontext_test_helper import ImageGeneratorKontextTestHelper
+
+
+class TestImageGeneratorKontext:
+    def test_image_generation_kontext(self):
+        ImageGeneratorKontextTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_dev_kontext.png",
+            output_image_path="output_dev_kontext.png",
+            model_config=ModelConfig.dev_kontext(),
+            steps=20,
+            seed=4869845,
+            height=384,
+            width=640,
+            guidance=2.5,
+            prompt="Make the hand fistbump the camera instead of showing a flat palm",
+            kontext_image_path="reference_upscaled.png",
+        )
--- a/tests/image_generation/test_generate_image_qwen_image.py
+++ b/tests/image_generation/test_generate_image_qwen_image.py
@ -0,0 +1,58 @@
+from mflux.config.model_config import ModelConfig
+from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
+
+
+class TestImageGeneratorQwenImage:
+    def test_qwen_image_generation_text_to_image(self):
+        ImageGeneratorTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_qwen_txt2img.png",
+            output_image_path="output_qwen_txt2img.png",
+            model_class=QwenImage,
+            model_config=ModelConfig.qwen_image(),
+            quantize=6,  # We should probably use at least 8-bit, but it doesn't run on 32GB machines
+            steps=20,
+            seed=42,
+            height=341,
+            width=768,
+            prompt="Luxury food photograph",
+            negative_prompt="ugly, blurry, low quality",
+            mismatch_threshold=0.35,  # Qwen models produce visually similar images with minor pixel differences
+        )
+
+    def test_qwen_image_generation_image_to_image(self):
+        ImageGeneratorTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_qwen_img2img.png",
+            output_image_path="output_qwen_img2img.png",
+            model_class=QwenImage,
+            model_config=ModelConfig.qwen_image(),
+            quantize=6,  # We should probably use at least 8-bit, but it doesn't run on 32GB machines
+            steps=20,
+            seed=44,
+            height=341,
+            width=768,
+            image_path="reference_dev_image_to_image.png",
+            image_strength=0.4,
+            prompt="Luxury food photograph of a burger",
+            negative_prompt="ugly, blurry, low quality",
+            mismatch_threshold=0.35,  # Qwen models produce visually similar images with minor pixel differences
+        )
+
+    def test_qwen_image_generation_lora(self):
+        ImageGeneratorTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_qwen_lora.png",
+            output_image_path="output_qwen_lora.png",
+            model_class=QwenImage,
+            model_config=ModelConfig.qwen_image(),
+            guidance=1.0,
+            quantize=6,  # We should probably use at least 8-bit, but it doesn't run on 32GB machines
+            steps=4,
+            seed=42,
+            height=341,
+            width=768,
+            prompt="Luxury food photograph",
+            negative_prompt="ugly, blurry, low quality",
+            lora_repo_id="lightx2v/Qwen-Image-Lightning",
+            lora_names=["Qwen-Image-Lightning-4steps-V2.0.safetensors"],
+            mismatch_threshold=0.65,  # LoRA tests have higher variance due to model updates
+        )
--- a/tests/image_generation/test_generate_image_qwen_image_edit.py
+++ b/tests/image_generation/test_generate_image_qwen_image_edit.py
@ -1,25 +1,11 @@
+import pytest
+
 from mflux.config.model_config import ModelConfig
-from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
 from mflux.models.qwen.variants.edit import QwenImageEdit
 from tests.image_generation.helpers.image_generation_edit_test_helper import ImageGeneratorEditTestHelper


-class TestImageGeneratorEdit:
-    def test_image_generation_flux_kontext(self):
-        ImageGeneratorEditTestHelper.assert_matches_reference_image(
-            reference_image_path="reference_dev_kontext.png",
-            output_image_path="output_dev_kontext.png",
-            model_class=Flux1Kontext,
-            model_config=ModelConfig.dev_kontext(),
-            steps=20,
-            seed=4869845,
-            height=384,
-            width=640,
-            guidance=2.5,
-            prompt="Make the hand fistbump the camera instead of showing a flat palm",
-            image_path="reference_upscaled.png",
-        )
-
+class TestImageGeneratorQwenImageEdit:
    def test_image_generation_qwen_edit(self):
        ImageGeneratorEditTestHelper.assert_matches_reference_image(
            reference_image_path="reference_qwen_edit.png",
@ -31,11 +17,13 @@ class TestImageGeneratorEdit:
            height=384,
            width=640,
            guidance=2.5,
-            quantize=6,
+            quantize=8,  # We should probably use at least 8-bit, but it doesn't run on 32GB machines
            prompt="Make the hand fistbump the camera instead of showing a flat palm",
            image_path="reference_upscaled.png",
+            mismatch_threshold=0.25,
        )

+    @pytest.mark.high_memory_requirement
    def test_image_generation_qwen_edit_multiple_images(self):
        ImageGeneratorEditTestHelper.assert_matches_reference_image(
            reference_image_path="reference_qwen_edit_multiple_images.png",
@ -47,7 +35,7 @@ class TestImageGeneratorEdit:
            height=384,
            width=640,
            guidance=2.5,
-            quantize=8,
+            quantize=8,  # We should probably use at least 8-bit, but it doesn't run on 32GB machines
            prompt="Make the hand fistbump the camera instead of showing a flat palm, and the man should wear this shirt. Maintain the original pose, body position, and overall stance.",
            image_paths=["reference_upscaled.png", "shirt.jpg"],
        )
--- a/uv.lock
+++ b/uv.lock