Qwen image edit (#274)

Co-authored-by: Filip Strand <filip@Host-022.local>

_typos.toml

@@ -2,4 +2,4 @@
 # use "typos --dump-config" to disocver other configs to override
 
 [default]
-extend-ignore-identifiers-re = ["EmbedND", "embed_nd", "alle"]
+extend-ignore-identifiers-re = ["EmbedND", "embed_nd", "alle", "thw"]

pyproject.toml

@@ -84,6 +84,8 @@ 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-edit = "mflux.generate_qwen_edit:main"
+mflux-generate-qwen-edit-plus = "mflux.generate_qwen_edit_plus:main"
 mflux-concept = "mflux.concept:main"
 mflux-concept-from-image = "mflux.concept_from_image:main"
 mflux-save = "mflux.save:main"

src/mflux/callbacks/instances/memory_saver.py

@@ -62,8 +62,12 @@ class MemorySaver(BeforeLoopCallback, InLoopCallback, AfterLoopCallback):
             self.model.clip_image_encoder = None
         if hasattr(self.model, "t5_image_encoder"):
             self.model.t5_image_encoder = None
-        if hasattr(self.model, "text_encoder"):
+        if hasattr(self.model, "text_encoder") and self.model.text_encoder is not None:
             self.model.text_encoder = None
+        if hasattr(self.model, "qwen_vl_encoder") and self.model.qwen_vl_encoder is not None:
+            self.model.qwen_vl_encoder = None
+        if hasattr(self.model, "qwen_vl_tokenizer") and self.model.qwen_vl_tokenizer is not None:
+            self.model.qwen_vl_tokenizer = None
         gc.collect()
         mx.clear_cache()
 

src/mflux/config/model_config.py

@@ -89,6 +89,16 @@ class ModelConfig:
     def qwen_image() -> "ModelConfig":
         return AVAILABLE_MODELS["qwen-image"]
 
+    @staticmethod
+    @lru_cache
+    def qwen_image_edit() -> "ModelConfig":
+        return AVAILABLE_MODELS["qwen-image-edit"]
+
+    @staticmethod
+    @lru_cache
+    def qwen_image_edit_plus() -> "ModelConfig":
+        return AVAILABLE_MODELS["qwen-image-edit-plus"]
+
     def x_embedder_input_dim(self) -> int:
         if "Fill" in self.model_name:
             return 384
@@ -130,7 +140,9 @@ class ModelConfig:
             # Infer from model_name substring - prefer longer matches (more specific)
             # Use case-insensitive matching for better compatibility
             model_name_lower = model_name.lower()
-            matching_bases = [(b, alias) for b in base_models for alias in b.aliases if alias and alias.lower() in model_name_lower]
+            matching_bases = [
+                (b, alias) for b in base_models for alias in b.aliases if alias and alias.lower() in model_name_lower
+            ]
 
             if matching_bases:
                 # Sort by alias length descending, then by priority ascending
@@ -300,4 +312,28 @@ AVAILABLE_MODELS = {
         requires_sigma_shift=None,
         priority=11,
     ),
+    "qwen-image-edit": ModelConfig(
+        aliases=["qwen-image-edit", "qwen-edit"],
+        model_name="Qwen/Qwen-Image-Edit",
+        base_model=None,
+        controlnet_model=None,
+        custom_transformer_model=None,
+        num_train_steps=None,
+        max_sequence_length=None,
+        supports_guidance=None,
+        requires_sigma_shift=None,
+        priority=12,
+    ),
+    "qwen-image-edit-plus": ModelConfig(
+        aliases=["qwen-image-edit-plus", "qwen-edit-plus", "qwen-edit-2509"],
+        model_name="Qwen/Qwen-Image-Edit-2509",
+        base_model=None,
+        controlnet_model=None,
+        custom_transformer_model=None,
+        num_train_steps=None,
+        max_sequence_length=None,
+        supports_guidance=None,
+        requires_sigma_shift=None,
+        priority=13,
+    ),
 }

src/mflux/error/exceptions.py

@@ -37,3 +37,7 @@ class CommandExecutionError(MFluxException):
         super().__init__(
             f"Command '{' '.join(cmd)}' exited with code {return_code}. See .stdout / .stderr for details."
         )
+
+
+class ReferenceVsOutputImageError(AssertionError):
+    """Raised when reference and output images don't match within the allowed threshold."""

src/mflux/generate_qwen_edit.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+
+from mflux.callbacks.callback_manager import CallbackManager
+from mflux.config.config import Config
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit
+from mflux.ui import defaults as ui_defaults
+from mflux.ui.cli.parsers import CommandLineParser
+from mflux.ui.prompt_utils import get_effective_prompt
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Generate an image using Flux Kontext with image conditioning.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=True)
+    parser.add_image_to_image_arguments(required=True)
+    parser.add_output_arguments()
+    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_KONTEXT
+
+    # 1. Load the model
+    qwen = QwenImageEdit(
+        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),
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    image_path=args.image_path,
+                ),
+            )
+
+            # 4. Save the image
+            output_path = Path(args.output.format(seed=seed))
+            image.save(path=output_path, 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()

src/mflux/generate_qwen_edit_plus.py

@@ -0,0 +1,81 @@
+from pathlib import Path
+
+from mflux.callbacks.callback_manager import CallbackManager
+from mflux.config.config import Config
+from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
+from mflux.models.qwen.variants.edit.qwen_image_edit_plus import QwenImageEditPlus
+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 Edit Plus with image conditioning.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=True)
+    parser.add_image_to_image_arguments(required=True)
+    parser.add_argument(
+        "--image-paths",
+        type=Path,
+        nargs="+",
+        default=None,
+        help="Local paths to multiple init images (for plus model). If not provided, uses --image-path as single image.",
+    )
+    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_KONTEXT
+
+    # 1. Load the model
+    qwen = QwenImageEditPlus(
+        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. Prepare image paths: use --image-paths if provided, otherwise fallback to --image-path
+            image_paths = None
+            if args.image_paths:
+                image_paths = [str(p) for p in args.image_paths]
+            elif args.image_path:
+                image_paths = [str(args.image_path)]
+
+            # 4. Generate an image for each seed value
+            image = qwen.generate_image(
+                seed=seed,
+                prompt=get_effective_prompt(args),
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    image_path=args.image_path,
+                ),
+                negative_prompt=get_effective_negative_prompt(args),
+                image_paths=image_paths,
+            )
+
+            # 5. Save the image
+            output_path = Path(args.output.format(seed=seed))
+            image.save(path=output_path, 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()

src/mflux/models/common/lora/mapping/lora_loader.py

@@ -11,7 +11,6 @@ from mflux.models.common.lora.mapping.lora_mapping import LoRATarget
 
 
 class LoRALoader:
-
     @staticmethod
     def load_and_apply_lora(
         lora_mapping: list[LoRATarget],
@@ -39,10 +38,7 @@ class LoRALoader:
 
     @staticmethod
     def _apply_single_lora(
-        transformer: nn.Module,
-        lora_file: str,
-        scale: float,
-        lora_mapping: list[LoRATarget]
+        transformer: nn.Module, lora_file: str, scale: float, lora_mapping: list[LoRATarget]
     ) -> None:
         # Load the LoRA weights
         if not Path(lora_file).exists():
@@ -58,8 +54,6 @@ class LoRALoader:
             print(f"❌ Failed to load LoRA file: {e}")
             return
 
-        print(f"🔍 DEBUG: Found {len(weights)} LoRA weights")
-
         # Apply LoRA using the provided mapping
         flat_mapping = LoRALoader._get_flat_mapping(lora_mapping)
         applied_count = LoRALoader._apply_lora_with_mapping(transformer, weights, scale, flat_mapping)
@@ -68,10 +62,7 @@ class LoRALoader:
 
     @staticmethod
     def _apply_lora_with_mapping(
-        transformer: nn.Module,
-        weights: dict,
-        scale: float,
-        lora_mappings: Dict[str, Tuple[str, str, bool]]
+        transformer: nn.Module, weights: dict, scale: float, lora_mappings: Dict[str, Tuple[str, str, bool]]
     ) -> int:
         applied_count = 0
         lora_data_by_target = {}
@@ -87,7 +78,7 @@ class LoRALoader:
                     # Extract block number from the weight key - try both . and _ separators
                     # This handles both standard LoRA formats (dot-separated) and other formats (underscore-separated)
                     # Find all numbers in the weight key
-                    numbers_in_key = re.findall(r'\d+', weight_key)
+                    numbers_in_key = re.findall(r"\d+", weight_key)
                     for num_str in numbers_in_key:
                         try:
                             test_block_idx = int(num_str)
@@ -119,8 +110,6 @@ class LoRALoader:
 
             lora_data_by_target[target_path][matrix_name] = (weight_value, transpose)
 
-        print(f"🔍 DEBUG: Found {len(lora_data_by_target)} LoRA target layers")
-
         # Apply LoRA to each target
         for target_path, lora_data in lora_data_by_target.items():
             if LoRALoader._apply_lora_matrices_to_target(transformer, target_path, lora_data, scale):
@@ -129,12 +118,7 @@ class LoRALoader:
         return applied_count
 
     @staticmethod
-    def _apply_lora_matrices_to_target(
-        transformer: nn.Module,
-        target_path: str,
-        lora_data: dict,
-        scale: float
-    ) -> bool:
+    def _apply_lora_matrices_to_target(transformer: nn.Module, target_path: str, lora_data: dict, scale: float) -> bool:
         # Navigate to the target layer
         current_module = transformer
         path_parts = target_path.split(".")
@@ -175,7 +159,7 @@ class LoRALoader:
 
         # Create new LoRA layer
         # Check if it's a linear layer (either nn.Linear, LoRALinear, or FusedLoRALinear)
-        is_linear = hasattr(current_module, 'weight')
+        is_linear = hasattr(current_module, "weight")
         is_lora_linear = isinstance(current_module, LoRALinear)
         is_fused_linear = isinstance(current_module, FusedLoRALinear)
 
@@ -184,11 +168,7 @@ class LoRALoader:
             if is_lora_linear:
                 print(f"   🔀 Fusing with existing LoRA at {target_path}")
                 # Create a temporary LoRA layer from the base linear of the existing LoRA
-                lora_layer = LoRALinear.from_linear(
-                    current_module.linear,
-                    r=lora_A.shape[1],
-                    scale=effective_scale
-                )
+                lora_layer = LoRALinear.from_linear(current_module.linear, r=lora_A.shape[1], scale=effective_scale)
                 # Set the LoRA matrices
                 lora_layer.lora_A = lora_A
                 lora_layer.lora_B = lora_B
@@ -197,18 +177,13 @@ class LoRALoader:
                     lora_layer.lora_B = lora_layer.lora_B * alpha_scale
 
                 # Create fused layer with the existing LoRA and the new one
-                fused_layer = FusedLoRALinear(
-                    base_linear=current_module.linear,
-                    loras=[current_module, lora_layer]
-                )
+                fused_layer = FusedLoRALinear(base_linear=current_module.linear, loras=[current_module, lora_layer])
                 replacement_layer = fused_layer
             elif is_fused_linear:
                 print(f"   🔀 Adding to existing fusion at {target_path}")
                 # Create a temporary LoRA layer from the base linear
                 lora_layer = LoRALinear.from_linear(
-                    current_module.base_linear,
-                    r=lora_A.shape[1],
-                    scale=effective_scale
+                    current_module.base_linear, r=lora_A.shape[1], scale=effective_scale
                 )
                 # Set the LoRA matrices
                 lora_layer.lora_A = lora_A
@@ -219,18 +194,13 @@ class LoRALoader:
 
                 # Add to existing fusion
                 fused_layer = FusedLoRALinear(
-                    base_linear=current_module.base_linear,
-                    loras=current_module.loras + [lora_layer]
+                    base_linear=current_module.base_linear, loras=current_module.loras + [lora_layer]
                 )
                 replacement_layer = fused_layer
             else:
                 # First LoRA on this layer
                 # Create LoRA layer
-                lora_layer = LoRALinear.from_linear(
-                    current_module,
-                    r=lora_A.shape[1],
-                    scale=effective_scale
-                )
+                lora_layer = LoRALinear.from_linear(current_module, r=lora_A.shape[1], scale=effective_scale)
                 # Set the LoRA matrices - use the correct dimensions from the LoRA file
                 lora_layer.lora_A = lora_A
                 lora_layer.lora_B = lora_B

src/mflux/models/common/weights/__init__.py

@@ -0,0 +1 @@
+"""Common weight mapping utilities."""

src/mflux/models/common/weights/mapping/__init__.py

@@ -0,0 +1,6 @@
+"""Weight mapping utilities for transforming HuggingFace weights to MLX structure."""
+
+from mflux.models.common.weights.mapping.weight_mapper import WeightMapper
+from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget
+
+__all__ = ["WeightMapping", "WeightTarget", "WeightMapper"]

src/mflux/models/common/weights/mapping/weight_mapper.py

@@ -0,0 +1,215 @@
+"""
+Weight Mapper - Applies declarative weight mappings to transform HF weights to MLX structure.
+
+Similar to LoRALoader, but for weight mapping instead of LoRA application.
+"""
+
+import re
+from typing import Callable, Dict, List, Optional
+
+import mlx.core as mx
+
+from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
+
+
+class WeightMapper:
+    """Maps HuggingFace weights to MLX nested structure using declarative mappings."""
+
+    @staticmethod
+    def apply_mapping(
+        hf_weights: Dict[str, mx.array],
+        mapping: List[WeightTarget],
+        num_blocks: Optional[int] = None,
+        num_layers: Optional[int] = None,
+    ) -> Dict:
+        """
+        Apply weight mapping to transform HF weights to MLX structure.
+
+        Args:
+            hf_weights: Raw HuggingFace weights (flat dict with dot-notation keys)
+            mapping: List of WeightTarget mappings
+            num_blocks: Number of transformer blocks (auto-detected if None)
+            num_layers: Number of text encoder layers (auto-detected if None)
+
+        Returns:
+            Nested dict structure matching MLX model
+        """
+        # Auto-detect number of blocks if not provided
+        if num_blocks is None:
+            num_blocks = WeightMapper._detect_num_blocks(hf_weights)
+
+        # Auto-detect number of layers if not provided
+        if num_layers is None:
+            num_layers = WeightMapper._detect_num_layers(hf_weights)
+
+        # Build flat mapping: HF pattern -> (MLX path, transform)
+        flat_mapping = WeightMapper._build_flat_mapping(mapping, num_blocks, num_layers)
+
+        # Map weights
+        mapped_weights = {}
+        mapped_count = 0
+        skipped_count = 0
+
+        for hf_key, hf_tensor in hf_weights.items():
+            # Try to find matching mapping
+            mlx_path, transform = WeightMapper._find_mapping(hf_key, flat_mapping)
+
+            if mlx_path:
+                # Apply transform if specified
+                tensor = hf_tensor
+                if transform:
+                    tensor = transform(tensor)
+
+                # Build nested structure
+                WeightMapper._set_nested_value(mapped_weights, mlx_path, tensor)
+                mapped_count += 1
+            else:
+                # Weight not in mapping - might be intentionally skipped (e.g., lm_head)
+                # or optional weight (e.g., conv_shortcut) - that's OK
+                skipped_count += 1
+
+        # Optional: uncomment for debugging
+        # print(f"✅ Mapped {mapped_count} weights, skipped {skipped_count}")
+
+        return mapped_weights
+
+    @staticmethod
+    def _detect_num_blocks(hf_weights: Dict[str, mx.array]) -> int:
+        """Detect number of transformer blocks from weight keys."""
+        block_numbers = set()
+        for key in hf_weights.keys():
+            # Match pattern: transformer_blocks.{number}.something
+            match = re.search(r"transformer_blocks\.(\d+)\.", key)
+            if match:
+                block_numbers.add(int(match.group(1)))
+
+        if block_numbers:
+            return max(block_numbers) + 1  # Blocks are 0-indexed
+        return 0
+
+    @staticmethod
+    def _detect_num_layers(hf_weights: Dict[str, mx.array]) -> int:
+        """Detect number of text encoder layers from weight keys."""
+        layer_numbers = set()
+        for key in hf_weights.keys():
+            # Match pattern: model.layers.{number}.something
+            match = re.search(r"model\.layers\.(\d+)\.", key)
+            if match:
+                layer_numbers.add(int(match.group(1)))
+
+        if layer_numbers:
+            return max(layer_numbers) + 1  # Layers are 0-indexed
+        return 28  # Default 28 layers for Qwen text encoder
+
+    @staticmethod
+    def _build_flat_mapping(
+        mapping: List[WeightTarget], num_blocks: int = 0, num_layers: int = 28
+    ) -> Dict[str, tuple[str, Optional[Callable[[mx.array], mx.array]]]]:
+        """
+        Build flat mapping from declarative targets.
+
+        Returns:
+            Dict mapping HF pattern -> (MLX path, transform)
+        """
+        flat = {}
+
+        for target in mapping:
+            # Expand placeholders for each pattern
+            for hf_pattern in target.hf_patterns:
+                # Check which placeholders are present
+                has_block = "{block}" in hf_pattern or "{block}" in target.mlx_path
+                has_i = "{i}" in hf_pattern or "{i}" in target.mlx_path
+                has_res = "{res}" in hf_pattern or "{res}" in target.mlx_path
+                has_layer = "{layer}" in hf_pattern or "{layer}" in target.mlx_path
+
+                # Handle multiple placeholders together
+                if has_block and has_res:
+                    # Up blocks: expand both {block} and {res}
+                    max_blocks = num_blocks if num_blocks > 0 else 4  # Default 4 for up_blocks
+                    for block_num in range(max_blocks):
+                        for res in range(3):  # 3 resnets per up_block
+                            concrete_hf = hf_pattern.replace("{block}", str(block_num)).replace("{res}", str(res))
+                            concrete_mlx = target.mlx_path.replace("{block}", str(block_num)).replace("{res}", str(res))
+                            flat[concrete_hf] = (concrete_mlx, target.transform)
+                elif has_block:
+                    # Expand {block} only (for transformer blocks or visual blocks)
+                    # Check if this is for visual blocks (32 blocks) or transformer blocks
+                    if "visual.blocks" in hf_pattern or "visual.blocks" in target.mlx_path:
+                        max_blocks = 32  # Visual blocks are always 32
+                    else:
+                        max_blocks = num_blocks if num_blocks > 0 else 4  # Default 4 for up_blocks
+                    for block_num in range(max_blocks):
+                        concrete_hf = hf_pattern.replace("{block}", str(block_num))
+                        concrete_mlx = target.mlx_path.replace("{block}", str(block_num))
+                        flat[concrete_hf] = (concrete_mlx, target.transform)
+                elif has_layer:
+                    # Expand {layer} for text encoder layers or visual blocks
+                    max_layers = num_layers if num_layers > 0 else 28  # Default 28 for text encoder
+                    for layer_num in range(max_layers):
+                        concrete_hf = hf_pattern.replace("{layer}", str(layer_num))
+                        concrete_mlx = target.mlx_path.replace("{layer}", str(layer_num))
+                        flat[concrete_hf] = (concrete_mlx, target.transform)
+                elif has_i:
+                    # Expand {i} only (for mid_block resnets)
+                    for i in range(2):  # 2 resnets in mid_block
+                        concrete_hf = hf_pattern.replace("{i}", str(i))
+                        concrete_mlx = target.mlx_path.replace("{i}", str(i))
+                        flat[concrete_hf] = (concrete_mlx, target.transform)
+                elif has_res:
+                    # This shouldn't happen for VAE (encoder down_blocks are explicit)
+                    # But handle it just in case
+                    if "up_block" in hf_pattern:
+                        for res in range(3):
+                            concrete_hf = hf_pattern.replace("{res}", str(res))
+                            concrete_mlx = target.mlx_path.replace("{res}", str(res))
+                            flat[concrete_hf] = (concrete_mlx, target.transform)
+                else:
+                    # No placeholder, use as-is
+                    flat[hf_pattern] = (target.mlx_path, target.transform)
+
+        return flat
+
+    @staticmethod
+    def _find_mapping(
+        hf_key: str, flat_mapping: Dict[str, tuple[str, Optional[Callable[[mx.array], mx.array]]]]
+    ) -> tuple[Optional[str], Optional[Callable[[mx.array], mx.array]]]:
+        """Find MLX path and transform for a given HF key."""
+        return flat_mapping.get(hf_key, (None, None))
+
+    @staticmethod
+    def _set_nested_value(d: Dict, path: str, value: mx.array):
+        """
+        Set value in nested dict using dot-notation path.
+
+        Creates nested structure as needed.
+        Handles both dict keys and list indices.
+        """
+        parts = path.split(".")
+        current = d
+        i = 0
+
+        while i < len(parts) - 1:
+            part = parts[i]
+
+            # Check if next part is a digit (list index)
+            if i + 1 < len(parts) and parts[i + 1].isdigit():
+                # This is a list, ensure it exists
+                if part not in current:
+                    current[part] = []
+                # Ensure list is large enough
+                idx = int(parts[i + 1])
+                while len(current[part]) <= idx:
+                    current[part].append({})
+                current = current[part][idx]
+                # Skip both the key and the index
+                i += 2
+            else:
+                # Regular dict key
+                if part not in current:
+                    current[part] = {}
+                current = current[part]
+                i += 1
+
+        # Set final value
+        final_key = parts[-1]
+        current[final_key] = value

src/mflux/models/common/weights/mapping/weight_mapping.py

@@ -0,0 +1,33 @@
+"""
+Base classes for declarative weight mapping.
+
+Similar to LoRA mapping, but for transforming HuggingFace weights to MLX structure.
+"""
+
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Protocol
+
+import mlx.core as mx
+
+
+@dataclass
+class WeightTarget:
+    """
+    Declarative weight mapping target.
+
+    Maps HuggingFace weight names to MLX nested structure paths.
+    """
+
+    mlx_path: str  # MLX nested path, e.g., "transformer_blocks.{block}.attn.to_q.weight"
+    hf_patterns: List[str]  # HuggingFace naming patterns, e.g., ["transformer_blocks.{block}.attn.to_q.weight"]
+    transform: Optional[Callable[[mx.array], mx.array]] = None  # Optional transform (reshape, transpose, etc.)
+    required: bool = True  # If False, weight is optional (may not exist in all models)
+
+
+class WeightMapping(Protocol):
+    """Protocol for weight mapping classes."""
+
+    @staticmethod
+    def get_mapping() -> List[WeightTarget]:
+        """Return list of weight mapping targets."""
+        return []

src/mflux/models/qwen/model/qwen_text_encoder/qwen_attention.py

@@ -42,18 +42,14 @@ class QwenAttention(nn.Module):
         position_embeddings: tuple[mx.array, mx.array] | None = None,
     ) -> mx.array:
         bsz, q_len, _ = hidden_states.shape
-
-        # Q, K, V projections
         query_states = self.q_proj(hidden_states)
         key_states = self.k_proj(hidden_states)
         value_states = self.v_proj(hidden_states)
 
-        # Reshape and transpose
         query_states = query_states.reshape(bsz, q_len, self.num_attention_heads, self.head_dim).transpose(0, 2, 1, 3)
         key_states = key_states.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(0, 2, 1, 3)
         value_states = value_states.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(0, 2, 1, 3)
 
-        # Apply RoPE
         query_states, key_states = QwenAttention._apply_multimodal_rotary_pos_emb(
             q=query_states,
             k=key_states,
@@ -61,15 +57,12 @@ class QwenAttention(nn.Module):
             mrope_section=self.rope_scaling["mrope_section"],
         )
 
-        # GQA expansion using repeat_interleave semantics
         if self.num_key_value_heads != self.num_attention_heads:
             key_states = QwenAttention._repeat_kv(key_states, self.num_key_value_groups)
             value_states = QwenAttention._repeat_kv(value_states, self.num_key_value_groups)
 
-        # Attention computation
         attn_weights = mx.matmul(query_states, key_states.transpose(0, 1, 3, 2)) * self.scaling
 
-        # Apply attention mask
         if attention_mask is not None:
             causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
             attn_weights = attn_weights + causal_mask
@@ -84,8 +77,6 @@ class QwenAttention(nn.Module):
     @staticmethod
     def _repeat_kv(hidden_states: mx.array, n_rep: int) -> mx.array:
         batch, num_key_value_heads, slen, head_dim = hidden_states.shape
-        if n_rep == 1:
-            return hidden_states
         hidden_states = mx.expand_dims(hidden_states, axis=2)
         hidden_states = mx.broadcast_to(hidden_states, (batch, num_key_value_heads, n_rep, slen, head_dim))
         return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
@@ -100,11 +91,10 @@ class QwenAttention(nn.Module):
     ) -> tuple[mx.array, mx.array]:
         mrope_section_doubled = [s * 2 for s in mrope_section]
 
+        cos, sin = position_embeddings
         cos_chunks = []
         sin_chunks = []
         start_idx = 0
-
-        cos, sin = position_embeddings
         for section_size in mrope_section_doubled:
             end_idx = start_idx + section_size
             cos_chunk = cos[..., start_idx:end_idx]
@@ -113,27 +103,29 @@ class QwenAttention(nn.Module):
             sin_chunks.append(sin_chunk)
             start_idx = end_idx
 
-        # For each chunk position, select the corresponding modality
-        # chunk 0 -> modality 0, chunk 1 -> modality 1, chunk 2 -> modality 2, etc.
         cos_selected = [chunk[i % 3] for i, chunk in enumerate(cos_chunks)]
         sin_selected = [chunk[i % 3] for i, chunk in enumerate(sin_chunks)]
 
-        # Concatenate back together
         cos_combined = mx.concatenate(cos_selected, axis=-1)
         sin_combined = mx.concatenate(sin_selected, axis=-1)
 
-        # Add head dimension if needed
         if unsqueeze_dim == 1:
             cos_combined = mx.expand_dims(cos_combined, axis=1)
             sin_combined = mx.expand_dims(sin_combined, axis=1)
-        elif unsqueeze_dim == 2:
-            cos_combined = mx.expand_dims(cos_combined, axis=2)
-            sin_combined = mx.expand_dims(sin_combined, axis=2)
 
-        # Apply rotary embedding with properly processed cos/sin
+        orig_q_dtype = q.dtype
+        orig_k_dtype = k.dtype
+        q = q.astype(mx.float32)
+        k = k.astype(mx.float32)
+        cos_combined = cos_combined.astype(mx.float32)
+        sin_combined = sin_combined.astype(mx.float32)
+
         q_embed = (q * cos_combined) + (QwenAttention._rotate_half(q) * sin_combined)
         k_embed = (k * cos_combined) + (QwenAttention._rotate_half(k) * sin_combined)
 
+        q_embed = q_embed.astype(orig_q_dtype)
+        k_embed = k_embed.astype(orig_k_dtype)
+
         return q_embed, k_embed
 
     @staticmethod

src/mflux/models/qwen/model/qwen_text_encoder/qwen_encoder.py

@@ -19,6 +19,8 @@ class QwenEncoder(nn.Module):
         self.vocab_size = vocab_size
         self.hidden_size = hidden_size
         self.num_hidden_layers = num_hidden_layers
+        self.image_token_id = 151655
+
         self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
         self.layers = [QwenEncoderLayer() for i in range(num_hidden_layers)]
         self.norm = QwenRMSNorm(hidden_size, eps=1e-6)
@@ -29,27 +31,69 @@ class QwenEncoder(nn.Module):
             rope_type="default",
         )
 
+        self.visual = None
+
+    def get_image_features(self, pixel_values: mx.array, image_grid_thw: mx.array) -> mx.array:
+        if self.visual is None:
+            raise RuntimeError("Vision transformer not initialized. Call load_visual_weights() first.")
+
+        pixel_values = pixel_values.astype(mx.float32)
+        image_embeds = self.visual(pixel_values, image_grid_thw)
+        original_split_sizes = image_grid_thw.prod(axis=-1).astype(mx.int32)
+        split_sizes = (original_split_sizes // 4).astype(mx.int32)
+        split_sizes = [int(s) for s in split_sizes.tolist()]
+        split_sizes = [s for s in split_sizes if s > 0]
+        image_embeds_split = []
+        start_idx = 0
+        for split_size in split_sizes:
+            end_idx = start_idx + split_size
+            image_embeds_split.append(image_embeds[start_idx:end_idx])
+            start_idx = end_idx
+        return image_embeds_split
+
     def __call__(
         self,
         input_ids: mx.array,
-        attention_mask: mx.array | None = None,
+        attention_mask: mx.array,
+        pixel_values: mx.array | None = None,
+        image_grid_thw: mx.array | None = None,
     ) -> mx.array:
         batch_size, seq_len = input_ids.shape
         inputs_embeds = self.embed_tokens(input_ids)
+
+        if pixel_values is not None and image_grid_thw is not None:
+            image_embeds_split = self.get_image_features(pixel_values, image_grid_thw)
+            image_embeds = mx.concatenate(image_embeds_split, axis=0)
+
+            image_positions = input_ids == self.image_token_id
+            n_image_tokens = mx.sum(image_positions).item()
+
+            if n_image_tokens > 0 and image_embeds.shape[0] >= n_image_tokens:
+                image_positions_flat = image_positions.flatten()
+                inputs_embeds_flat = inputs_embeds.reshape(-1, inputs_embeds.shape[-1])
+                image_embeds_to_use = image_embeds
+
+                new_embeds_list = []
+                image_idx = 0
+                for i in range(len(image_positions_flat)):
+                    if image_positions_flat[i] and image_idx < image_embeds_to_use.shape[0]:
+                        new_embeds_list.append(image_embeds_to_use[image_idx])
+                        image_idx += 1
+                    else:
+                        new_embeds_list.append(inputs_embeds_flat[i])
+
+                new_embeds = mx.stack(new_embeds_list, axis=0)
+                inputs_embeds = new_embeds.reshape(inputs_embeds.shape)
         cache_position = mx.arange(seq_len, dtype=mx.int32)
         position_ids = mx.expand_dims(mx.expand_dims(cache_position, axis=0), axis=0)
         position_ids = mx.broadcast_to(position_ids, (3, batch_size, seq_len))
-        if attention_mask is not None:
-            padding_mask = mx.where(
-                attention_mask == 1,
-                mx.zeros_like(attention_mask).astype(mx.float32),
-                mx.ones_like(attention_mask).astype(mx.float32) * (-float("inf")),
-            )
-            padding_mask = mx.expand_dims(mx.expand_dims(padding_mask, axis=1), axis=1)
-        else:
-            padding_mask = None
+        padding_mask = mx.where(
+            attention_mask == 1,
+            mx.zeros_like(attention_mask).astype(mx.float32),
+            mx.ones_like(attention_mask).astype(mx.float32) * (-float("inf")),
+        )
+        padding_mask = mx.expand_dims(mx.expand_dims(padding_mask, axis=1), axis=1)
 
-        # Create causal triangular mask
         idx = mx.arange(seq_len, dtype=mx.int32)
         j = mx.expand_dims(idx, axis=0)
         i = mx.expand_dims(idx, axis=1)
@@ -57,18 +101,13 @@ class QwenEncoder(nn.Module):
         zeros_2d = mx.zeros((seq_len, seq_len)).astype(mx.float32)
         neginf_2d = mx.ones((seq_len, seq_len)).astype(mx.float32) * (-float("inf"))
         causal_tri_mask = mx.where(tri_bool, neginf_2d, zeros_2d)
-        causal_tri_mask = mx.expand_dims(mx.expand_dims(causal_tri_mask, axis=0), axis=0)  # [1,1,S,S]
+        causal_tri_mask = mx.expand_dims(mx.expand_dims(causal_tri_mask, axis=0), axis=0)
         causal_tri_mask = mx.broadcast_to(causal_tri_mask, (batch_size, 1, seq_len, seq_len))
-
-        # Final attention mask: causal triangle + padding mask
-        if padding_mask is not None:
-            attention_mask_4d = causal_tri_mask + padding_mask
-        else:
-            attention_mask_4d = causal_tri_mask
-
+        attention_mask_4d = causal_tri_mask + padding_mask
         hidden_states = inputs_embeds
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
-        for layer in self.layers:
+        for i, layer in enumerate(self.layers):
             hidden_states = layer(hidden_states, attention_mask_4d, position_embeddings)
+
         hidden_states = self.norm(hidden_states)
         return hidden_states

src/mflux/models/qwen/model/qwen_text_encoder/qwen_patch_merger.py

@@ -0,0 +1,32 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class PatchMerger(nn.Module):
+    def __init__(self, context_dim: int, hidden_size: int, spatial_merge_size: int = 2):
+        super().__init__()
+        self.spatial_merge_size = spatial_merge_size
+        self.hidden_size_merged = context_dim * (spatial_merge_size**2)
+        self.ln_q = nn.RMSNorm(context_dim, eps=1e-6)
+        self.mlp_0 = nn.Linear(self.hidden_size_merged, self.hidden_size_merged, bias=True)
+        self.mlp_1 = nn.Linear(self.hidden_size_merged, hidden_size, bias=True)
+
+    def __call__(self, x: mx.array, grid_thw: mx.array) -> mx.array:
+        if not hasattr(self, "_weights_logged"):
+            self._weights_logged = True
+        x = self.ln_q(x)
+        merged_patches = []
+        offset = 0
+        for t, h, w in grid_thw:
+            t, h, w = int(t), int(h), int(w)
+            num_patches_this_image = t * h * w
+            x_this_image = x[offset : offset + num_patches_this_image]
+            x_merged = x_this_image.reshape(-1, self.hidden_size_merged)
+            merged_patches.append(x_merged)
+            offset += num_patches_this_image
+
+        x = mx.concatenate(merged_patches, axis=0)
+        x = self.mlp_0(x)
+        x = nn.gelu(x)
+        x = self.mlp_1(x)
+        return x

src/mflux/models/qwen/model/qwen_text_encoder/qwen_text_encoder.py

@@ -12,15 +12,17 @@ class QwenTextEncoder(nn.Module):
     def __call__(
         self,
         input_ids: mx.array,
-        attention_mask: mx.array | None = None,
+        attention_mask: mx.array,
     ) -> tuple[mx.array, mx.array]:
         hidden_states = self.encoder(input_ids, attention_mask)
+
         prompt_embeds, encoder_attention_mask = QwenTextEncoder._process_text_embeddings_mlx(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
             drop_idx=34,
             dtype=mx.bfloat16,
         )
+
         return prompt_embeds, encoder_attention_mask
 
     @staticmethod

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_attention.py

@@ -0,0 +1,83 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class VisionAttention(nn.Module):
+    def __init__(self, embed_dim: int = 1280, num_heads: int = 16):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=True)
+        self.proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+    def _rotate_half(self, x: mx.array) -> mx.array:
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return mx.concatenate([-x2, x1], axis=-1)
+
+    def _apply_rope(self, x: mx.array, cos: mx.array, sin: mx.array) -> mx.array:
+        cos_expanded = mx.expand_dims(cos, axis=0)
+        sin_expanded = mx.expand_dims(sin, axis=0)
+        rotated = (x * cos_expanded) + (self._rotate_half(x) * sin_expanded)
+        return rotated
+
+    def __call__(self, x: mx.array, position_embeddings=None, cu_seqlens=None) -> mx.array:
+        seq_len, embed_dim = x.shape
+
+        qkv = self.qkv(x)
+        qkv = qkv.reshape(seq_len, 3, self.num_heads, self.head_dim)
+        q, k, v = mx.split(qkv, 3, axis=1)
+        q = q.squeeze(1).transpose(1, 0, 2)
+        k = k.squeeze(1).transpose(1, 0, 2)
+        v = v.squeeze(1).transpose(1, 0, 2)
+
+        if position_embeddings is not None:
+            cos_emb, sin_emb = position_embeddings
+            if cos_emb.shape[0] != seq_len:
+                cos_emb = cos_emb[:seq_len]
+                sin_emb = sin_emb[:seq_len]
+
+            q = self._apply_rope(q, cos_emb, sin_emb)
+            k = self._apply_rope(k, cos_emb, sin_emb)
+
+        # Process attention chunks if cu_seqlens is provided (windowed attention)
+        if cu_seqlens is not None and len(cu_seqlens) > 2:
+            # Split Q, K, V into chunks based on cu_seqlens
+            # cu_seqlens is cumulative, so lengths[i] = cu_seqlens[i+1] - cu_seqlens[i]
+            lengths = [int((cu_seqlens[i + 1] - cu_seqlens[i]).item()) for i in range(len(cu_seqlens) - 1)]
+
+            # Split tensors (q,k,v are [heads, seq, head_dim])
+            q_chunks = []
+            k_chunks = []
+            v_chunks = []
+            offset = 0
+            for length in lengths:
+                q_chunks.append(q[:, offset : offset + length, :])
+                k_chunks.append(k[:, offset : offset + length, :])
+                v_chunks.append(v[:, offset : offset + length, :])
+                offset += length
+
+            # Process each chunk separately
+            attn_outputs = []
+            scale = 1.0 / (self.head_dim**0.5)
+            for q_chunk, k_chunk, v_chunk in zip(q_chunks, k_chunks, v_chunks):
+                # Compute attention for this chunk
+                scores = mx.matmul(q_chunk, k_chunk.transpose(0, 2, 1)) * scale
+                attn_weights = mx.softmax(scores, axis=-1)
+                attn_chunk = mx.matmul(attn_weights, v_chunk)  # [heads, chunk_len, head_dim]
+                attn_outputs.append(attn_chunk)
+
+            # Concatenate chunks back together
+            attn_output = mx.concatenate(attn_outputs, axis=1)  # [heads, seq, head_dim]
+        else:
+            # Full attention (no chunking)
+            scale = 1.0 / (self.head_dim**0.5)
+            scores = mx.matmul(q, k.transpose(0, 2, 1)) * scale  # [heads, seq, seq]
+            attn_weights = mx.softmax(scores, axis=-1)
+            attn_output = mx.matmul(attn_weights, v)  # [heads, seq, head_dim]
+
+        # Reshape and project
+        attn_output = attn_output.transpose(1, 0, 2).reshape(seq_len, embed_dim)  # [seq, embed_dim]
+        return self.proj(attn_output)

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_block.py

@@ -0,0 +1,24 @@
+import mlx.core as mx
+from mlx import nn
+
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_attention import VisionAttention
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_mlp import VisionMLP
+
+
+class VisionBlock(nn.Module):
+    def __init__(self, embed_dim: int = 1280, num_heads: int = 16, mlp_ratio: float = 2.671875):
+        super().__init__()
+        self.norm1 = nn.RMSNorm(embed_dim, eps=1e-6)
+        self.norm2 = nn.RMSNorm(embed_dim, eps=1e-6)
+        self.attn = VisionAttention(embed_dim, num_heads)
+        mlp_hidden_dim = int(embed_dim * mlp_ratio)
+        self.mlp = VisionMLP(embed_dim, mlp_hidden_dim)
+
+    def __call__(self, x: mx.array, position_embeddings=None, cu_seqlens=None) -> mx.array:
+        normed1 = self.norm1(x)
+        attn_out = self.attn(normed1, position_embeddings, cu_seqlens)
+        x = x + attn_out
+        normed2 = self.norm2(x)
+        mlp_out = self.mlp(normed2)
+        x = x + mlp_out
+        return x

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_encoder.py

@@ -0,0 +1,54 @@
+import mlx.core as mx
+from mlx import nn
+
+from mflux.models.qwen.model.qwen_text_encoder.qwen_encoder import QwenEncoder
+
+
+class QwenVisionLanguageEncoder(nn.Module):
+    def __init__(self, encoder=None):
+        super().__init__()
+        self.encoder = encoder or QwenEncoder()
+        self.edit_template_start_idx = 64
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        attention_mask: mx.array | None = None,
+        pixel_values: mx.array | None = None,
+        image_grid_thw: mx.array | None = None,
+    ) -> tuple[mx.array, mx.array]:
+        hidden_states = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+        )
+
+        trimmed = []
+        batch = hidden_states.shape[0]
+        for i in range(batch):
+            valid_len = int(mx.sum(attention_mask[i]).item())
+            trimmed.append(hidden_states[i, :valid_len, :])
+        drop_idx = self.edit_template_start_idx
+        trimmed_after_drop = [t[drop_idx:] if t.shape[0] > drop_idx else t for t in trimmed]
+        trimmed = trimmed_after_drop
+        max_len = max(t.shape[0] for t in trimmed) if trimmed else 0
+        hidden_dim = hidden_states.shape[2]
+        padded_embeds = []
+        padded_masks = []
+        for t in trimmed:
+            cur_len = t.shape[0]
+            if cur_len < max_len:
+                pad_e = mx.zeros((max_len - cur_len, hidden_dim), dtype=t.dtype)
+                t_pad = mx.concatenate([t, pad_e], axis=0)
+                pad_m = mx.concatenate(
+                    [mx.ones(cur_len, dtype=mx.int32), mx.zeros(max_len - cur_len, dtype=mx.int32)], axis=0
+                )
+            else:
+                t_pad = t
+                pad_m = mx.ones(cur_len, dtype=mx.int32)
+            padded_embeds.append(t_pad)
+            padded_masks.append(pad_m)
+        prompt_embeds = mx.stack(padded_embeds, axis=0) if padded_embeds else hidden_states
+        encoder_attention_mask = mx.stack(padded_masks, axis=0) if padded_masks else attention_mask
+        return prompt_embeds, encoder_attention_mask

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_language_prompt_encoder.py

@@ -0,0 +1,93 @@
+from typing import Optional, Union
+
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+
+from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_language_encoder import QwenVisionLanguageEncoder
+from mflux.models.qwen.tokenizer.qwen_tokenizer import TokenizerQwen
+from mflux.models.qwen.tokenizer.qwen_vision_language_tokenizer import QwenVisionLanguageTokenizer
+
+
+class QwenVisionLanguagePromptEncoder:
+    @staticmethod
+    def encode_prompt_with_image(
+        prompt: str,
+        image: Union[Image.Image, np.ndarray, str],
+        negative_prompt: Optional[str] = None,
+        prompt_cache: Optional[dict] = None,
+        qwen_vl_tokenizer: Optional[QwenVisionLanguageTokenizer] = None,
+        qwen_vl_encoder: Optional[QwenVisionLanguageEncoder] = None,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]:
+        if qwen_vl_tokenizer is None or qwen_vl_encoder is None:
+            raise ValueError("qwen_vl_tokenizer and qwen_vl_encoder are required")
+
+        input_ids, attention_mask, pixel_values, image_grid_thw = qwen_vl_tokenizer.tokenize_with_image(
+            prompt=prompt, image=image
+        )
+
+        prompt_embeds, prompt_mask = qwen_vl_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            image_grid_thw=image_grid_thw,
+        )
+
+        if negative_prompt is not None and negative_prompt != "":
+            neg_input_ids, neg_attention_mask, neg_pixel_values, neg_image_grid_thw = (
+                qwen_vl_tokenizer.tokenize_with_image(prompt=negative_prompt, image=image)
+            )
+
+            neg_prompt_embeds, neg_prompt_mask = qwen_vl_encoder(
+                input_ids=neg_input_ids,
+                attention_mask=neg_attention_mask,
+                pixel_values=neg_pixel_values,
+                image_grid_thw=neg_image_grid_thw,
+            )
+        else:
+            # If no negative prompt provided, mirror shapes with zeros masks/embeds to disable CFG cleanly
+            neg_prompt_embeds, neg_prompt_mask = prompt_embeds, prompt_mask
+
+        # 3. Return the result including image_grid_thw for conditioning
+        result = (prompt_embeds, prompt_mask, neg_prompt_embeds, neg_prompt_mask, image_grid_thw)
+        return result
+
+    @staticmethod
+    def encode_prompt_auto(
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        image: Optional[Union[Image.Image, np.ndarray, str]] = None,
+        prompt_cache: Optional[dict] = None,
+        qwen_tokenizer: Optional[TokenizerQwen] = None,
+        qwen_text_encoder: Optional[QwenTextEncoder] = None,
+        qwen_vl_tokenizer: Optional[QwenVisionLanguageTokenizer] = None,
+        qwen_vl_encoder: Optional[QwenVisionLanguageEncoder] = None,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array]:
+        if prompt_cache is None:
+            prompt_cache = {}
+
+        if image is not None and qwen_vl_tokenizer is not None and qwen_vl_encoder is not None:
+            # Use vision-language encoding for edit tasks
+            return QwenVisionLanguagePromptEncoder.encode_prompt_with_image(
+                prompt=prompt,
+                image=image,
+                negative_prompt=negative_prompt,
+                prompt_cache=prompt_cache,
+                qwen_vl_tokenizer=qwen_vl_tokenizer,
+                qwen_vl_encoder=qwen_vl_encoder,
+            )
+        else:
+            # Fall back to text-only encoding for regular generation
+            if qwen_tokenizer is None or qwen_text_encoder is None:
+                raise ValueError(
+                    "Text-only components (qwen_tokenizer, qwen_text_encoder) are required when no image is provided"
+                )
+
+            return QwenVisionLanguagePromptEncoder.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt or "",
+                prompt_cache=prompt_cache,
+                qwen_tokenizer=qwen_tokenizer,
+                qwen_text_encoder=qwen_text_encoder,
+            )

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_mlp.py

@@ -0,0 +1,15 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class VisionMLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=True)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=True)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gate = nn.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_patch_embed.py

@@ -0,0 +1,34 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class VisionPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 14,
+        temporal_patch_size: int = 2,
+        in_channels: int = 3,
+        embed_dim: int = 1280,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.in_channels = in_channels
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv3d(
+            in_channels=in_channels,
+            out_channels=embed_dim,
+            kernel_size=[temporal_patch_size, patch_size, patch_size],
+            stride=[temporal_patch_size, patch_size, patch_size],
+            bias=False,
+        )
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        batch_size = hidden_states.shape[0]
+        hidden_states = hidden_states.reshape(
+            batch_size, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+        )
+        hidden_states = hidden_states.transpose(0, 2, 3, 4, 1)
+        output = self.proj(hidden_states)
+        return output.reshape(batch_size, self.embed_dim)

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_rotary_embedding.py

@@ -0,0 +1,16 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class VisionRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        inv_freq = 1.0 / (theta ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
+        self.inv_freq = inv_freq
+
+    def __call__(self, max_grid_size: int) -> mx.array:
+        positions = mx.arange(max_grid_size, dtype=mx.float32)
+        freqs = mx.outer(positions, self.inv_freq)
+        return freqs

src/mflux/models/qwen/model/qwen_text_encoder/qwen_vision_transformer.py

@@ -0,0 +1,149 @@
+import mlx.core as mx
+import numpy as np
+from mlx import nn
+
+from mflux.models.qwen.model.qwen_text_encoder.qwen_patch_merger import PatchMerger
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_block import VisionBlock
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_patch_embed import VisionPatchEmbed
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_rotary_embedding import VisionRotaryEmbedding
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 14,
+        temporal_patch_size: int = 2,
+        in_channels: int = 3,
+        embed_dim: int = 1280,
+        depth: int = 32,
+        num_heads: int = 16,
+        mlp_ratio: float = 2.671875,
+        hidden_size: int = 3584,
+        spatial_merge_size: int = 2,
+        window_size: int = 112,
+        fullatt_block_indexes: list = None,
+    ):
+        super().__init__()
+
+        self.patch_embed = VisionPatchEmbed(patch_size, temporal_patch_size, in_channels, embed_dim)
+        self.spatial_merge_size = spatial_merge_size
+        self.window_size = window_size
+        self.fullatt_block_indexes = fullatt_block_indexes or [7, 15, 23, 31]
+        self.spatial_merge_unit = spatial_merge_size * spatial_merge_size
+        self.patch_size = patch_size
+
+        head_dim = embed_dim // num_heads
+        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
+
+        self.blocks = [VisionBlock(embed_dim, num_heads, mlp_ratio) for _ in range(depth)]
+        self.merger = PatchMerger(embed_dim, hidden_size, spatial_merge_size)
+
+    def get_window_index(self, grid_thw: mx.array):
+        window_index = []
+        cu_window_seqlens = [0]
+        window_index_id = 0
+        vit_merger_window_size = self.window_size // self.patch_size // self.spatial_merge_size
+
+        for t, grid_h, grid_w in grid_thw:
+            t, grid_h, grid_w = int(t), int(grid_h), int(grid_w)
+            llm_grid_h = grid_h // self.spatial_merge_size
+            llm_grid_w = grid_w // self.spatial_merge_size
+
+            index = mx.arange(t * llm_grid_h * llm_grid_w).reshape(t, llm_grid_h, llm_grid_w)
+
+            pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size
+            pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size
+            num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size
+            num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size
+
+            index_padded = mx.pad(index, ((0, 0), (0, pad_h), (0, pad_w)), constant_values=-100)
+
+            index_padded = index_padded.reshape(
+                t, num_windows_h, vit_merger_window_size, num_windows_w, vit_merger_window_size
+            )
+            index_padded = mx.transpose(index_padded, (0, 1, 3, 2, 4)).reshape(
+                t, num_windows_h * num_windows_w, vit_merger_window_size, vit_merger_window_size
+            )
+
+            seqlens = mx.sum((index_padded != -100).astype(mx.int32), axis=(2, 3)).reshape(-1)
+            index_padded_flat = index_padded.reshape(-1)
+
+            index_padded_np = np.array(index_padded_flat)
+            index_new = mx.array(index_padded_np[index_padded_np != -100])
+
+            window_index.append(index_new + window_index_id)
+            cu_seqlens_tmp = mx.cumsum(seqlens) * self.spatial_merge_unit + cu_window_seqlens[-1]
+            cu_window_seqlens.extend(cu_seqlens_tmp.tolist())
+            window_index_id += t * llm_grid_h * llm_grid_w
+
+        window_index = mx.concatenate(window_index, axis=0)
+        cu_window_seqlens = mx.array(cu_window_seqlens, dtype=mx.int32)
+
+        return window_index, cu_window_seqlens
+
+    def rot_pos_emb(self, grid_thw: mx.array) -> mx.array:
+        pos_ids = []
+        for t, h, w in grid_thw:
+            t, h, w = int(t), int(h), int(w)
+            hpos_ids = mx.repeat(mx.arange(h)[..., None], w, axis=1)
+            wpos_ids = mx.repeat(mx.arange(w)[None, ...], h, axis=0)
+            merge_h = h // self.spatial_merge_size
+            merge_w = w // self.spatial_merge_size
+            hpos_ids = hpos_ids.reshape(merge_h, self.spatial_merge_size, merge_w, self.spatial_merge_size)
+            wpos_ids = wpos_ids.reshape(merge_h, self.spatial_merge_size, merge_w, self.spatial_merge_size)
+            hpos_ids = mx.transpose(hpos_ids, (0, 2, 1, 3))
+            wpos_ids = mx.transpose(wpos_ids, (0, 2, 1, 3))
+            hpos_ids = hpos_ids.reshape(-1)
+            wpos_ids = wpos_ids.reshape(-1)
+            pos_id_pair = mx.stack([hpos_ids, wpos_ids], axis=-1)
+            pos_id_pair = mx.tile(pos_id_pair, (t, 1))
+            pos_ids.append(pos_id_pair)
+        pos_ids = mx.concatenate(pos_ids, axis=0)
+        max_grid_size = int(mx.max(grid_thw[:, 1:]).item())
+        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+        h_indices = pos_ids[:, 0].astype(mx.int32)
+        w_indices = pos_ids[:, 1].astype(mx.int32)
+        h_emb = rotary_pos_emb_full[h_indices]
+        w_emb = rotary_pos_emb_full[w_indices]
+        rotary_pos_emb = mx.stack([h_emb, w_emb], axis=1)
+        rotary_pos_emb = rotary_pos_emb.reshape(rotary_pos_emb.shape[0], -1)
+        return rotary_pos_emb
+
+    def __call__(self, pixel_values: mx.array, grid_thw: mx.array) -> mx.array:
+        hidden_states = self.patch_embed(pixel_values)
+        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+        window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+        cu_window_seqlens_unique = [cu_window_seqlens[0].item()]
+        for i in range(1, len(cu_window_seqlens)):
+            if cu_window_seqlens[i].item() != cu_window_seqlens_unique[-1]:
+                cu_window_seqlens_unique.append(cu_window_seqlens[i].item())
+        cu_window_seqlens = mx.array(cu_window_seqlens_unique, dtype=mx.int32)
+        seq_len = hidden_states.shape[0]
+        cu_seqlens = []
+        offset = 0
+        for t, h, w in grid_thw:
+            t, h, w = int(t), int(h), int(w)
+            length = t * h * w
+            offset += length
+            cu_seqlens.append(offset)
+        cu_seqlens = mx.array([0] + cu_seqlens, dtype=mx.int32)
+        seq_len = hidden_states.shape[0]
+        num_groups = seq_len // self.spatial_merge_unit
+        hidden_states_grouped = hidden_states.reshape(num_groups, self.spatial_merge_unit, -1)
+        hidden_states_grouped = hidden_states_grouped[window_index.astype(mx.int32), :, :]
+        hidden_states = hidden_states_grouped.reshape(seq_len, -1)
+        rotary_pos_emb_grouped = rotary_pos_emb.reshape(num_groups, self.spatial_merge_unit, -1)
+        rotary_pos_emb_grouped = rotary_pos_emb_grouped[window_index.astype(mx.int32), :, :]
+        rotary_pos_emb = rotary_pos_emb_grouped.reshape(seq_len, -1)
+        emb = mx.concatenate([rotary_pos_emb, rotary_pos_emb], axis=-1)
+        position_embeddings = (mx.cos(emb), mx.sin(emb))
+        for layer_num, block in enumerate(self.blocks):
+            if layer_num in self.fullatt_block_indexes:
+                cu_seqlens_now = cu_seqlens
+            else:
+                cu_seqlens_now = cu_window_seqlens
+            hidden_states = block(hidden_states, position_embeddings, cu_seqlens_now)
+        hidden_states = self.merger(hidden_states, grid_thw)
+        reverse_indices = mx.argsort(window_index.astype(mx.int32))
+        hidden_states = hidden_states[reverse_indices.astype(mx.int32), :]
+        return hidden_states

src/mflux/models/qwen/model/qwen_transformer/qwen_attention.py

@@ -2,8 +2,7 @@ from __future__ import annotations
 
 import mlx.core as mx
 from mlx import nn
-
-from mflux.models.flux.model.flux_transformer.common.attention_utils import AttentionUtils
+from mlx.core.fast import scaled_dot_product_attention
 
 
 class QwenAttention(nn.Module):
@@ -12,24 +11,16 @@ class QwenAttention(nn.Module):
         self.dim = dim
         self.num_heads = num_heads
         self.head_dim = head_dim
-
-        # Attention projections for image stream
         self.to_q = nn.Linear(dim, dim)
         self.to_k = nn.Linear(dim, dim)
         self.to_v = nn.Linear(dim, dim)
-
-        # Attention projections for text stream
         self.add_q_proj = nn.Linear(dim, dim)
         self.add_k_proj = nn.Linear(dim, dim)
         self.add_v_proj = nn.Linear(dim, dim)
-
-        # Query/Key normalization
         self.norm_q = nn.RMSNorm(self.head_dim, eps=1e-6)
         self.norm_k = nn.RMSNorm(self.head_dim, eps=1e-6)
         self.norm_added_q = nn.RMSNorm(self.head_dim, eps=1e-6)
         self.norm_added_k = nn.RMSNorm(self.head_dim, eps=1e-6)
-
-        # Output projections
         self.attn_to_out = [nn.Linear(dim, dim)]
         self.to_add_out = nn.Linear(dim, dim)
 
@@ -39,112 +30,129 @@ class QwenAttention(nn.Module):
         txt_modulated: mx.array,
         encoder_hidden_states_mask: mx.array | None,
         image_rotary_emb: tuple[mx.array, mx.array],
+        block_idx: int | None = None,
     ) -> tuple[mx.array, mx.array]:
-        # 1a. Compute Q,K,V for image stream (hidden_states)
-        img_q, img_k, img_v = AttentionUtils.process_qkv(
-            hidden_states=img_modulated,
-            to_q=self.to_q,
-            to_k=self.to_k,
-            to_v=self.to_v,
-            norm_q=self.norm_q,
-            norm_k=self.norm_k,
-            num_heads=self.num_heads,
-            head_dim=self.head_dim,
-        )
+        img_query = self.to_q(img_modulated)
+        img_key = self.to_k(img_modulated)
+        img_value = self.to_v(img_modulated)
 
-        # 1b. Compute Q,K,V for text stream (encoder_hidden_states)
-        txt_q, txt_k, txt_v = AttentionUtils.process_qkv(
-            hidden_states=txt_modulated,
-            to_q=self.add_q_proj,
-            to_k=self.add_k_proj,
-            to_v=self.add_v_proj,
-            norm_q=self.norm_added_q,
-            norm_k=self.norm_added_k,
-            num_heads=self.num_heads,
-            head_dim=self.head_dim,
-        )
+        txt_query = self.add_q_proj(txt_modulated)
+        txt_key = self.add_k_proj(txt_modulated)
+        txt_value = self.add_v_proj(txt_modulated)
 
-        # 1c. Concatenate results [text, image] like Flux
-        joint_q = mx.concatenate([txt_q, img_q], axis=2)
-        joint_k = mx.concatenate([txt_k, img_k], axis=2)
-        joint_v = mx.concatenate([txt_v, img_v], axis=2)
+        img_query = mx.reshape(img_query, (img_query.shape[0], img_query.shape[1], self.num_heads, self.head_dim))
+        img_key = mx.reshape(img_key, (img_key.shape[0], img_key.shape[1], self.num_heads, self.head_dim))
+        img_value = mx.reshape(img_value, (img_value.shape[0], img_value.shape[1], self.num_heads, self.head_dim))
 
-        # 1d. Apply RoPE to concatenated Q,K
-        joint_q, joint_k = QwenAttention._apply_rotary_embeddings_joint(
-            joint_q=joint_q,
-            joint_k=joint_k,
-            txt_seq_len=txt_q.shape[2],
-            image_rotary_emb=image_rotary_emb,
-        )
+        txt_query = mx.reshape(txt_query, (txt_query.shape[0], txt_query.shape[1], self.num_heads, self.head_dim))
+        txt_key = mx.reshape(txt_key, (txt_key.shape[0], txt_key.shape[1], self.num_heads, self.head_dim))
+        txt_value = mx.reshape(txt_value, (txt_value.shape[0], txt_value.shape[1], self.num_heads, self.head_dim))
 
-        # 2. Compute attention with optional masking
-        mask = AttentionUtils.convert_key_padding_mask_to_additive_mask(
+        if self.norm_q is not None:
+            img_query = self.norm_q(img_query)
+        if self.norm_k is not None:
+            img_key = self.norm_k(img_key)
+        if self.norm_added_q is not None:
+            txt_query = self.norm_added_q(txt_query)
+        if self.norm_added_k is not None:
+            txt_key = self.norm_added_k(txt_key)
+
+        if image_rotary_emb is not None:
+            (img_cos, img_sin), (txt_cos, txt_sin) = image_rotary_emb
+            img_query = QwenAttention._apply_rope_qwen(img_query, img_cos, img_sin)
+            img_key = QwenAttention._apply_rope_qwen(img_key, img_cos, img_sin)
+            txt_query = QwenAttention._apply_rope_qwen(txt_query, txt_cos, txt_sin)
+            txt_key = QwenAttention._apply_rope_qwen(txt_key, txt_cos, txt_sin)
+
+        joint_query = mx.concatenate([txt_query, img_query], axis=1)
+        joint_key = mx.concatenate([txt_key, img_key], axis=1)
+        joint_value = mx.concatenate([txt_value, img_value], axis=1)
+
+        seq_txt = txt_modulated.shape[1]
+        mask = self._convert_mask_for_qwen(
             mask=encoder_hidden_states_mask,
-            joint_seq_len=joint_q.shape[2],
-            txt_seq_len=txt_q.shape[2],
+            joint_seq_len=joint_query.shape[1],
+            txt_seq_len=seq_txt,
         )
 
-        # 3. Compute attention
-        hidden_states = AttentionUtils.compute_attention(
-            query=joint_q,
-            key=joint_k,
-            value=joint_v,
-            batch_size=1,
-            num_heads=self.num_heads,
-            head_dim=self.head_dim,
+        hidden_states = self._compute_attention_qwen(
+            query=joint_query,
+            key=joint_key,
+            value=joint_value,
             mask=mask,
+            block_idx=block_idx,
         )
 
-        # 4. Separate the results
-        txt_seq_len = txt_modulated.shape[1]
-        txt_attn_output = hidden_states[:, :txt_seq_len]
-        img_attn_output = hidden_states[:, txt_seq_len:]
-
-        # 5. Project the output (Flux-style)
-        img_attn_output = self.attn_to_out[0](img_attn_output.astype(mx.float32)).astype(img_modulated.dtype)
-        txt_attn_output = self.to_add_out(txt_attn_output.astype(mx.float32)).astype(txt_modulated.dtype)
+        txt_attn_output = hidden_states[:, :seq_txt, :]
+        img_attn_output = hidden_states[:, seq_txt:, :]
+        img_attn_output = self.attn_to_out[0](img_attn_output)
+        txt_attn_output = self.to_add_out(txt_attn_output)
 
         return img_attn_output, txt_attn_output
 
+    def _compute_attention_qwen(
+        self,
+        query: mx.array,
+        key: mx.array,
+        value: mx.array,
+        mask: mx.array | None = None,
+        block_idx: int | None = None,
+    ) -> mx.array:
+        query_bhsd = mx.transpose(query, (0, 2, 1, 3))
+        key_bhsd = mx.transpose(key, (0, 2, 1, 3))
+        value_bhsd = mx.transpose(value, (0, 2, 1, 3))
+        head_dim = query.shape[-1]
+        scale_value = 1.0 / (head_dim**0.5)
+        hidden_states_bhsd = scaled_dot_product_attention(
+            query_bhsd, key_bhsd, value_bhsd, scale=scale_value, mask=mask
+        )
+        hidden_states = mx.transpose(hidden_states_bhsd, (0, 2, 1, 3))
+        batch_size = hidden_states.shape[0]
+        seq_len = hidden_states.shape[1]
+        hidden_states = mx.reshape(hidden_states, (batch_size, seq_len, self.num_heads * self.head_dim))
+        hidden_states = hidden_states.astype(query.dtype)
+        return hidden_states
+
     @staticmethod
-    def _apply_rotary_embeddings_joint(
-        joint_q: mx.array,
-        joint_k: mx.array,
+    def _convert_mask_for_qwen(
+        mask: mx.array | None,
+        joint_seq_len: int,
         txt_seq_len: int,
-        image_rotary_emb: tuple[mx.array, mx.array],
-    ) -> tuple[mx.array, mx.array]:
-        img_rot, txt_rot = image_rotary_emb
+    ) -> mx.array | None:
+        if mask is None:
+            return None
 
-        # Extract separate parts
-        txt_q = joint_q[:, :, :txt_seq_len, :]
-        img_q = joint_q[:, :, txt_seq_len:, :]
-        txt_k = joint_k[:, :, :txt_seq_len, :]
-        img_k = joint_k[:, :, txt_seq_len:, :]
+        bsz = mask.shape[0]
+        img_seq_len = joint_seq_len - txt_seq_len
 
-        # Prepare cos/sin for text and image
-        img_cos = img_rot[..., 0, 0].reshape(img_rot.shape[2], img_rot.shape[3])
-        img_sin = img_rot[..., 1, 0].reshape(img_rot.shape[2], img_rot.shape[3])
-        txt_cos = txt_rot[..., 0, 0].reshape(txt_rot.shape[2], txt_rot.shape[3])
-        txt_sin = txt_rot[..., 1, 0].reshape(txt_rot.shape[2], txt_rot.shape[3])
+        ones_img = mx.ones((bsz, img_seq_len), dtype=mx.float32)
+        joint_mask = mx.concatenate([mask.astype(mx.float32), ones_img], axis=1)
 
-        # Transpose [B,H,S,D] -> [B,S,H,D] for RoPE application
-        img_q_bshd = mx.transpose(img_q, (0, 2, 1, 3))
-        img_k_bshd = mx.transpose(img_k, (0, 2, 1, 3))
-        txt_q_bshd = mx.transpose(txt_q, (0, 2, 1, 3))
-        txt_k_bshd = mx.transpose(txt_k, (0, 2, 1, 3))
+        if mx.all(joint_mask >= 0.999):
+            return None
 
-        # Apply RoPE
-        img_q_bshd, img_k_bshd = AttentionUtils.apply_rope_bshd(img_q_bshd, img_k_bshd, img_cos, img_sin)
-        txt_q_bshd, txt_k_bshd = AttentionUtils.apply_rope_bshd(txt_q_bshd, txt_k_bshd, txt_cos, txt_sin)
+        additive = (1.0 - joint_mask) * (-1e9)
+        return additive.reshape((additive.shape[0], 1, 1, additive.shape[1]))
 
-        # Back to [B,H,S,D]
-        img_q = mx.transpose(img_q_bshd, (0, 2, 1, 3))
-        img_k = mx.transpose(img_k_bshd, (0, 2, 1, 3))
-        txt_q = mx.transpose(txt_q_bshd, (0, 2, 1, 3))
-        txt_k = mx.transpose(txt_k_bshd, (0, 2, 1, 3))
+    @staticmethod
+    def _apply_rope_qwen(x: mx.array, cos_vals: mx.array, sin_vals: mx.array) -> mx.array:
+        x_float = x.astype(mx.float32)
+        x_reshaped = mx.reshape(x_float, (*x.shape[:-1], -1, 2))
 
-        # Concatenate back [text, image]
-        joint_q = mx.concatenate([txt_q, img_q], axis=2)
-        joint_k = mx.concatenate([txt_k, img_k], axis=2)
+        x_real = x_reshaped[..., 0]
+        x_imag = x_reshaped[..., 1]
 
-        return joint_q, joint_k
+        freqs_cos = cos_vals[None, :, None, :]
+        freqs_sin = sin_vals[None, :, None, :]
+
+        if freqs_cos.shape[-1] != x_real.shape[-1]:
+            freqs_cos = freqs_cos[..., : x_real.shape[-1]]
+            freqs_sin = freqs_sin[..., : x_real.shape[-1]]
+
+        out_real = x_real * freqs_cos - x_imag * freqs_sin
+        out_imag = x_real * freqs_sin + x_imag * freqs_cos
+
+        out_pairs = mx.stack([out_real, out_imag], axis=-1)
+        x_out = mx.reshape(out_pairs, (*x.shape[:-1], -1))
+
+        return x_out.astype(x.dtype)

src/mflux/models/qwen/model/qwen_transformer/qwen_feed_forward.py

@@ -5,11 +5,11 @@ from mlx import nn
 class QwenFeedForward(nn.Module):
     def __init__(self, dim: int = 3072):
         super().__init__()
-        self.mlp_in = nn.Linear(dim, 4 * dim)
-        self.mlp_out = nn.Linear(4 * dim, dim)
+        self.mlp_in = nn.Linear(dim, 4 * dim, bias=True)
+        self.mlp_out = nn.Linear(4 * dim, dim, bias=True)
 
     def __call__(self, hidden_states: mx.array) -> mx.array:
-        hidden_states = self.mlp_in(hidden_states.astype(mx.float32))
+        hidden_states = self.mlp_in(hidden_states)
         hidden_states = nn.gelu_approx(hidden_states)
         hidden_states = self.mlp_out(hidden_states)
         return hidden_states

src/mflux/models/qwen/model/qwen_transformer/qwen_layer_norm.py

@@ -9,13 +9,10 @@ class QwenLayerNorm(nn.Module):
         self.norm1 = nn.LayerNorm(dims=dim, eps=1e-6, affine=False)
 
     def __call__(self, hidden_states: mx.array, text_embeddings: mx.array) -> tuple[mx.array, mx.array, mx.array]:
-        # Compute modulation parameters from text embeddings
-        mod_params = self.mod_linear(nn.silu(text_embeddings))
+        temb_silu = nn.silu(text_embeddings)
+        mod_params = self.mod_linear(temb_silu)
         mod1, mod2 = mx.split(mod_params, 2, axis=-1)
-
-        # Stage 1: Apply normalization and modulation for attention
         normed1 = self.norm1(hidden_states)
         shift1, scale1, gate1 = mx.split(mod1, 3, axis=-1)
         normed_stage1 = normed1 * (1 + scale1[:, None, :]) + shift1[:, None, :]
-
         return normed_stage1, gate1, mod2

src/mflux/models/qwen/model/qwen_transformer/qwen_rope.py

@@ -6,114 +6,108 @@ from mlx import nn
 
 
 class QwenEmbedRopeMLX(nn.Module):
-    """
-    Faithful MLX port of Qwen's RoPE embedding helper used in the flux_transformer.
-
-    Exposes a forward(video_fhw, txt_seq_lens) API that returns rotation matrices
-    for image and text streams with shape [1, 1, S, D/2, 2, 2].
-    """
-
-    def __init__(self, theta: int, axes_dim: list[int], scale_rope: bool = True):
+    def __init__(self, theta: int, axes_dim: list[int], scale_rope: bool = False):
         super().__init__()
         self.theta = theta
         self.axes_dim = axes_dim
         self.scale_rope = scale_rope
 
-        # Precompute positive and negative index caches up to 1024 as in reference
-        pos_index = np.arange(1024, dtype=np.int32)
-        neg_index = (np.arange(1024, dtype=np.int32)[::-1] * -1) - 1
+        pos_index = np.arange(4096, dtype=np.int32)
+        neg_index = (np.arange(4096, dtype=np.int32)[::-1] * -1) - 1
 
-        self._pos_cos = []
-        self._pos_sin = []
-        self._neg_cos = []
-        self._neg_sin = []
-        for dim in axes_dim:
-            cos_p, sin_p = self._rope_params(pos_index, dim)
-            cos_n, sin_n = self._rope_params(neg_index, dim)
-            self._pos_cos.append(cos_p)
-            self._pos_sin.append(sin_p)
-            self._neg_cos.append(cos_n)
-            self._neg_sin.append(sin_n)
+        self.pos_freqs = np.concatenate(
+            [
+                self._rope_params(pos_index, self.axes_dim[0], self.theta),
+                self._rope_params(pos_index, self.axes_dim[1], self.theta),
+                self._rope_params(pos_index, self.axes_dim[2], self.theta),
+            ],
+            axis=1,
+        )
+        self.neg_freqs = np.concatenate(
+            [
+                self._rope_params(neg_index, self.axes_dim[0], self.theta),
+                self._rope_params(neg_index, self.axes_dim[1], self.theta),
+                self._rope_params(neg_index, self.axes_dim[2], self.theta),
+            ],
+            axis=1,
+        )
 
-    @staticmethod
-    def _rope_params(index: np.ndarray, dim: int, theta: int = 10000) -> tuple[np.ndarray, np.ndarray]:
+    def _rope_params(self, index: np.ndarray, dim: int, theta: int) -> np.ndarray:
         assert dim % 2 == 0
         scales = np.arange(0, dim, 2, dtype=np.float32) / dim
         omega = 1.0 / (theta**scales)
-        out = np.outer(index.astype(np.float32), omega).astype(np.float32)
-        return np.cos(out), np.sin(out)
+        freqs = np.outer(index.astype(np.float32), omega)
 
-    def _build_video_freqs(self, frame: int, height: int, width: int) -> tuple[np.ndarray, np.ndarray]:
-        dim_f, dim_h, dim_w = self.axes_dim
+        cos_freqs = np.cos(freqs)
+        sin_freqs = np.sin(freqs)
 
-        cos_f = self._pos_cos[0][:frame].reshape(frame, 1, 1, -1)
-        sin_f = self._pos_sin[0][:frame].reshape(frame, 1, 1, -1)
+        return np.stack([cos_freqs, sin_freqs], axis=-1)
+
+    def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> tuple[np.ndarray, np.ndarray]:
+        seq_lens = frame * height * width
+
+        axes_splits = [x // 2 for x in self.axes_dim]
+        freqs_pos = np.split(self.pos_freqs, np.cumsum(axes_splits)[:-1], axis=1)
+        freqs_neg = np.split(self.neg_freqs, np.cumsum(axes_splits)[:-1], axis=1)
+
+        freqs_frame_raw = freqs_pos[0][idx : idx + frame]
+        freqs_frame = freqs_frame_raw.reshape(frame, 1, 1, -1, 2)
+        freqs_frame = np.broadcast_to(freqs_frame, (frame, height, width, freqs_frame.shape[-2], 2))
 
         if self.scale_rope:
-            cos_h = np.concatenate(
-                [self._neg_cos[1][-(height - height // 2) :], self._pos_cos[1][: height // 2]], axis=0
-            )
-            sin_h = np.concatenate(
-                [self._neg_sin[1][-(height - height // 2) :], self._pos_sin[1][: height // 2]], axis=0
+            freqs_height = np.concatenate(
+                [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], axis=0
             )
         else:
-            cos_h = self._pos_cos[1][:height]
-            sin_h = self._pos_sin[1][:height]
-        cos_h = cos_h.reshape(1, height, 1, -1)
-        sin_h = sin_h.reshape(1, height, 1, -1)
+            freqs_height = freqs_pos[1][:height]
+        freqs_height = freqs_height.reshape(1, height, 1, -1, 2)
+        freqs_height = np.broadcast_to(freqs_height, (frame, height, width, freqs_height.shape[-2], 2))
 
         if self.scale_rope:
-            cos_w = np.concatenate([self._neg_cos[2][-(width - width // 2) :], self._pos_cos[2][: width // 2]], axis=0)
-            sin_w = np.concatenate([self._neg_sin[2][-(width - width // 2) :], self._pos_sin[2][: width // 2]], axis=0)
+            freqs_width = np.concatenate([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], axis=0)
         else:
-            cos_w = self._pos_cos[2][:width]
-            sin_w = self._pos_sin[2][:width]
-        cos_w = cos_w.reshape(1, 1, width, -1)
-        sin_w = sin_w.reshape(1, 1, width, -1)
+            freqs_width = freqs_pos[2][:width]
+        freqs_width = freqs_width.reshape(1, 1, width, -1, 2)
+        freqs_width = np.broadcast_to(freqs_width, (frame, height, width, freqs_width.shape[-2], 2))
 
-        cos = np.concatenate(
-            [
-                np.broadcast_to(cos_f, (frame, height, width, cos_f.shape[-1])),
-                np.broadcast_to(cos_h, (frame, height, width, cos_h.shape[-1])),
-                np.broadcast_to(cos_w, (frame, height, width, cos_w.shape[-1])),
-            ],
-            axis=-1,
+        freqs = np.concatenate([freqs_frame, freqs_height, freqs_width], axis=-2)
+        freqs = freqs.reshape(seq_lens, -1, 2)
+
+        cos_freqs = freqs[..., 0]
+        sin_freqs = freqs[..., 1]
+
+        return cos_freqs, sin_freqs
+
+    def __call__(
+        self,
+        video_fhw: tuple[int, int, int] | list[tuple[int, int, int]],
+        txt_seq_lens: list[int],
+    ) -> tuple[tuple[mx.array, mx.array], tuple[mx.array, mx.array]]:
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+
+        vid_cos_list = []
+        vid_sin_list = []
+        max_vid_index = 0
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            cos_v, sin_v = self._compute_video_freqs(frame, height, width, idx)
+            vid_cos_list.append(cos_v)
+            vid_sin_list.append(sin_v)
+
+            if self.scale_rope:
+                max_vid_index = max(height // 2, width // 2, max_vid_index)
+            else:
+                max_vid_index = max(height, width, max_vid_index)
+
+        vid_cos = np.concatenate(vid_cos_list, axis=0)
+        vid_sin = np.concatenate(vid_sin_list, axis=0)
+
+        max_len = max(txt_seq_lens)
+        txt_cos = self.pos_freqs[max_vid_index : max_vid_index + max_len, :, 0]
+        txt_sin = self.pos_freqs[max_vid_index : max_vid_index + max_len, :, 1]
+
+        return (
+            (mx.array(vid_cos.astype(np.float32)), mx.array(vid_sin.astype(np.float32))),
+            (mx.array(txt_cos.astype(np.float32)), mx.array(txt_sin.astype(np.float32))),
         )
-        sin = np.concatenate(
-            [
-                np.broadcast_to(sin_f, (frame, height, width, sin_f.shape[-1])),
-                np.broadcast_to(sin_h, (frame, height, width, sin_h.shape[-1])),
-                np.broadcast_to(sin_w, (frame, height, width, sin_w.shape[-1])),
-            ],
-            axis=-1,
-        )
-
-        # Flatten to [S, D/2]
-        cos = cos.reshape(-1, cos.shape[-1])
-        sin = sin.reshape(-1, sin.shape[-1])
-        return cos, sin
-
-    def __call__(self, video_fhw: tuple[int, int, int] | list[tuple[int, int, int]], txt_seq_lens: list[int]):
-        if isinstance(video_fhw, list):
-            video_fhw = video_fhw[0]
-        frame, height, width = video_fhw
-
-        cos_v, sin_v = self._build_video_freqs(frame, height, width)
-
-        # Build text freqs using max_vid_index rule
-        max_vid_index = max(height // 2, width // 2) if self.scale_rope else max(height, width)
-        txt_len = max(txt_seq_lens)
-        # Combine axes cos/sin for positions starting at max_vid_index
-        cos_full = np.concatenate(self._pos_cos, axis=1)  # [1024, 128]
-        sin_full = np.concatenate(self._pos_sin, axis=1)
-        cos_t = cos_full[max_vid_index : max_vid_index + txt_len]
-        sin_t = sin_full[max_vid_index : max_vid_index + txt_len]
-
-        def to_rot(cos: np.ndarray, sin: np.ndarray) -> mx.array:
-            row0 = np.stack([cos, -sin], axis=-1)
-            row1 = np.stack([sin, cos], axis=-1)
-            rot = np.stack([row0, row1], axis=-2)
-            rot = rot[None, None, ...].astype(np.float32)
-            return mx.array(rot)
-
-        return to_rot(cos_v, sin_v), to_rot(cos_t, sin_t)

src/mflux/models/qwen/model/qwen_transformer/qwen_time_text_embed.py

@@ -8,10 +8,10 @@ from mflux.models.qwen.model.qwen_transformer.qwen_timesteps import QwenTimestep
 class QwenTimeTextEmbed(nn.Module):
     def __init__(self, timestep_proj_dim: int = 256, inner_dim: int = 3072):
         super().__init__()
-        self.time_proj = QwenTimesteps(proj_dim=timestep_proj_dim)
+        self.time_proj = QwenTimesteps(proj_dim=timestep_proj_dim, scale=1000.0)
         self.timestep_embedder = QwenTimestepEmbedding(proj_dim=timestep_proj_dim, inner_dim=inner_dim)
 
     def __call__(self, timestep: mx.array, hidden_states: mx.array) -> mx.array:
-        time_proj = self.time_proj(timestep)
-        time_emb = self.timestep_embedder(time_proj.astype(hidden_states.dtype))
-        return time_emb
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.astype(hidden_states.dtype))
+        return timesteps_emb

src/mflux/models/qwen/model/qwen_transformer/qwen_timestep_embedding.py

@@ -5,10 +5,11 @@ from mlx import nn
 class QwenTimestepEmbedding(nn.Module):
     def __init__(self, proj_dim: int, inner_dim: int):
         super().__init__()
-        self.linear_1 = nn.Linear(proj_dim, inner_dim)
-        self.linear_2 = nn.Linear(inner_dim, inner_dim)
+        self.linear_1 = nn.Linear(proj_dim, inner_dim, bias=True)
+        self.linear_2 = nn.Linear(inner_dim, inner_dim, bias=True)
 
     def __call__(self, x: mx.array) -> mx.array:
-        x = nn.silu(self.linear_1(x))
+        x = self.linear_1(x)
+        x = nn.silu(x)
         x = self.linear_2(x)
         return x

src/mflux/models/qwen/model/qwen_transformer/qwen_timesteps.py

@@ -1,3 +1,5 @@
+import math
+
 import mlx.core as mx
 from mlx import nn
 
@@ -6,16 +8,27 @@ class QwenTimesteps(nn.Module):
     def __init__(self, proj_dim: int = 256, scale: float = 1000.0):
         super().__init__()
         self.proj_dim = proj_dim
+        self.flip_sin_to_cos = True
+        self.downscale_freq_shift = 0.0
         self.scale = scale
 
     def __call__(self, timesteps: mx.array) -> mx.array:
+        assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
         half_dim = self.proj_dim // 2
-        max_period = 10000.0
-        exponent = -mx.log(mx.array(max_period)) * mx.arange(0, half_dim, dtype=mx.float32)
-        exponent = exponent / (half_dim - 0.0)
-        freqs = mx.exp(exponent)
-        emb = timesteps.astype(mx.float32)[:, None] * freqs[None, :]
+        max_period = 10000
+
+        exponent = -math.log(max_period) * mx.arange(0, half_dim, dtype=mx.float32)
+        exponent = exponent / (half_dim - self.downscale_freq_shift)
+
+        timesteps_dtype = timesteps.dtype
+        emb = mx.exp(exponent).astype(timesteps_dtype)
+
+        emb = timesteps[:, None].astype(mx.float32) * emb[None, :]
         emb = self.scale * emb
         emb = mx.concatenate([mx.sin(emb), mx.cos(emb)], axis=-1)
-        emb = mx.concatenate([emb[:, half_dim:], emb[:, :half_dim]], axis=-1)
+
+        if self.flip_sin_to_cos:
+            emb = mx.concatenate([emb[:, half_dim:], emb[:, :half_dim]], axis=-1)
+
         return emb

src/mflux/models/qwen/model/qwen_transformer/qwen_transformer.py

@@ -9,6 +9,7 @@ from mflux.models.flux.model.flux_transformer.ada_layer_norm_continuous import A
 from mflux.models.qwen.model.qwen_transformer.qwen_rope import QwenEmbedRopeMLX
 from mflux.models.qwen.model.qwen_transformer.qwen_time_text_embed import QwenTimeTextEmbed
 from mflux.models.qwen.model.qwen_transformer.qwen_transformer_block import QwenTransformerBlock
+from mflux.models.qwen.model.qwen_transformer.qwen_transformer_rms_norm import QwenTransformerRMSNorm
 
 
 class QwenTransformer(nn.Module):
@@ -25,7 +26,7 @@ class QwenTransformer(nn.Module):
         super().__init__()
         self.inner_dim = num_attention_heads * attention_head_dim
         self.img_in = nn.Linear(in_channels, self.inner_dim)
-        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = QwenTransformerRMSNorm(joint_attention_dim, eps=1e-6)
         self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
         self.time_text_embed = QwenTimeTextEmbed(timestep_proj_dim=256, inner_dim=self.inner_dim)
         self.pos_embed = QwenEmbedRopeMLX(theta=10000, axes_dim=[16, 56, 56], scale_rope=True)
@@ -40,14 +41,22 @@ class QwenTransformer(nn.Module):
         hidden_states: mx.array,
         encoder_hidden_states: mx.array,
         encoder_hidden_states_mask: mx.array,
+        qwen_image_ids: mx.array | None = None,
+        cond_image_grid: tuple[int, int, int] | None = None,
     ) -> mx.array:
-        # 1. Create embeddings
         hidden_states = self.img_in(hidden_states)
-        encoder_hidden_states = self.txt_in(self.txt_norm(encoder_hidden_states))
-        text_embeddings = QwenTransformer._compute_text_embeddings(t, hidden_states, self.time_text_embed, config)
-        image_rotary_embeddings = QwenTransformer._compute_rotary_embeddings(encoder_hidden_states_mask, self.pos_embed, config)  # fmt: off
-
-        # 2. Run the transformer blocks
+        batch_size = hidden_states.shape[0]
+        timestep = QwenTransformer._compute_timestep(t, config)
+        timestep = mx.broadcast_to(timestep, (batch_size,)).astype(hidden_states.dtype)
+        encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+        text_embeddings = self.time_text_embed(timestep, hidden_states)
+        image_rotary_embeddings = QwenTransformer._compute_rotary_embeddings(
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
+            pos_embed=self.pos_embed,
+            config=config,
+            cond_image_grid=cond_image_grid,
+        )
         for idx, block in enumerate(self.transformer_blocks):
             encoder_hidden_states, hidden_states = QwenTransformer._apply_transformer_block(
                 idx=idx,
@@ -58,8 +67,6 @@ class QwenTransformer(nn.Module):
                 text_embeddings=text_embeddings,
                 image_rotary_embeddings=image_rotary_embeddings,
             )
-
-        # 3. Apply output normalization and projection
         hidden_states = self.norm_out(hidden_states, text_embeddings)
         hidden_states = self.proj_out(hidden_states)
         return hidden_states
@@ -80,30 +87,56 @@ class QwenTransformer(nn.Module):
             encoder_hidden_states_mask=encoder_hidden_states_mask,
             text_embeddings=text_embeddings,
             image_rotary_emb=image_rotary_embeddings,
+            block_idx=idx,
         )
 
     @staticmethod
-    def _compute_text_embeddings(
-        t: int,
-        hidden_states: mx.array,
-        time_text_embed: QwenTimeTextEmbed,
+    def _compute_timestep(
+        t: int | float,
         config: RuntimeConfig,
     ) -> mx.array:
-        time_step = config.scheduler.sigmas[t]
-        batch = hidden_states.shape[0]
-        timestep = mx.array(np.full((batch,), time_step, dtype=np.float32))
-        text_embeddings = time_text_embed(timestep, hidden_states)
-        return text_embeddings
+        """
+        Compute timestep tensor from step index or value.
+        """
+        if isinstance(t, int):
+            if t < len(config.scheduler.sigmas):
+                timestep_idx = t
+                time_step = config.scheduler.sigmas[timestep_idx]
+            else:
+                timestep_idx = None
+                for idx, ts in enumerate(config.scheduler.timesteps):
+                    if abs(int(ts.item()) - t) < 1:
+                        timestep_idx = idx
+                        break
+                if timestep_idx is None:
+                    time_step = t / 1000.0
+                else:
+                    time_step = config.scheduler.sigmas[timestep_idx]
+        else:
+            timestep_idx = None
+            time_step = t
+
+        timestep = mx.array(np.full((1,), time_step, dtype=np.float32))
+        return timestep
 
     @staticmethod
     def _compute_rotary_embeddings(
         encoder_hidden_states_mask: mx.array,
         pos_embed: QwenEmbedRopeMLX,
         config: RuntimeConfig,
+        cond_image_grid: tuple[int, int, int] | list[tuple[int, int, int]] | None = None,
     ) -> tuple[mx.array, mx.array]:
         latent_height = config.height // 16
         latent_width = config.width // 16
-        img_shapes = [(1, latent_height, latent_width)]
+
+        if cond_image_grid is None:
+            img_shapes = [(1, latent_height, latent_width)]
+        else:
+            if isinstance(cond_image_grid, list):
+                img_shapes = [(1, latent_height, latent_width)] + cond_image_grid
+            else:
+                img_shapes = [(1, latent_height, latent_width), cond_image_grid]
+
         txt_seq_lens = [int(mx.sum(encoder_hidden_states_mask[i]).item()) for i in range(encoder_hidden_states_mask.shape[0])]  # fmt: off
         img_rotary_emb, txt_rotary_emb = pos_embed(video_fhw=img_shapes, txt_seq_lens=txt_seq_lens)
         return img_rotary_emb, txt_rotary_emb

src/mflux/models/qwen/model/qwen_transformer/qwen_transformer_block.py

@@ -5,18 +5,23 @@ from mlx import nn
 
 from mflux.models.qwen.model.qwen_transformer.qwen_attention import QwenAttention
 from mflux.models.qwen.model.qwen_transformer.qwen_feed_forward import QwenFeedForward
-from mflux.models.qwen.model.qwen_transformer.qwen_layer_norm import QwenLayerNorm
 
 
 class QwenTransformerBlock(nn.Module):
     def __init__(self, dim: int = 3072, num_heads: int = 24, head_dim: int = 128):
         super().__init__()
-        self.img_norm1 = QwenLayerNorm(dim=dim)
-        self.txt_norm1 = QwenLayerNorm(dim=dim)
+
+        self.img_mod_silu = nn.SiLU()
+        self.img_mod_linear = nn.Linear(dim, 6 * dim, bias=True)
+        self.img_norm1 = nn.LayerNorm(dims=dim, eps=1e-6, affine=False)
         self.attn = QwenAttention(dim=dim, num_heads=num_heads, head_dim=head_dim)
         self.img_norm2 = nn.LayerNorm(dims=dim, eps=1e-6, affine=False)
-        self.txt_norm2 = nn.LayerNorm(dims=dim, eps=1e-6, affine=False)
         self.img_ff = QwenFeedForward(dim=dim)
+
+        self.txt_mod_silu = nn.SiLU()
+        self.txt_mod_linear = nn.Linear(dim, 6 * dim, bias=True)
+        self.txt_norm1 = nn.LayerNorm(dims=dim, eps=1e-6, affine=False)
+        self.txt_norm2 = nn.LayerNorm(dims=dim, eps=1e-6, affine=False)
         self.txt_ff = QwenFeedForward(dim=dim)
 
     def __call__(
@@ -26,51 +31,46 @@ class QwenTransformerBlock(nn.Module):
         encoder_hidden_states_mask: mx.array | None,
         text_embeddings: mx.array,
         image_rotary_emb: tuple[mx.array, mx.array],
+        block_idx: int | None = None,
     ) -> tuple[mx.array, mx.array]:
-        # 1. Compute stage 1 normalization and modulation
-        img_modulated, img_gate1, img_mod2 = self.img_norm1(hidden_states, text_embeddings)
-        txt_modulated, txt_gate1, txt_mod2 = self.txt_norm1(encoder_hidden_states, text_embeddings)
+        img_mod_params = self.img_mod_linear(self.img_mod_silu(text_embeddings))
+        txt_mod_params = self.txt_mod_linear(self.txt_mod_silu(text_embeddings))
+
+        img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1)
+        txt_mod1, txt_mod2 = mx.split(txt_mod_params, 2, axis=-1)
+
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = QwenTransformerBlock._modulate(img_normed, img_mod1)
+
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = QwenTransformerBlock._modulate(txt_normed, txt_mod1)
 
-        # 2. Compute attention
         img_attn_output, txt_attn_output = self.attn(
             img_modulated=img_modulated,
             txt_modulated=txt_modulated,
             encoder_hidden_states_mask=encoder_hidden_states_mask,
             image_rotary_emb=image_rotary_emb,
+            block_idx=block_idx,
         )
 
-        # 3. Apply attention residual and feed forward
-        hidden_states = QwenTransformerBlock._apply_residual_and_feed_forward(
-            hidden_states=hidden_states,
-            output=img_attn_output,
-            gate_attn=img_gate1,
-            mod_params=img_mod2,
-            norm2=self.img_norm2,
-            ff_module=self.img_ff,
-        )
-        encoder_hidden_states = QwenTransformerBlock._apply_residual_and_feed_forward(
-            hidden_states=encoder_hidden_states,
-            output=txt_attn_output,
-            gate_attn=txt_gate1,
-            mod_params=txt_mod2,
-            norm2=self.txt_norm2,
-            ff_module=self.txt_ff,
-        )
+        hidden_states = hidden_states + img_gate1 * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = QwenTransformerBlock._modulate(img_normed2, img_mod2)
+
+        img_mlp_output = self.img_ff(img_modulated2)
+
+        hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = QwenTransformerBlock._modulate(txt_normed2, txt_mod2)
+        txt_mlp_output = self.txt_ff(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
 
         return encoder_hidden_states, hidden_states
 
     @staticmethod
-    def _apply_residual_and_feed_forward(
-        hidden_states: mx.array,
-        output: mx.array,
-        gate_attn: mx.array,
-        mod_params: mx.array,
-        norm2: nn.LayerNorm,
-        ff_module: QwenFeedForward,
-    ) -> mx.array:
-        hidden_states = hidden_states + gate_attn[:, None, :] * output
-        shift2, scale2, gate_ff = mx.split(mod_params, 3, axis=-1)
-        normed = norm2(hidden_states)
-        modulated = normed * (1 + scale2[:, None, :]) + shift2[:, None, :]
-        ff_output = ff_module(modulated)
-        return hidden_states + gate_ff[:, None, :] * ff_output.astype(hidden_states.dtype)
+    def _modulate(x: mx.array, mod_params: mx.array) -> tuple[mx.array, mx.array]:
+        shift, scale, gate = mx.split(mod_params, 3, axis=-1)
+        return x * (1 + scale[:, None, :]) + shift[:, None, :], gate[:, None, :]

src/mflux/models/qwen/model/qwen_transformer/qwen_transformer_rms_norm.py

@@ -0,0 +1,23 @@
+import mlx.core as mx
+from mlx import nn
+
+
+class QwenTransformerRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = mx.ones((dim,))
+        self.eps = eps
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        input_dtype = hidden_states.dtype
+        variance = mx.power(hidden_states.astype(mx.float32), 2).mean(axis=-1, keepdims=True)
+        hidden_states = hidden_states * mx.rsqrt(variance + self.eps)
+
+        if self.weight is not None:
+            if self.weight.dtype in [mx.bfloat16, mx.float16]:
+                hidden_states = hidden_states.astype(self.weight.dtype)
+            hidden_states = hidden_states * self.weight
+            if hidden_states.dtype != input_dtype:
+                hidden_states = hidden_states.astype(input_dtype)
+
+        return hidden_states

src/mflux/models/qwen/model/qwen_vae/qwen_image_attention_block_3d.py

@@ -8,7 +8,7 @@ class QwenImageAttentionBlock3D(nn.Module):
     def __init__(self, dim: int):
         super().__init__()
         self.dim = dim
-        self.norm = QwenImageRMSNorm(dim, images=True)  # Attention blocks use images=True
+        self.norm = QwenImageRMSNorm(dim, images=True)
         self.to_qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, stride=1, padding=0)
         self.proj = nn.Conv2d(dim, dim, kernel_size=1, stride=1, padding=0)
 
@@ -21,7 +21,7 @@ class QwenImageAttentionBlock3D(nn.Module):
         x_5d = self.norm(x_5d)
         x = mx.squeeze(x_5d, axis=2)
         x = mx.transpose(x, (0, 2, 3, 1))
-        qkv = self.to_qkv(x)  # [bt, h, w, 3c]
+        qkv = self.to_qkv(x)
         qkv = mx.transpose(qkv, (0, 3, 1, 2))
         qkv = mx.reshape(qkv, (batch_size * time, 1, channels * 3, height * width))
         qkv = mx.transpose(qkv, (0, 1, 3, 2))

src/mflux/models/qwen/model/qwen_vae/qwen_image_causal_conv_3d.py

@@ -24,21 +24,14 @@ class QwenImageCausalConv3D(nn.Module):
         self.kernel_size = kernel_size
 
     def __call__(self, x: mx.array) -> mx.array:
-        if isinstance(self.padding, int):
-            pad_t = pad_h = pad_w = self.padding
-        elif isinstance(self.padding, (tuple, list)) and len(self.padding) == 3:
-            pad_t, pad_h, pad_w = self.padding
-        else:
-            pad_t = pad_h = pad_w = self.padding
-
-        # Apply causal padding in channels-first format (B, C, T, H, W)
+        pad_t = pad_h = pad_w = self.padding
         if pad_t > 0 or pad_h > 0 or pad_w > 0:
             pad_spec = [
-                (0, 0),  # Batch dimension (B)
-                (0, 0),  # Channel dimension (C)
-                (2 * pad_t, 0),  # Time dimension (T) - causal: pad past, not future
-                (pad_h, pad_h),  # Height dimension (H)
-                (pad_w, pad_w),  # Width dimension (W)
+                (0, 0),
+                (0, 0),
+                (2 * pad_t, 0),
+                (pad_h, pad_h),
+                (pad_w, pad_w),
             ]
             x = mx.pad(x, pad_spec)
 

src/mflux/models/qwen/model/qwen_vae/qwen_image_decoder_3d.py

@@ -16,8 +16,8 @@ class QwenImageDecoder3D(nn.Module):
         self.up_block1 = QwenImageUpBlock3D(192, 384, num_res_blocks=2, upsample_mode="upsample3d")
         self.up_block2 = QwenImageUpBlock3D(192, 192, num_res_blocks=2, upsample_mode="upsample2d")
         self.up_block3 = QwenImageUpBlock3D(96, 96, num_res_blocks=2, upsample_mode=None)
-        self.norm_out = QwenImageRMSNorm(96, images=False)  # Decoder norm_out uses images=False
-        self.conv_out = QwenImageCausalConv3D(96, 3, 3, 1, 1)  # RGB output
+        self.norm_out = QwenImageRMSNorm(96, images=False)
+        self.conv_out = QwenImageCausalConv3D(96, 3, 3, 1, 1)
 
     def __call__(self, x: mx.array) -> mx.array:
         x = self.conv_in(x)

src/mflux/models/qwen/model/qwen_vae/qwen_image_down_block_3d.py

@@ -10,7 +10,7 @@ class QwenImageDownBlock3D(nn.Module):
         super().__init__()
         self.resnets = []
         current_dim = in_channels
-        for _ in range(num_res_blocks):  # FIXED: Use num_res_blocks, not num_res_blocks + 1
+        for _ in range(num_res_blocks):
             self.resnets.append(QwenImageResBlock3D(current_dim, out_channels))
             current_dim = out_channels
 

src/mflux/models/qwen/model/qwen_vae/qwen_image_encoder_3d.py

@@ -23,10 +23,7 @@ class QwenImageEncoder3D(nn.Module):
 
         down_blocks = []
         for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
-            if i < len(self.temporal_downsample):
-                downsample_mode = "downsample3d" if self.temporal_downsample[i] else "downsample2d"
-            else:
-                downsample_mode = "downsample2d"
+            downsample_mode = "downsample3d" if self.temporal_downsample[i] else "downsample2d"
             # Do not downsample on the final stage only
             if i == len(dims) - 2:
                 downsample_mode = None
@@ -39,14 +36,14 @@ class QwenImageEncoder3D(nn.Module):
         self.down_blocks = down_blocks
 
         self.mid_block = QwenImageMidBlock3D(dims[-1], num_layers=1)
-        self.norm_out = QwenImageRMSNorm(dims[-1], images=False)  # Encoder norm_out uses images=False
-        self.conv_out = QwenImageCausalConv3D(dims[-1], 32, 3, 1, 1)  # Changed z_dim to 32
+        self.norm_out = QwenImageRMSNorm(dims[-1], images=False)
+        self.conv_out = QwenImageCausalConv3D(dims[-1], 32, 3, 1, 1)
 
     def __call__(self, x: mx.array) -> mx.array:
         x = self.conv_in(x)
         for stage_idx, down_block in enumerate(self.down_blocks):
-            for res_idx, resnet in enumerate(down_block.resnets):
-                if stage_idx == 3:
+            if stage_idx == 3:
+                for resnet in down_block.resnets:
                     residual = x
                     n1 = resnet.norm1(x)
                     a1 = nn.silu(n1)
@@ -54,14 +51,11 @@ class QwenImageEncoder3D(nn.Module):
                     n2 = resnet.norm2(c1)
                     a2 = nn.silu(n2)
                     c2 = resnet.conv2(a2)
-                    if resnet.skip_conv is not None:
-                        residual = resnet.skip_conv(residual)
-                    y = c2 + residual
-                    x = y
-                else:
-                    x = resnet(x)
-            if down_block.downsamplers is not None:
-                x = down_block.downsamplers[0](x)
+                    x = c2 + residual
+                if down_block.downsamplers is not None:
+                    x = down_block.downsamplers[0](x)
+            else:
+                x = down_block(x)
 
         x = self.mid_block(x)
         norm_in = x

src/mflux/models/qwen/qwen_edit_initializer.py

@@ -0,0 +1,109 @@
+from pathlib import Path
+
+from transformers import Qwen2TokenizerFast
+
+from mflux.config.model_config import ModelConfig
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_language_encoder import QwenVisionLanguageEncoder
+from mflux.models.qwen.qwen_initializer import QwenImageInitializer
+from mflux.models.qwen.tokenizer.qwen_vision_language_processor import QwenVisionLanguageProcessor
+from mflux.models.qwen.tokenizer.qwen_vision_language_tokenizer import QwenVisionLanguageTokenizer
+from mflux.utils.download import snapshot_download
+
+
+class QwenImageEditInitializer:
+    @staticmethod
+    def init(
+        qwen_model,
+        model_config: ModelConfig,
+        quantize: int | None,
+        local_path: str | None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+    ) -> None:
+        # 1. Initialize the base Qwen Image model (VAE, transformer, text encoder, etc.)
+        QwenImageInitializer.init(
+            qwen_model=qwen_model,
+            model_config=model_config,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
+        )
+
+        # 2. Add vision-language components for edit functionality
+        QwenImageEditInitializer._init_vision_language_components(
+            qwen_model=qwen_model,
+            repo_id=model_config.model_name,
+            local_path=local_path,
+            model_config=model_config,
+        )
+
+    @staticmethod
+    def _init_vision_language_components(
+        qwen_model,
+        repo_id: str,
+        local_path: str | None,
+        model_config: ModelConfig | None = None,
+    ) -> None:
+        # 1. Download or get cached tokenizer
+        root_path = Path(local_path) if local_path else QwenImageEditInitializer._download_vl_processor(repo_id)
+
+        # 2. Load only the tokenizer (we implement image processing ourselves)
+        tokenizer_path = root_path / "tokenizer"
+        if not tokenizer_path.exists():
+            tokenizer_path = root_path
+
+        try:
+            tokenizer = Qwen2TokenizerFast.from_pretrained(
+                pretrained_model_name_or_path=tokenizer_path,
+                local_files_only=True,
+            )
+        except OSError:
+            tokenizer = Qwen2TokenizerFast.from_pretrained(repo_id)
+
+        # 3. Create our MLX processor with the tokenizer
+        processor = QwenVisionLanguageProcessor(tokenizer=tokenizer)
+
+        # 4. Initialize vision-language tokenizer wrapper
+        if model_config is not None:
+            model_name_lower = model_config.model_name.lower()
+            use_picture_prefix = (
+                "plus" in model_name_lower
+                or "2509" in model_config.model_name  # Check original case for "2509"
+                or (hasattr(model_config, "use_picture_prefix") and model_config.use_picture_prefix)
+            )
+        else:
+            # Fallback: check repo_id
+            use_picture_prefix = "plus" in repo_id.lower() or "2509" in repo_id
+
+        qwen_model.qwen_vl_tokenizer = QwenVisionLanguageTokenizer(
+            processor=processor,
+            max_length=1024,
+            use_picture_prefix=use_picture_prefix,
+        )
+
+        # 5. Initialize vision-language encoder (integrated approach like Diffusers)
+        qwen_model.qwen_vl_encoder = QwenVisionLanguageEncoder(encoder=qwen_model.text_encoder.encoder)
+
+    @staticmethod
+    def _download_vl_processor(repo_id: str) -> Path:
+        return Path(
+            snapshot_download(
+                repo_id=repo_id,
+                allow_patterns=[
+                    "processor/**",
+                    "preprocessor_config.json",
+                    "tokenizer/**",
+                    "added_tokens.json",
+                    "chat_template.jinja",
+                    "tokenizer.json",
+                    "tokenizer_config.json",
+                    "vocab.json",
+                    "merges.txt",
+                ],
+            )
+        )

src/mflux/models/qwen/qwen_initializer.py

@@ -5,8 +5,8 @@ from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenText
 from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
 from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE
 from mflux.models.qwen.tokenizer.qwen_tokenizer_handler import QwenTokenizerHandler
-from mflux.models.qwen.weights.qwen_weight_handler import QwenWeightHandler
 from mflux.models.qwen.weights.qwen_lora_mapping import QwenLoRAMapping
+from mflux.models.qwen.weights.qwen_weight_handler import QwenWeightHandler
 from mflux.models.qwen.weights.qwen_weight_util import QwenWeightUtil
 
 
@@ -68,7 +68,3 @@ class QwenImageInitializer:
                 lora_files=qwen_model.lora_paths,
                 lora_scales=qwen_model.lora_scales,
             )
-        else:
-            print("⚠️  No LoRA paths provided, skipping LoRA setup")
-
-        print("🔧 === END QWEN LORA WEIGHT SETUP DEBUG ===\n")

src/mflux/models/qwen/tokenizer/qwen_image_processor.py

@@ -0,0 +1,162 @@
+import math
+from typing import Optional, Union
+
+import numpy as np
+from PIL import Image
+
+OPENAI_CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
+OPENAI_CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int = 28,
+    min_pixels: int = 56 * 56,
+    max_pixels: int = 14 * 14 * 4 * 1280,
+) -> tuple[int, int]:
+    if max(height, width) / min(height, width) > 200:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+        )
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = max(factor, math.floor(height / beta / factor) * factor)
+        w_bar = max(factor, math.floor(width / beta / factor) * factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+    return h_bar, w_bar
+
+
+class QwenImageProcessor:
+    def __init__(
+        self,
+        min_pixels: int = 56 * 56,
+        max_pixels: int = 28 * 28 * 1280,
+        patch_size: int = 14,
+        temporal_patch_size: int = 2,
+        merge_size: int = 2,
+        image_mean: Optional[list[float]] = None,
+        image_std: Optional[list[float]] = None,
+    ):
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.merge_size = merge_size
+        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+
+    def _preprocess(
+        self,
+        image: Image.Image,
+        resized_height: Optional[int] = None,
+        resized_width: Optional[int] = None,
+    ) -> tuple[np.ndarray, tuple[int, int, int]]:
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        height, width = image.size[1], image.size[0]
+
+        if resized_height is None or resized_width is None:
+            resized_height, resized_width = smart_resize(
+                height,
+                width,
+                factor=self.patch_size * self.merge_size,
+                min_pixels=self.min_pixels,
+                max_pixels=self.max_pixels,
+            )
+        if (height, width) != (resized_height, resized_width):
+            image = image.resize((resized_width, resized_height), Image.BICUBIC)
+
+        image_np = np.array(image).astype(np.float32)
+        image_np = image_np / 255.0
+
+        mean_np = np.array(self.image_mean, dtype=np.float32)
+        std_np = np.array(self.image_std, dtype=np.float32)
+        image_np = (image_np - mean_np) / std_np
+
+        image_np = image_np.transpose(2, 0, 1)
+        patches = image_np[np.newaxis]  # Shape: (1, channel, height, width)
+
+        if patches.shape[0] % self.temporal_patch_size != 0:
+            repeats = np.repeat(
+                patches[-1][np.newaxis],
+                self.temporal_patch_size - (patches.shape[0] % self.temporal_patch_size),
+                axis=0,
+            )
+            patches = np.concatenate([patches, repeats], axis=0)
+
+        channel = patches.shape[1]
+        grid_t = patches.shape[0] // self.temporal_patch_size
+        grid_h = resized_height // self.patch_size
+        grid_w = resized_width // self.patch_size
+
+        patches = patches.reshape(
+            grid_t,
+            self.temporal_patch_size,
+            channel,
+            grid_h // self.merge_size,
+            self.merge_size,
+            self.patch_size,
+            grid_w // self.merge_size,
+            self.merge_size,
+            self.patch_size,
+        )
+
+        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+
+        flatten_patches = patches.reshape(
+            grid_t * grid_h * grid_w,
+            channel * self.temporal_patch_size * self.patch_size * self.patch_size,
+        )
+
+        return flatten_patches, (grid_t, grid_h, grid_w)
+
+    def preprocess(
+        self,
+        images: Union[Image.Image, list[Image.Image]],
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if not isinstance(images, list):
+            images = [images]
+
+        pixel_values_list = []
+        vision_grid_thws = []
+
+        for image in images:
+            patches, image_grid_thw = self._preprocess(image)
+            pixel_values_list.append(patches)
+            vision_grid_thws.append([image_grid_thw[0], image_grid_thw[1], image_grid_thw[2]])
+
+        # Concatenate all patches from all images along the patch dimension
+        pixel_values = np.concatenate(pixel_values_list, axis=0) if pixel_values_list else np.array([])
+
+        vision_grid_thws = np.array(vision_grid_thws)
+
+        return pixel_values, vision_grid_thws
+
+    def get_number_of_image_patches(
+        self,
+        height: int,
+        width: int,
+        min_pixels: Optional[int] = None,
+        max_pixels: Optional[int] = None,
+    ) -> int:
+        min_pixels = min_pixels if min_pixels is not None else self.min_pixels
+        max_pixels = max_pixels if max_pixels is not None else self.max_pixels
+
+        factor = self.patch_size * self.merge_size
+        resized_height, resized_width = smart_resize(
+            height,
+            width,
+            factor,
+            min_pixels=min_pixels,
+            max_pixels=max_pixels,
+        )
+        grid_h = resized_height // self.patch_size
+        grid_w = resized_width // self.patch_size
+        return grid_h * grid_w

src/mflux/models/qwen/tokenizer/qwen_tokenizer.py

@@ -14,7 +14,7 @@ class TokenizerQwen:
         tokens = self.tokenizer(
             formatted_text,
             max_length=self.max_length + self.template_start_idx,
-            padding=True,
+            padding=False,
             truncation=True,
             return_tensors="np",
         )

src/mflux/models/qwen/tokenizer/qwen_tokenizer_handler.py

@@ -7,13 +7,6 @@ from mflux.utils.download import snapshot_download
 
 
 class QwenTokenizerHandler:
-    """
-    Handler for Qwen tokenizer, following the same pattern as TokenizerHandler.
-
-    This loads the Qwen2Tokenizer from the model repository and wraps it
-    in our TokenizerQwen class for consistent interface.
-    """
-
     def __init__(
         self,
         repo_id: str,
@@ -34,7 +27,6 @@ class QwenTokenizerHandler:
 
     @staticmethod
     def _download_or_get_cached_tokenizer(repo_id: str) -> Path:
-        """Download or get cached tokenizer files."""
         return Path(
             snapshot_download(
                 repo_id=repo_id,

src/mflux/models/qwen/tokenizer/qwen_vision_language_processor.py

@@ -0,0 +1,103 @@
+from typing import Optional, Union
+
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+
+from mflux.models.qwen.tokenizer.qwen_image_processor import QwenImageProcessor
+
+
+class QwenVisionLanguageProcessor:
+    def __init__(
+        self,
+        tokenizer,
+        image_processor: Optional[QwenImageProcessor] = None,
+        image_token: str = "<|image_pad|>",
+        video_token: str = "<|video_pad|>",
+    ):
+        self.tokenizer = tokenizer
+        self.image_processor = image_processor or QwenImageProcessor()
+
+        self.image_token = image_token
+        self.video_token = video_token
+        self.image_token_id = (
+            tokenizer.image_token_id
+            if hasattr(tokenizer, "image_token_id")
+            else tokenizer.convert_tokens_to_ids(self.image_token)
+        )
+        self.video_token_id = (
+            tokenizer.video_token_id
+            if hasattr(tokenizer, "video_token_id")
+            else tokenizer.convert_tokens_to_ids(self.video_token)
+        )
+
+    def __call__(
+        self,
+        images: Optional[Union[Image.Image, list[Image.Image]]] = None,
+        text: Optional[Union[str, list[str]]] = None,
+        padding: bool = True,
+        return_tensors: Optional[str] = None,
+    ) -> dict:
+        image_inputs = {}
+        if images is not None:
+            pixel_values, image_grid_thw = self.image_processor.preprocess(images)
+            image_inputs = {
+                "pixel_values": pixel_values,
+                "image_grid_thw": image_grid_thw,
+            }
+
+        if text is not None:
+            if not isinstance(text, list):
+                text = [text]
+            text = text.copy()
+
+            if images is not None:
+                merge_length = self.image_processor.merge_size**2
+                index = 0
+                for i in range(len(text)):
+                    while self.image_token in text[i]:
+                        if index < len(image_grid_thw):
+                            num_image_tokens = int(np.prod(image_grid_thw[index])) // merge_length
+                            text[i] = text[i].replace(
+                                self.image_token,
+                                "<|placeholder|>" * num_image_tokens,
+                                1,
+                            )
+                            index += 1
+                        else:
+                            break
+                    text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+            text_inputs = self.tokenizer(
+                text,
+                padding=padding,
+                return_tensors="pt" if return_tensors == "pt" else "np",
+            )
+
+            if return_tensors == "pt":
+                import torch
+
+                if isinstance(text_inputs["input_ids"], torch.Tensor):
+                    input_ids = mx.array(text_inputs["input_ids"].numpy())
+                    attention_mask = mx.array(text_inputs["attention_mask"].numpy())
+                else:
+                    input_ids = mx.array(text_inputs["input_ids"])
+                    attention_mask = mx.array(text_inputs["attention_mask"])
+            else:
+                if isinstance(text_inputs["input_ids"], np.ndarray):
+                    input_ids = mx.array(text_inputs["input_ids"])
+                    attention_mask = mx.array(text_inputs["attention_mask"])
+                else:
+                    input_ids = mx.array(np.array(text_inputs["input_ids"]))
+                    attention_mask = mx.array(np.array(text_inputs["attention_mask"]))
+        else:
+            input_ids = None
+            attention_mask = None
+
+        result = {**image_inputs}
+        if input_ids is not None:
+            result["input_ids"] = input_ids
+        if attention_mask is not None:
+            result["attention_mask"] = attention_mask
+
+        return result

src/mflux/models/qwen/tokenizer/qwen_vision_language_tokenizer.py

@@ -0,0 +1,142 @@
+import math
+from pathlib import Path
+from typing import Union
+
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+
+from mflux.models.qwen.tokenizer.qwen_vision_language_processor import QwenVisionLanguageProcessor
+
+
+class QwenVisionLanguageTokenizer:
+    def __init__(
+        self,
+        processor: QwenVisionLanguageProcessor,
+        max_length: int = 1024,
+        use_picture_prefix: bool = True,
+    ):
+        self.processor = processor
+        self.max_length = max_length
+        self.use_picture_prefix = use_picture_prefix
+
+        if use_picture_prefix:
+            self.edit_template = (
+                "<|im_start|>system\n"
+                "Describe the key features of the input image (color, shape, size, texture, objects, background), "
+                "then explain how the user's text instruction should alter or modify the image. "
+                "Generate a new image that meets the user's requirements while maintaining consistency "
+                "with the original input where appropriate.<|im_end|>\n"
+                "<|im_start|>user\n"
+                "{}<|im_end|>\n"
+                "<|im_start|>assistant\n"
+            )
+        else:
+            self.edit_template = (
+                "<|im_start|>system\n"
+                "Describe the key features of the input image (color, shape, size, texture, objects, background), "
+                "then explain how the user's text instruction should alter or modify the image. "
+                "Generate a new image that meets the user's requirements while maintaining consistency "
+                "with the original input where appropriate.<|im_end|>\n"
+                "<|im_start|>user\n"
+                "<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n"
+                "<|im_start|>assistant\n"
+            )
+        self.edit_template_start_idx = 64
+
+    def tokenize_with_image(
+        self,
+        prompt: str,
+        image: Union[Image.Image, np.ndarray, str, list],
+        vl_width: int | None = None,
+        vl_height: int | None = None,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array]:
+        # Normalize image to list format
+        if not isinstance(image, list):
+            images = [image]
+        else:
+            images = image
+
+        # Format prompt based on tokenizer mode
+        if self.use_picture_prefix:
+            # Edit Plus format: Add "Picture N:" prefix for each image
+            # For multiple images: "Picture 1: ... Picture 2: ... Picture N: ..."
+            img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
+            base_img_prompt = ""
+            for i in range(len(images)):
+                base_img_prompt += img_prompt_template.format(i + 1)
+            formatted_text = self.edit_template.format(base_img_prompt + prompt)
+        else:
+            # Regular Edit format: Vision tokens already in template
+            # Just format with user prompt directly
+            formatted_text = self.edit_template.format(prompt)
+
+        # Process images: convert to PIL Images and resize to CONDITION_IMAGE_SIZE
+        CONDITION_IMAGE_SIZE = 384 * 384
+
+        processed_images = []
+        for img in images:
+            # Convert to PIL Image
+            if isinstance(img, (str, Path)):
+                img = Image.open(img).convert("RGB")
+            elif isinstance(img, np.ndarray):
+                img = Image.fromarray(img)
+            elif not isinstance(img, Image.Image):
+                raise ValueError(f"Unsupported image type: {type(img)}")
+
+            # Resize to CONDITION_IMAGE_SIZE (384×384) maintaining aspect ratio
+            img_w, img_h = img.size
+            ratio = img_w / img_h
+            condition_width = math.sqrt(CONDITION_IMAGE_SIZE * ratio)
+            condition_height = condition_width / ratio
+            condition_width = round(condition_width / 32) * 32
+            condition_height = round(condition_height / 32) * 32
+
+            img = img.resize((int(condition_width), int(condition_height)), Image.BICUBIC)
+            processed_images.append(img)
+
+        # Use our MLX processor for both text and images
+        model_inputs = self.processor(
+            text=[formatted_text],
+            images=processed_images,
+            padding=True,
+            return_tensors=None,  # Return numpy/MLX arrays, not PyTorch
+        )
+
+        grid_thw = model_inputs["image_grid_thw"][0]
+        factor = 14 * 2
+        self._vl_image_width = int(grid_thw[2]) * factor
+        self._vl_image_height = int(grid_thw[1]) * factor
+
+        # Convert to MLX arrays if needed
+        input_ids = model_inputs["input_ids"]
+        attention_mask = model_inputs["attention_mask"]
+        pixel_values = mx.array(model_inputs["pixel_values"])
+        image_grid_thw = mx.array(model_inputs["image_grid_thw"])
+
+        return input_ids, attention_mask, pixel_values, image_grid_thw
+
+    def tokenize_text_only(self, prompt: str) -> tuple[mx.array, mx.array]:
+        # Use the regular text-only template
+        text_template = (
+            "<|im_start|>system\n"
+            "Describe the image by detailing the color, shape, size, texture, quantity, text, "
+            "spatial relationships of the objects and background:<|im_end|>\n"
+            "<|im_start|>user\n{}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        )
+
+        formatted_text = text_template.format(prompt)
+        tokens = self.processor.tokenizer(
+            formatted_text,
+            max_length=self.max_length + 34,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+        )
+
+        # Convert PyTorch tensors to MLX arrays
+        input_ids = mx.array(tokens["input_ids"].numpy())
+        attention_mask = mx.array(tokens["attention_mask"].numpy())
+
+        return input_ids, attention_mask

src/mflux/models/qwen/variants/edit/__init__.py

@@ -0,0 +1,5 @@
+# Qwen Image Edit variant
+from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit
+from mflux.models.qwen.variants.edit.qwen_image_edit_plus import QwenImageEditPlus
+
+__all__ = ["QwenImageEdit", "QwenImageEditPlus"]

src/mflux/models/qwen/variants/edit/qwen_image_edit.py

@@ -0,0 +1,263 @@
+import math
+
+import mlx.core as mx
+from mlx import nn
+from PIL import Image
+from tqdm import tqdm
+
+from mflux.callbacks.callbacks import Callbacks
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
+from mflux.latent_creator.latent_creator import LatentCreator
+from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder
+from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
+from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE
+from mflux.models.qwen.qwen_edit_initializer import QwenImageEditInitializer
+from mflux.models.qwen.variants.edit.utils.qwen_edit_util import QwenEditUtil
+from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.generated_image import GeneratedImage
+from mflux.post_processing.image_util import ImageUtil
+
+
+class QwenImageEdit(nn.Module):
+    vae: QwenVAE
+    transformer: QwenTransformer
+    text_encoder: QwenTextEncoder
+
+    def __init__(
+        self,
+        quantize: int | None = None,
+        local_path: str | None = None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+    ):
+        super().__init__()
+        QwenImageEditInitializer.init(
+            qwen_model=self,
+            model_config=ModelConfig.qwen_image_edit(),
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
+        )
+
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        config: Config,
+        negative_prompt: str | None = None,
+    ) -> GeneratedImage:
+        runtime_config, vl_width, vl_height = self._compute_dimensions(config)
+        timesteps = runtime_config.scheduler.timesteps
+        time_steps = tqdm(range(len(timesteps)))
+
+        # 1. Create initial latents
+        latents = LatentCreator.create(
+            seed=seed,
+            height=runtime_config.height,
+            width=runtime_config.width,
+        )
+
+        # 2. Encode the prompt
+        prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = (
+            self._debug_encode_prompts_with_image(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                image_path=runtime_config.image_path,
+                runtime_config=runtime_config,
+                vl_width=vl_width,
+                vl_height=vl_height,
+            )
+        )
+
+        # 3. Generate image conditioning latents
+        image_path = str(runtime_config.image_path) if runtime_config.image_path else None
+        static_image_latents, qwen_image_ids, cond_h_patches, cond_w_patches, num_images = (
+            QwenEditUtil.create_image_conditioning_latents(
+                vae=self.vae,
+                height=runtime_config.height,
+                width=runtime_config.width,
+                image_paths=image_path,
+                vl_width=vl_width,
+                vl_height=vl_height,
+            )
+        )
+
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=runtime_config,
+        )
+
+        for t in time_steps:
+            try:
+                # 4.t Concatenate the updated latents with the static image latents
+                hidden_states = mx.concatenate([latents, static_image_latents], axis=1)
+                hidden_states_neg = mx.concatenate([latents, static_image_latents], axis=1)  # noqa: F841
+
+                # 5.t Predict the noise
+                noise = self.transformer(
+                    t=t,
+                    config=runtime_config,
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_hidden_states_mask=prompt_mask,
+                    qwen_image_ids=qwen_image_ids,
+                    cond_image_grid=(1, cond_h_patches, cond_w_patches),
+                )[:, : latents.shape[1]]
+
+                noise_negative = self.transformer(
+                    t=t,
+                    config=runtime_config,
+                    hidden_states=hidden_states_neg,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    encoder_hidden_states_mask=negative_prompt_mask,
+                    qwen_image_ids=qwen_image_ids,
+                    cond_image_grid=(1, cond_h_patches, cond_w_patches),
+                )[:, : latents.shape[1]]
+
+                guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, runtime_config.guidance)
+
+                # 6.t Take one denoise step
+                latents = runtime_config.scheduler.step(
+                    model_output=guided_noise,
+                    timestep=t,
+                    sample=latents,
+                )
+
+                # (Optional) Call subscribers in-loop
+                Callbacks.in_loop(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=runtime_config,
+                    time_steps=time_steps,
+                )
+
+                # (Optional) Evaluate to enable progress tracking
+                mx.eval(latents)
+
+            except KeyboardInterrupt:  # noqa: PERF203
+                Callbacks.interruption(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=runtime_config,
+                    time_steps=time_steps,
+                )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(timesteps)}")
+
+        # (Optional) Call subscribers after loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=runtime_config,
+        )
+
+        # 7. Decode the latent array and return the image
+        latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width)
+        decoded = self.vae.decode(latents)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            config=runtime_config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            image_path=runtime_config.image_path,
+            generation_time=time_steps.format_dict["elapsed"],
+            negative_prompt=negative_prompt,
+        )
+
+    def _debug_encode_prompts_with_image(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        image_path: str,
+        runtime_config,
+        vl_width: int | None = None,
+        vl_height: int | None = None,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array]:
+        tokenizer = self.qwen_vl_tokenizer
+        image = Image.open(image_path).convert("RGB")
+
+        pos_input_ids, pos_attention_mask, pos_pixel_values, pos_image_grid_thw = tokenizer.tokenize_with_image(
+            prompt, image, vl_width=vl_width, vl_height=vl_height
+        )
+
+        pos_hidden_states = self.qwen_vl_encoder(
+            input_ids=pos_input_ids,
+            attention_mask=pos_attention_mask,
+            pixel_values=pos_pixel_values,
+            image_grid_thw=pos_image_grid_thw,
+        )
+
+        neg_prompt = negative_prompt if negative_prompt is not None else ""
+        neg_input_ids, neg_attention_mask, neg_pixel_values, neg_image_grid_thw = tokenizer.tokenize_with_image(
+            neg_prompt, image, vl_width=vl_width, vl_height=vl_height
+        )
+
+        neg_hidden_states = self.qwen_vl_encoder(
+            input_ids=neg_input_ids,
+            attention_mask=neg_attention_mask,
+            pixel_values=neg_pixel_values,
+            image_grid_thw=neg_image_grid_thw,
+        )
+
+        return (
+            pos_hidden_states[0].astype(mx.float16),  # prompt_embeds
+            pos_hidden_states[1].astype(mx.float16),  # prompt_mask
+            neg_hidden_states[0].astype(mx.float16),  # negative_prompt_embeds
+            neg_hidden_states[1].astype(mx.float16),  # negative_prompt_mask
+        )
+
+    def _compute_dimensions(
+        self,
+        config: Config,
+    ) -> tuple[RuntimeConfig, int, int]:
+        image = ImageUtil.load_image(config.image_path).convert("RGB")
+        image_size = image.size
+
+        target_area = 1024 * 1024
+        ratio = image_size[0] / image_size[1]
+        calculated_width = math.sqrt(target_area * ratio)
+        calculated_height = calculated_width / ratio
+        calculated_width = round(calculated_width / 32) * 32
+        calculated_height = round(calculated_height / 32) * 32
+
+        use_height = config.height or calculated_height
+        use_width = config.width or calculated_width
+
+        vae_scale_factor = 8
+        multiple_of = vae_scale_factor * 2
+        use_width = use_width // multiple_of * multiple_of
+        use_height = use_height // multiple_of * multiple_of
+
+        final_config = Config(
+            height=use_height,
+            width=use_width,
+            image_path=config.image_path,
+            num_inference_steps=config.num_inference_steps,
+            guidance=config.guidance,
+            scheduler=config.scheduler_str,
+        )
+        runtime_config = RuntimeConfig(final_config, self.model_config)
+
+        vl_width = calculated_width
+        vl_height = calculated_height
+
+        return runtime_config, int(vl_width), int(vl_height)

src/mflux/models/qwen/variants/edit/qwen_image_edit_plus.py

@@ -0,0 +1,285 @@
+import math
+
+import mlx.core as mx
+from mlx import nn
+from tqdm import tqdm
+
+from mflux.callbacks.callbacks import Callbacks
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.config.runtime_config import RuntimeConfig
+from mflux.error.exceptions import StopImageGenerationException
+from mflux.latent_creator.latent_creator import LatentCreator
+from mflux.models.qwen.model.qwen_text_encoder.qwen_text_encoder import QwenTextEncoder
+from mflux.models.qwen.model.qwen_transformer.qwen_transformer import QwenTransformer
+from mflux.models.qwen.model.qwen_vae.qwen_vae import QwenVAE
+from mflux.models.qwen.qwen_edit_initializer import QwenImageEditInitializer
+from mflux.models.qwen.variants.edit.utils.qwen_edit_util import QwenEditUtil
+from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.generated_image import GeneratedImage
+from mflux.post_processing.image_util import ImageUtil
+
+
+class QwenImageEditPlus(nn.Module):
+    vae: QwenVAE
+    transformer: QwenTransformer
+    text_encoder: QwenTextEncoder
+
+    def __init__(
+        self,
+        quantize: int | None = None,
+        local_path: str | None = None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+    ):
+        super().__init__()
+        QwenImageEditInitializer.init(
+            qwen_model=self,
+            model_config=ModelConfig.qwen_image_edit_plus(),
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
+        )
+
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        config: Config,
+        negative_prompt: str | None = None,
+        image_paths: list[str] | None = None,
+    ) -> GeneratedImage:
+        if image_paths is None:
+            image_paths = [config.image_path]
+
+        runtime_config, vl_width, vl_height, vae_width, vae_height = self._compute_dimensions(config, image_paths)
+        timesteps = runtime_config.scheduler.timesteps
+        time_steps = tqdm(range(len(timesteps)))
+
+        # 1. Create initial latents
+        latents = LatentCreator.create(
+            seed=seed,
+            height=runtime_config.height,
+            width=runtime_config.width,
+        )
+
+        # 2. Encode the prompt
+        prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = self._encode_prompts_with_image_plus(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image_paths=image_paths,
+            runtime_config=runtime_config,
+            vl_width=vl_width,
+            vl_height=vl_height,
+        )
+
+        # 3. Generate image conditioning latents
+        static_image_latents, qwen_image_ids, cond_h_patches, cond_w_patches, num_images = (
+            QwenEditUtil.create_image_conditioning_latents(
+                vae=self.vae,
+                height=vae_height,
+                width=vae_width,
+                image_paths=image_paths,
+                vl_width=vl_width,
+                vl_height=vl_height,
+            )
+        )
+
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=runtime_config,
+        )
+
+        for t in time_steps:
+            try:
+                # 4.t Concatenate the updated latents with the static image latents
+                hidden_states = mx.concatenate([latents, static_image_latents], axis=1)
+                hidden_states_neg = mx.concatenate([latents, static_image_latents], axis=1)
+
+                # 5.t Predict the noise
+                if num_images > 1:
+                    cond_image_grid = [(1, cond_h_patches, cond_w_patches) for _ in range(num_images)]
+                else:
+                    cond_image_grid = (1, cond_h_patches, cond_w_patches)
+
+                noise = self.transformer(
+                    t=t,
+                    config=runtime_config,
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_hidden_states_mask=prompt_mask,
+                    qwen_image_ids=qwen_image_ids,
+                    cond_image_grid=cond_image_grid,
+                )[:, : latents.shape[1]]
+                noise_negative = self.transformer(
+                    t=t,
+                    config=runtime_config,
+                    hidden_states=hidden_states_neg,
+                    encoder_hidden_states=negative_prompt_embeds,
+                    encoder_hidden_states_mask=negative_prompt_mask,
+                    qwen_image_ids=qwen_image_ids,
+                    cond_image_grid=cond_image_grid,
+                )[:, : latents.shape[1]]
+                guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, runtime_config.guidance)
+
+                # 6.t Take one denoise step
+                latents = runtime_config.scheduler.step(
+                    model_output=guided_noise,
+                    timestep=t,
+                    sample=latents,
+                )
+
+                # (Optional) Call subscribers in-loop
+                Callbacks.in_loop(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=runtime_config,
+                    time_steps=time_steps,
+                )
+
+                # (Optional) Evaluate to enable progress tracking
+                mx.eval(latents)
+
+            except KeyboardInterrupt:  # noqa: PERF203
+                Callbacks.interruption(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=runtime_config,
+                    time_steps=time_steps,
+                )
+                raise StopImageGenerationException(f"Stopping image generation at step {t + 1}/{len(timesteps)}")
+
+        # (Optional) Call subscribers after loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=runtime_config,
+        )
+
+        # 7. Decode the latent array and return the image
+        latents = ArrayUtil.unpack_latents(latents=latents, height=runtime_config.height, width=runtime_config.width)
+        decoded = self.vae.decode(latents)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            config=runtime_config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            image_path=runtime_config.image_path,
+            generation_time=time_steps.format_dict["elapsed"],
+            negative_prompt=negative_prompt,
+        )
+
+    def _encode_prompts_with_image_plus(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        image_paths: list[str],
+        runtime_config,
+        vl_width: int | None = None,
+        vl_height: int | None = None,
+    ) -> tuple[mx.array, mx.array, mx.array, mx.array]:
+        tokenizer = self.qwen_vl_tokenizer
+        pos_input_ids, pos_attention_mask, pos_pixel_values, pos_image_grid_thw = tokenizer.tokenize_with_image(
+            prompt, image_paths, vl_width=vl_width, vl_height=vl_height
+        )
+
+        pos_hidden_states = self.qwen_vl_encoder(
+            input_ids=pos_input_ids,
+            attention_mask=pos_attention_mask,
+            pixel_values=pos_pixel_values,
+            image_grid_thw=pos_image_grid_thw,
+        )
+        mx.eval(pos_hidden_states[0])
+        mx.eval(pos_hidden_states[1])
+
+        neg_prompt = negative_prompt if negative_prompt is not None else ""
+
+        # Use real MLX tokenizer for negative prompt
+        neg_input_ids, neg_attention_mask, neg_pixel_values, neg_image_grid_thw = tokenizer.tokenize_with_image(
+            neg_prompt, image_paths, vl_width=vl_width, vl_height=vl_height
+        )
+
+        neg_hidden_states = self.qwen_vl_encoder(
+            input_ids=neg_input_ids,
+            attention_mask=neg_attention_mask,
+            pixel_values=neg_pixel_values,
+            image_grid_thw=neg_image_grid_thw,
+        )
+        mx.eval(neg_hidden_states[0])
+        mx.eval(neg_hidden_states[1])
+
+        final_prompt_embeds = pos_hidden_states[0].astype(mx.float16)
+        final_prompt_mask = pos_hidden_states[1].astype(mx.float16)
+
+        return (
+            final_prompt_embeds,  # prompt_embeds
+            final_prompt_mask,  # prompt_mask
+            neg_hidden_states[0].astype(mx.float16),  # negative_prompt_embeds
+            neg_hidden_states[1].astype(mx.float16),  # negative_prompt_mask
+        )
+
+    def _compute_dimensions(
+        self,
+        config: Config,
+        image_paths: list[str],
+    ) -> tuple[RuntimeConfig, int, int, int, int]:
+        last_image = ImageUtil.load_image(image_paths[-1]).convert("RGB")
+        image_size = last_image.size
+
+        target_area = 1024 * 1024
+        ratio = image_size[0] / image_size[1]
+        calculated_width = math.sqrt(target_area * ratio)
+        calculated_height = calculated_width / ratio
+        calculated_width = round(calculated_width / 32) * 32
+        calculated_height = round(calculated_height / 32) * 32
+
+        use_height = config.height or calculated_height
+        use_width = config.width or calculated_width
+
+        vae_scale_factor = 8
+        multiple_of = vae_scale_factor * 2
+        use_width = use_width // multiple_of * multiple_of
+        use_height = use_height // multiple_of * multiple_of
+
+        final_config = Config(
+            height=use_height,
+            width=use_width,
+            image_path=config.image_path,
+            num_inference_steps=config.num_inference_steps,
+            guidance=config.guidance,
+            scheduler=config.scheduler_str,
+        )
+        runtime_config = RuntimeConfig(final_config, self.model_config)
+
+        CONDITION_IMAGE_SIZE = 384 * 384
+        condition_ratio = image_size[0] / image_size[1]
+        vl_width = math.sqrt(CONDITION_IMAGE_SIZE * condition_ratio)
+        vl_height = vl_width / condition_ratio
+        vl_width = round(vl_width / 32) * 32
+        vl_height = round(vl_height / 32) * 32
+
+        VAE_IMAGE_SIZE = 1024 * 1024
+        vae_ratio = image_size[0] / image_size[1]
+        vae_width = math.sqrt(VAE_IMAGE_SIZE * vae_ratio)
+        vae_height = vae_width / vae_ratio
+        vae_width = round(vae_width / 32) * 32
+        vae_height = round(vae_height / 32) * 32
+
+        return runtime_config, int(vl_width), int(vl_height), int(vae_width), int(vae_height)

src/mflux/models/qwen/variants/edit/utils/__init__.py

@@ -0,0 +1 @@
+# Qwen Edit utilities

src/mflux/models/qwen/variants/edit/utils/qwen_edit_util.py

@@ -0,0 +1,106 @@
+import os
+
+import mlx.core as mx
+
+from mflux.latent_creator.latent_creator import LatentCreator
+from mflux.post_processing.array_util import ArrayUtil
+
+
+class QwenEditUtil:
+    @staticmethod
+    def create_image_conditioning_latents(
+        vae,
+        height: int,
+        width: int,
+        image_paths: list[str] | str,
+        vl_width: int | None = None,
+        vl_height: int | None = None,
+    ) -> tuple[mx.array, mx.array, int, int, int]:
+        if not isinstance(image_paths, list):
+            image_paths = [str(image_paths)]
+
+        if vl_width is not None and vl_height is not None:
+            calc_w, calc_h = vl_width, vl_height
+        else:
+            from mflux.post_processing.image_util import ImageUtil
+
+            pil_image = ImageUtil.load_image(image_paths[-1]).convert("RGB")
+            img_w, img_h = pil_image.size
+            target_area_env = os.getenv("MFLUX_QWEN_VL_TARGET_AREA")
+            target_area = int(target_area_env) if target_area_env is not None else 1024 * 1024
+            ratio = img_w / img_h
+            calc_w = int(round(((target_area * ratio) ** 0.5) / 32) * 32)
+            calc_h = int(round((calc_w / ratio) / 32) * 32)
+
+        all_image_latents = []
+        for image_path in image_paths:
+            input_image = LatentCreator.encode_image(
+                vae=vae,
+                image_path=image_path,
+                height=calc_h,
+                width=calc_w,
+            )
+
+            image_latents = ArrayUtil.pack_latents(
+                latents=input_image,
+                height=calc_h,
+                width=calc_w,
+                num_channels_latents=16,
+            )
+            all_image_latents.append(image_latents)
+
+        image_latents = mx.concatenate(all_image_latents, axis=1)
+
+        all_image_ids = []
+        for _ in image_paths:
+            image_ids = QwenEditUtil._create_image_ids(
+                height=calc_h,
+                width=calc_w,
+            )
+            all_image_ids.append(image_ids)
+        image_ids = mx.concatenate(all_image_ids, axis=1)
+
+        cond_h_patches = calc_h // 16
+        cond_w_patches = calc_w // 16
+        num_images = len(image_paths)
+        return image_latents, image_ids, cond_h_patches, cond_w_patches, num_images
+
+    @staticmethod
+    def _create_image_ids(
+        height: int,
+        width: int,
+    ) -> mx.array:
+        latent_height = height // 16
+        latent_width = width // 16
+
+        image_ids = mx.zeros((latent_height, latent_width, 3))
+
+        row_coords = mx.arange(0, latent_height)[:, None]
+        row_coords = mx.broadcast_to(row_coords, (latent_height, latent_width))
+        image_ids = mx.concatenate(
+            [
+                image_ids[:, :, :1],
+                row_coords[:, :, None],
+                image_ids[:, :, 2:],
+            ],
+            axis=2,
+        )
+
+        col_coords = mx.arange(0, latent_width)[None, :]
+        col_coords = mx.broadcast_to(col_coords, (latent_height, latent_width))
+        image_ids = mx.concatenate(
+            [
+                image_ids[:, :, :2],
+                col_coords[:, :, None],
+            ],
+            axis=2,
+        )
+
+        image_ids = mx.reshape(image_ids, (latent_height * latent_width, 3))
+
+        first_dim = mx.ones((image_ids.shape[0], 1))
+        image_ids = mx.concatenate([first_dim, image_ids[:, 1:]], axis=1)
+
+        image_ids = mx.expand_dims(image_ids, axis=0)
+
+        return image_ids

src/mflux/models/qwen/variants/edit/utils/qwen_prompt_rewriter.py

@@ -0,0 +1,116 @@
+from typing import Optional
+
+
+class QwenPromptRewriter:
+    @staticmethod
+    def detect_language(prompt: str) -> str:
+        chinese_ranges = [
+            ("\u4e00", "\u9fff"),
+        ]
+        for char in prompt:
+            if any(start <= char <= end for start, end in chinese_ranges):
+                return "zh"
+        return "en"
+
+    @staticmethod
+    def enhance_prompt_en(original_prompt: str) -> str:
+        # Remove extra whitespace and normalize
+        prompt = original_prompt.strip()
+
+        # Add style and quality enhancers if not present
+        quality_terms = ["ultra hd", "4k", "cinematic", "high quality", "detailed"]
+        has_quality = any(term in prompt.lower() for term in quality_terms)
+
+        if not has_quality:
+            prompt += ", ultra HD, 4K, cinematic composition, high quality, detailed"
+
+        # Enhance editing-specific instructions
+        if any(word in prompt.lower() for word in ["change", "replace", "modify", "edit"]):
+            if "maintain" not in prompt.lower() and "keep" not in prompt.lower():
+                prompt += ", maintain original composition and lighting"
+
+        # Add style consistency for artistic transformations
+        art_styles = ["watercolor", "oil painting", "sketch", "digital art", "anime", "cartoon"]
+        if any(style in prompt.lower() for style in art_styles):
+            prompt += ", consistent artistic style throughout"
+
+        return prompt
+
+    @staticmethod
+    def enhance_prompt_zh(original_prompt: str) -> str:
+        # Remove extra whitespace and normalize
+        prompt = original_prompt.strip()
+
+        # Add quality enhancers for Chinese prompts
+        quality_terms = ["超清", "4k", "高质量", "精细", "电影级"]
+        has_quality = any(term in prompt for term in quality_terms)
+
+        if not has_quality:
+            prompt += "，超清，4K，电影级构图，高质量，精细"
+
+        # Enhance editing-specific instructions
+        edit_terms = ["改变", "替换", "修改", "编辑", "变成"]
+        if any(term in prompt for term in edit_terms):
+            if "保持" not in prompt and "维持" not in prompt:
+                prompt += "，保持原有构图和光线"
+
+        return prompt
+
+    @staticmethod
+    def enhance_edit_prompt(prompt: str, context: Optional[str] = None) -> str:
+        language = QwenPromptRewriter.detect_language(prompt)
+
+        if language == "zh":
+            enhanced = QwenPromptRewriter.enhance_prompt_zh(prompt)
+        else:
+            enhanced = QwenPromptRewriter.enhance_prompt_en(prompt)
+
+        # Add edit-specific stability improvements
+        edit_keywords = {
+            "en": ["transform", "change", "modify", "edit", "replace", "add", "remove"],
+            "zh": ["转换", "改变", "修改", "编辑", "替换", "添加", "移除", "变成"],
+        }
+
+        current_keywords = edit_keywords[language]
+        has_edit_instruction = any(keyword in enhanced.lower() for keyword in current_keywords)
+
+        if has_edit_instruction:
+            if language == "zh":
+                if "保持细节" not in enhanced:
+                    enhanced += "，保持细节准确"
+            else:
+                if "preserve details" not in enhanced.lower():
+                    enhanced += ", preserve fine details"
+
+        return enhanced
+
+    @staticmethod
+    def create_edit_template(
+        action: str, target: str, replacement: Optional[str] = None, style: Optional[str] = None, language: str = "en"
+    ) -> str:
+        if language == "zh":
+            templates = {
+                "replace": f"将{target}替换为{replacement}" if replacement else f"替换{target}",
+                "add": f"在图像中添加{target}",
+                "remove": f"从图像中移除{target}",
+                "transform": f"将图像转换为{target}风格",
+            }
+            quality_suffix = "，保持原有构图，超清，4K，精细细节"
+        else:
+            templates = {
+                "replace": f"Replace {target} with {replacement}" if replacement else f"Replace {target}",
+                "add": f"Add {target} to the image",
+                "remove": f"Remove {target} from the image",
+                "transform": f"Transform the image into {target} style",
+            }
+            quality_suffix = ", maintain original composition, ultra HD, 4K, fine details"
+
+        base_prompt = templates.get(action, target)
+
+        if style:
+            if language == "zh":
+                base_prompt += f"，{style}风格"
+            else:
+                base_prompt += f", {style} style"
+
+        return base_prompt + quality_suffix

src/mflux/models/qwen/variants/txt2img/qwen_image.py

@@ -52,10 +52,6 @@ class QwenImage(nn.Module):
         prompt: str,
         config: Config,
         negative_prompt: str | None = None,
-        prompt_embeds: mx.array | None = None,
-        prompt_mask: mx.array | None = None,
-        negative_prompt_embeds: mx.array | None = None,
-        negative_prompt_mask: mx.array | None = None,
     ) -> GeneratedImage:
         # 0. Create a new runtime config based on the model type and input parameters
         runtime_config = RuntimeConfig(config, self.model_config)
@@ -74,7 +70,7 @@ class QwenImage(nn.Module):
             ),
         )
 
-        # 2. Encode the prompt
+        # 2. Encode the prompt (using native MLX encoding)
         prompt_embeds, prompt_mask, negative_prompt_embeds, negative_prompt_mask = QwenPromptEncoder.encode_prompt(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -111,7 +107,7 @@ class QwenImage(nn.Module):
                     encoder_hidden_states=negative_prompt_embeds,
                     encoder_hidden_states_mask=negative_prompt_mask,
                 )
-                guided_noise = QwenImage._compute_guided_noise(noise, noise_negative, runtime_config.guidance)
+                guided_noise = QwenImage.compute_guided_noise(noise, noise_negative, runtime_config.guidance)
 
                 # 4.t Take one denoise step
                 latents = runtime_config.scheduler.step(
@@ -173,7 +169,7 @@ class QwenImage(nn.Module):
         QwenModelSaver.save_model(self, self.bits, base_path)
 
     @staticmethod
-    def _compute_guided_noise(
+    def compute_guided_noise(
         noise: mx.array,
         noise_negative: mx.array,
         guidance: float,

src/mflux/models/qwen/weights/qwen_weight_handler.py

@@ -1,11 +1,17 @@
 import json
+from collections.abc import Callable
 from pathlib import Path
 
 import mlx.core as mx
+import torch
+from mlx.utils import tree_unflatten
 from safetensors.mlx import load_file as mlx_load_file
+from safetensors.torch import load_file as torch_load_file
 
+from mflux.models.common.weights.mapping.weight_mapper import WeightMapper
+from mflux.models.common.weights.mapping.weight_mapping import WeightTarget
 from mflux.models.flux.weights.weight_handler import MetaData, WeightHandler
-from mflux.models.qwen.weights.qwen_weight_util import QwenWeightUtil
+from mflux.models.qwen.weights.qwen_weight_mapping import QwenWeightMapping
 
 
 class QwenWeightHandler:
@@ -21,13 +27,6 @@ class QwenWeightHandler:
         self.vae = vae
         self.meta_data = meta_data
 
-    def num_transformer_blocks(self) -> int:
-        return (
-            len(self.transformer["transformer_blocks"])
-            if self.transformer and "transformer_blocks" in self.transformer
-            else 0
-        )
-
     @staticmethod
     def load_regular_weights(
         repo_id: str | None = None,
@@ -36,9 +35,12 @@ class QwenWeightHandler:
         # Load the weights from disk, huggingface cache, or download from huggingface
         root_path = Path(local_path) if local_path else WeightHandler.download_or_get_cached_weights(repo_id)
 
+        # Determine if we should load visual weights (for Edit model)
+        load_visual_weights = repo_id and "edit" in repo_id.lower()
+
         # Load the weights
-        transformer, quantization_level, mflux_version = QwenWeightHandler.load_transformer(root_path=root_path)
-        qwen_text_encoder, _, _ = QwenWeightHandler._load_qwen_text_encoder(root_path=root_path)
+        transformer, quantization_level, mflux_version = QwenWeightHandler._load_transformer(root_path=root_path)
+        qwen_text_encoder, _, _ = QwenWeightHandler._load_qwen_text_encoder(root_path=root_path, load_visual_weights=load_visual_weights)  # fmt: off
         vae, _, _ = QwenWeightHandler._load_vae(root_path=root_path)
 
         return QwenWeightHandler(
@@ -54,612 +56,54 @@ class QwenWeightHandler:
         )
 
     @staticmethod
-    def load_transformer(root_path: Path) -> tuple[dict, int | None, str | None]:
-        flat = QwenWeightHandler._load_safetensors_shards(root_path / "transformer", loading_mode="multi_glob")
-
-        # Check if this is a saved quantized model (with metadata) or HuggingFace weights
-        # Saved models from mflux-save are already in MLX structure and don't need manual mapping
-        quantization_level = None
-        mflux_version = None
-
-        # Try to get metadata from the first weight shard
-        import mlx.core as mx
-        from mlx.utils import tree_unflatten
-
-        file_glob = sorted((root_path / "transformer").glob("*.safetensors"))
-        if file_glob:
-            data = mx.load(str(file_glob[0]), return_metadata=True)
-            if len(data) > 1:
-                quantization_level = data[1].get("quantization_level")
-                mflux_version = data[1].get("mflux_version")
-
-        # If this is a saved model (has metadata), use tree_unflatten directly
-        if quantization_level is not None or mflux_version is not None:
-            return tree_unflatten(list(flat.items())), quantization_level, mflux_version
-
-        # Otherwise, it's HuggingFace weights that need manual mapping
-        mapped_weights = QwenWeightHandler._manual_transformer_mapping(flat)
-        return mapped_weights, None, None
+    def _load_transformer(root_path: Path) -> tuple[dict, int | None, str | None]:
+        return QwenWeightHandler._load_component(
+            root_path=root_path,
+            component_name="transformer",
+            loading_mode="multi_glob",
+            mapping_getter=QwenWeightMapping.get_transformer_mapping,
+        )
 
     @staticmethod
-    def _load_qwen_text_encoder(root_path: Path) -> tuple[dict, int | None, str | None]:
-        import mlx.core as mx
-        from mlx.utils import tree_unflatten
-
-        # Check for saved model metadata FIRST to determine loading mode
-        quantization_level = None
-        mflux_version = None
-        file_glob = sorted((root_path / "text_encoder").glob("*.safetensors"))
-        if file_glob:
-            data = mx.load(str(file_glob[0]), return_metadata=True)
-            if len(data) > 1:
-                quantization_level = data[1].get("quantization_level")
-                mflux_version = data[1].get("mflux_version")
-
-        # If this is a saved model, load without expecting index.json
-        if quantization_level is not None or mflux_version is not None:
-            all_weights = QwenWeightHandler._load_safetensors_shards(
-                root_path / "text_encoder", loading_mode="multi_glob"
-            )
-            return tree_unflatten(list(all_weights.items())), quantization_level, mflux_version
-
-        # Otherwise, it's HuggingFace weights that need manual mapping
-        all_weights = QwenWeightHandler._load_safetensors_shards(root_path / "text_encoder", loading_mode="multi_json")
-        mapped_weights = QwenWeightHandler._manual_text_encoder_mapping(all_weights)
-        return mapped_weights, None, None
+    def _load_qwen_text_encoder(
+        root_path: Path, load_visual_weights: bool = False
+    ) -> tuple[dict, int | None, str | None]:
+        return QwenWeightHandler._load_component(
+            root_path=root_path,
+            component_name="text_encoder",
+            loading_mode="multi_json",
+            mapping_getter=QwenWeightMapping.get_text_encoder_mapping,
+        )
 
     @staticmethod
     def _load_vae(root_path: Path) -> tuple[dict, int | None, str | None]:
-        import mlx.core as mx
-        from mlx.utils import tree_unflatten
+        return QwenWeightHandler._load_component(
+            root_path=root_path,
+            component_name="vae",
+            loading_mode="single",
+            mapping_getter=QwenWeightMapping.get_vae_mapping,
+        )
 
-        weights = QwenWeightHandler._load_safetensors_shards(root_path / "vae", loading_mode="single")
+    @staticmethod
+    def _load_component(
+        root_path: Path,
+        component_name: str,
+        loading_mode: str,
+        mapping_getter: Callable[[], list[WeightTarget]],
+    ) -> tuple[dict, int | None, str | None]:
+        component_path = root_path / component_name
+        weights = QwenWeightHandler._load_safetensors_shards(component_path, loading_mode=loading_mode)
 
-        # Check for saved model metadata
-        quantization_level = None
-        mflux_version = None
-        file_glob = sorted((root_path / "vae").glob("*.safetensors"))
-        if file_glob:
-            data = mx.load(str(file_glob[0]), return_metadata=True)
-            if len(data) > 1:
-                quantization_level = data[1].get("quantization_level")
-                mflux_version = data[1].get("mflux_version")
-
-        # If this is a saved model, use tree_unflatten directly
+        # Check if this is a saved model (has metadata)
+        quantization_level, mflux_version = QwenWeightHandler._detect_metadata(component_path)
         if quantization_level is not None or mflux_version is not None:
-            return tree_unflatten(list(weights.items())), quantization_level, mflux_version
+            return QwenWeightHandler._load_saved_model_weights(weights, None, quantization_level, mflux_version)
 
-        # Otherwise, it's HuggingFace weights that need manual mapping and reshaping
-        reshaped_weights = [QwenWeightUtil.reshape_weights(k, v) for k, v in weights.items()]
-        reshaped_weights = QwenWeightUtil.flatten(reshaped_weights)
-        weights = dict(reshaped_weights)
-        mapped_weights = QwenWeightHandler._manual_flux_style_mapping(weights)
+        # Otherwise, it's HuggingFace weights that need mapping
+        mapping = mapping_getter()
+        mapped_weights = WeightMapper.apply_mapping(weights, mapping)
         return mapped_weights, None, None
 
-    @staticmethod
-    def _manual_flux_style_mapping(diffusers_weights: dict) -> dict:
-        weights = {}
-
-        # 1. Simple direct mappings (like Flux does)
-        weights["decoder"] = {}
-
-        # conv_in: decoder.conv_in.weight -> decoder.conv_in.conv3d.weight
-        weights["decoder"]["conv_in"] = {
-            "conv3d": {
-                "weight": diffusers_weights["decoder.conv_in.weight"],
-                "bias": diffusers_weights["decoder.conv_in.bias"],
-            }
-        }
-
-        # conv_out: decoder.conv_out.weight -> decoder.conv_out.conv3d.weight
-        weights["decoder"]["conv_out"] = {
-            "conv3d": {
-                "weight": diffusers_weights["decoder.conv_out.weight"],
-                "bias": diffusers_weights["decoder.conv_out.bias"],
-            }
-        }
-
-        # norm_out: decoder.norm_out.gamma -> decoder.norm_out.weight (flatten to 1D like MLX expects)
-        dec_gamma = diffusers_weights["decoder.norm_out.gamma"]
-        if len(dec_gamma.shape) > 1:
-            dec_gamma = mx.reshape(dec_gamma, (dec_gamma.shape[0],))
-        weights["decoder"]["norm_out"] = {"weight": dec_gamma}
-
-        # post_quant_conv: post_quant_conv.weight -> post_quant_conv.conv3d.weight
-        weights["post_quant_conv"] = {
-            "conv3d": {
-                "weight": diffusers_weights["post_quant_conv.weight"],
-                "bias": diffusers_weights["post_quant_conv.bias"],
-            }
-        }
-
-        # 2. Mid block (manual structure building)
-        weights["decoder"]["mid_block"] = {}
-
-        # mid_block.resnets (list of 2 resnets)
-        weights["decoder"]["mid_block"]["resnets"] = [{}, {}]
-        for i in range(2):
-            resnet = weights["decoder"]["mid_block"]["resnets"][i]
-            # conv1.weight -> conv1.conv3d.weight
-            resnet["conv1"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"decoder.mid_block.resnets.{i}.conv1.weight"],
-                    "bias": diffusers_weights[f"decoder.mid_block.resnets.{i}.conv1.bias"],
-                }
-            }
-            # conv2.weight -> conv2.conv3d.weight
-            resnet["conv2"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"decoder.mid_block.resnets.{i}.conv2.weight"],
-                    "bias": diffusers_weights[f"decoder.mid_block.resnets.{i}.conv2.bias"],
-                }
-            }
-            # norm gammas -> 1D weights
-            g1 = diffusers_weights[f"decoder.mid_block.resnets.{i}.norm1.gamma"]
-            if len(g1.shape) > 1:
-                g1 = mx.reshape(g1, (g1.shape[0],))
-            g2 = diffusers_weights[f"decoder.mid_block.resnets.{i}.norm2.gamma"]
-            if len(g2.shape) > 1:
-                g2 = mx.reshape(g2, (g2.shape[0],))
-            resnet["norm1"] = {"weight": g1}
-            resnet["norm2"] = {"weight": g2}
-
-        # mid_block.attentions (list of 1 attention)
-        weights["decoder"]["mid_block"]["attentions"] = [{}]
-        attn = weights["decoder"]["mid_block"]["attentions"][0]
-        g = diffusers_weights["decoder.mid_block.attentions.0.norm.gamma"]
-        if len(g.shape) > 1:
-            g = mx.reshape(g, (g.shape[0],))
-        attn["norm"] = {"weight": g}
-        # Note: to_qkv and proj are Conv2d, not Conv3d - no conv3d wrapper
-        attn["to_qkv"] = {
-            "weight": diffusers_weights["decoder.mid_block.attentions.0.to_qkv.weight"],
-            "bias": diffusers_weights["decoder.mid_block.attentions.0.to_qkv.bias"],
-        }
-        attn["proj"] = {
-            "weight": diffusers_weights["decoder.mid_block.attentions.0.proj.weight"],
-            "bias": diffusers_weights["decoder.mid_block.attentions.0.proj.bias"],
-        }
-
-        # 3. Up blocks (manual structure building)
-        for block_idx in range(4):
-            up_block_key = f"up_block{block_idx}"
-            weights["decoder"][up_block_key] = {}
-
-            # resnets (list of 3 resnets)
-            weights["decoder"][up_block_key]["resnets"] = [{}, {}, {}]
-            for res_idx in range(3):
-                resnet = weights["decoder"][up_block_key]["resnets"][res_idx]
-
-                # conv1.weight -> conv1.conv3d.weight
-                resnet["conv1"] = {
-                    "conv3d": {
-                        "weight": diffusers_weights[f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.conv1.weight"],
-                        "bias": diffusers_weights[f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.conv1.bias"],
-                    }
-                }
-                # conv2.weight -> conv2.conv3d.weight
-                resnet["conv2"] = {
-                    "conv3d": {
-                        "weight": diffusers_weights[f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.conv2.weight"],
-                        "bias": diffusers_weights[f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.conv2.bias"],
-                    }
-                }
-                # norm gammas -> 1D weights
-                g1 = diffusers_weights[f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.norm1.gamma"]
-                if len(g1.shape) > 1:
-                    g1 = mx.reshape(g1, (g1.shape[0],))
-                g2 = diffusers_weights[f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.norm2.gamma"]
-                if len(g2.shape) > 1:
-                    g2 = mx.reshape(g2, (g2.shape[0],))
-                resnet["norm1"] = {"weight": g1}
-                resnet["norm2"] = {"weight": g2}
-
-                # Handle optional conv_shortcut -> skip_conv (only exists for some resnets)
-                shortcut_key = f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.conv_shortcut.weight"
-                if shortcut_key in diffusers_weights:
-                    resnet["skip_conv"] = {
-                        "conv3d": {
-                            "weight": diffusers_weights[shortcut_key],
-                            "bias": diffusers_weights[
-                                f"decoder.up_blocks.{block_idx}.resnets.{res_idx}.conv_shortcut.bias"
-                            ],
-                        }
-                    }
-
-            # upsamplers (only for blocks 0, 1, 2)
-            if block_idx <= 2:
-                weights["decoder"][up_block_key]["upsamplers"] = [{}]
-                upsampler = weights["decoder"][up_block_key]["upsamplers"][0]
-
-                # resample.1.weight -> resample_conv.weight (Conv2d, no conv3d wrapper)
-                upsampler["resample_conv"] = {
-                    "weight": diffusers_weights[f"decoder.up_blocks.{block_idx}.upsamplers.0.resample.1.weight"],
-                    "bias": diffusers_weights[f"decoder.up_blocks.{block_idx}.upsamplers.0.resample.1.bias"],
-                }
-
-                # time_conv (only for blocks 0, 1)
-                if block_idx <= 1:
-                    upsampler["time_conv"] = {
-                        "conv3d": {
-                            "weight": diffusers_weights[f"decoder.up_blocks.{block_idx}.upsamplers.0.time_conv.weight"],
-                            "bias": diffusers_weights[f"decoder.up_blocks.{block_idx}.upsamplers.0.time_conv.bias"],
-                        }
-                    }
-
-        # 4. Encoder mappings (mirror structure of model parameter names)
-        weights["encoder"] = {}
-
-        # conv_in: encoder.conv_in.weight -> encoder.conv_in.conv3d.weight
-        weights["encoder"]["conv_in"] = {
-            "conv3d": {
-                "weight": diffusers_weights["encoder.conv_in.weight"],
-                "bias": diffusers_weights["encoder.conv_in.bias"],
-            }
-        }
-
-        # conv_out: encoder.conv_out.weight -> encoder.conv_out.conv3d.weight
-        weights["encoder"]["conv_out"] = {
-            "conv3d": {
-                "weight": diffusers_weights["encoder.conv_out.weight"],
-                "bias": diffusers_weights["encoder.conv_out.bias"],
-            }
-        }
-
-        # norm_out: encoder.norm_out.gamma -> encoder.norm_out.weight (flatten to 1D)
-        enc_gamma = diffusers_weights["encoder.norm_out.gamma"]
-        if len(enc_gamma.shape) > 1:
-            enc_gamma = mx.reshape(enc_gamma, (enc_gamma.shape[0],))
-        weights["encoder"]["norm_out"] = {"weight": enc_gamma}
-
-        # mid_block
-        weights["encoder"]["mid_block"] = {}
-        # mid_block.attentions (list of 1)
-        weights["encoder"]["mid_block"]["attentions"] = [{}]
-        enc_attn = weights["encoder"]["mid_block"]["attentions"][0]
-        g = diffusers_weights["encoder.mid_block.attentions.0.norm.gamma"]
-        if len(g.shape) > 1:
-            g = mx.reshape(g, (g.shape[0],))
-        enc_attn["norm"] = {"weight": g}
-        enc_attn["to_qkv"] = {
-            "weight": diffusers_weights["encoder.mid_block.attentions.0.to_qkv.weight"],
-            "bias": diffusers_weights["encoder.mid_block.attentions.0.to_qkv.bias"],
-        }
-        enc_attn["proj"] = {
-            "weight": diffusers_weights["encoder.mid_block.attentions.0.proj.weight"],
-            "bias": diffusers_weights["encoder.mid_block.attentions.0.proj.bias"],
-        }
-        # mid_block.resnets (list of 2)
-        weights["encoder"]["mid_block"]["resnets"] = [{}, {}]
-        for i in range(2):
-            res = weights["encoder"]["mid_block"]["resnets"][i]
-            g1 = diffusers_weights[f"encoder.mid_block.resnets.{i}.norm1.gamma"]
-            if len(g1.shape) > 1:
-                g1 = mx.reshape(g1, (g1.shape[0],))
-            g2 = diffusers_weights[f"encoder.mid_block.resnets.{i}.norm2.gamma"]
-            if len(g2.shape) > 1:
-                g2 = mx.reshape(g2, (g2.shape[0],))
-            res["norm1"] = {"weight": g1}
-            res["norm2"] = {"weight": g2}
-            res["conv1"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.mid_block.resnets.{i}.conv1.weight"],
-                    "bias": diffusers_weights[f"encoder.mid_block.resnets.{i}.conv1.bias"],
-                }
-            }
-            res["conv2"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.mid_block.resnets.{i}.conv2.weight"],
-                    "bias": diffusers_weights[f"encoder.mid_block.resnets.{i}.conv2.bias"],
-                }
-            }
-
-        # down_blocks - simplified mapping like decoder
-        # From safetensor keys, encoder has flattened indices 0-9 that need to be grouped into 4 stages
-        # Stage 0: indices 0,1 (2 resnets, no downsampler)
-        # Stage 1: indices 2,3,4 (downsampler at 2, resnets at 3,4)
-        # Stage 2: indices 5,6,7 (downsampler at 5, resnets at 6,7)
-        # Stage 3: indices 8,9 (downsampler at 8, resnet at 9)
-
-        weights["encoder"]["down_blocks"] = [{}, {}, {}, {}]
-
-        # Stage 0: 2 resnets + downsampler (2D)
-        weights["encoder"]["down_blocks"][0]["resnets"] = [{}, {}]
-        for res_idx in range(2):
-            flat_idx = res_idx  # 0, 1
-            resnet = weights["encoder"]["down_blocks"][0]["resnets"][res_idx]
-            g1 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm1.gamma"]
-            if len(g1.shape) > 1:
-                g1 = mx.reshape(g1, (g1.shape[0],))
-            g2 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm2.gamma"]
-            if len(g2.shape) > 1:
-                g2 = mx.reshape(g2, (g2.shape[0],))
-            resnet["norm1"] = {"weight": g1}
-            resnet["norm2"] = {"weight": g2}
-            resnet["conv1"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.bias"],
-                }
-            }
-            resnet["conv2"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.bias"],
-                }
-            }
-        # Downsampler appears after indices 0,1 at flattened index 2
-        weights["encoder"]["down_blocks"][0]["downsamplers"] = [{}]
-        d0 = weights["encoder"]["down_blocks"][0]["downsamplers"][0]
-        d0["resample_conv"] = {
-            "weight": diffusers_weights["encoder.down_blocks.2.resample.1.weight"],
-            "bias": diffusers_weights["encoder.down_blocks.2.resample.1.bias"],
-        }
-
-        # Stage 1: downsampler + 2 resnets
-        weights["encoder"]["down_blocks"][1]["resnets"] = [{}, {}]
-        weights["encoder"]["down_blocks"][1]["downsamplers"] = [{}]
-        # Downsampler at flattened index 5 (after stage0's 0,1,2 and stage1's 3,4)
-        d = weights["encoder"]["down_blocks"][1]["downsamplers"][0]
-        d["resample_conv"] = {
-            "weight": diffusers_weights["encoder.down_blocks.5.resample.1.weight"],
-            "bias": diffusers_weights["encoder.down_blocks.5.resample.1.bias"],
-        }
-        # Resnets at indices 3, 4
-        for res_idx in range(2):
-            flat_idx = 3 + res_idx  # 3, 4
-            resnet = weights["encoder"]["down_blocks"][1]["resnets"][res_idx]
-            g1 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm1.gamma"]
-            if len(g1.shape) > 1:
-                g1 = mx.reshape(g1, (g1.shape[0],))
-            g2 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm2.gamma"]
-            if len(g2.shape) > 1:
-                g2 = mx.reshape(g2, (g2.shape[0],))
-            resnet["norm1"] = {"weight": g1}
-            resnet["norm2"] = {"weight": g2}
-            resnet["conv1"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.bias"],
-                }
-            }
-            resnet["conv2"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.bias"],
-                }
-            }
-            # Skip conv only for first resnet (index 3)
-            if flat_idx == 3:
-                resnet["skip_conv"] = {
-                    "conv3d": {
-                        "weight": diffusers_weights["encoder.down_blocks.3.conv_shortcut.weight"],
-                        "bias": diffusers_weights["encoder.down_blocks.3.conv_shortcut.bias"],
-                    }
-                }
-
-        # Stage 2: downsampler + 2 resnets
-        weights["encoder"]["down_blocks"][2]["resnets"] = [{}, {}]
-        weights["encoder"]["down_blocks"][2]["downsamplers"] = [{}]
-        # Downsampler at flattened index 8
-        d = weights["encoder"]["down_blocks"][2]["downsamplers"][0]
-        d["resample_conv"] = {
-            "weight": diffusers_weights["encoder.down_blocks.8.resample.1.weight"],
-            "bias": diffusers_weights["encoder.down_blocks.8.resample.1.bias"],
-        }
-        d["time_conv"] = {
-            "conv3d": {
-                "weight": diffusers_weights["encoder.down_blocks.8.time_conv.weight"],
-                "bias": diffusers_weights["encoder.down_blocks.8.time_conv.bias"],
-            }
-        }
-        # Resnets at indices 6, 7
-        for res_idx in range(2):
-            flat_idx = 6 + res_idx  # 6, 7
-            resnet = weights["encoder"]["down_blocks"][2]["resnets"][res_idx]
-            g1 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm1.gamma"]
-            if len(g1.shape) > 1:
-                g1 = mx.reshape(g1, (g1.shape[0],))
-            g2 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm2.gamma"]
-            if len(g2.shape) > 1:
-                g2 = mx.reshape(g2, (g2.shape[0],))
-            resnet["norm1"] = {"weight": g1}
-            resnet["norm2"] = {"weight": g2}
-            resnet["conv1"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.bias"],
-                }
-            }
-            resnet["conv2"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.bias"],
-                }
-            }
-            # Skip conv only for first resnet (index 6)
-            if flat_idx == 6:
-                resnet["skip_conv"] = {
-                    "conv3d": {
-                        "weight": diffusers_weights["encoder.down_blocks.6.conv_shortcut.weight"],
-                        "bias": diffusers_weights["encoder.down_blocks.6.conv_shortcut.bias"],
-                    }
-                }
-
-        # Stage 3: no downsampler + 2 resnets (like Diffusers)
-        weights["encoder"]["down_blocks"][3]["resnets"] = [{}, {}]
-        # No downsampler in final stage (mirrors diffusers)
-        # Resnets at indices 9, 10
-        for res_idx in range(2):
-            flat_idx = 9 + res_idx  # 9, 10
-            resnet = weights["encoder"]["down_blocks"][3]["resnets"][res_idx]
-            g1 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm1.gamma"]
-            if len(g1.shape) > 1:
-                g1 = mx.reshape(g1, (g1.shape[0],))
-            g2 = diffusers_weights[f"encoder.down_blocks.{flat_idx}.norm2.gamma"]
-            if len(g2.shape) > 1:
-                g2 = mx.reshape(g2, (g2.shape[0],))
-            resnet["norm1"] = {"weight": g1}
-            resnet["norm2"] = {"weight": g2}
-            resnet["conv1"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv1.bias"],
-                }
-            }
-            resnet["conv2"] = {
-                "conv3d": {
-                    "weight": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.weight"],
-                    "bias": diffusers_weights[f"encoder.down_blocks.{flat_idx}.conv2.bias"],
-                }
-            }
-            # No skip_conv for stage 3 (channels don't change: 384->384)
-
-        # 5. Quant conv
-        weights["quant_conv"] = {
-            "conv3d": {
-                "weight": diffusers_weights["quant_conv.weight"],
-                "bias": diffusers_weights["quant_conv.bias"],
-            }
-        }
-
-        return weights
-
-    @staticmethod
-    def _manual_transformer_mapping(diffusers_weights: dict) -> dict:
-        weights = {}
-
-        # 1. Top-level mappings (exact MLX parameter names)
-        weights["img_in"] = {"weight": diffusers_weights["img_in.weight"], "bias": diffusers_weights["img_in.bias"]}
-        weights["txt_norm"] = {"weight": diffusers_weights["txt_norm.weight"]}
-        weights["txt_in"] = {"weight": diffusers_weights["txt_in.weight"], "bias": diffusers_weights["txt_in.bias"]}
-
-        # 2. Time text embedder (exact MLX structure)
-        weights["time_text_embed"] = {
-            "timestep_embedder": {
-                "linear_1": {
-                    "weight": diffusers_weights["time_text_embed.timestep_embedder.linear_1.weight"],
-                    "bias": diffusers_weights["time_text_embed.timestep_embedder.linear_1.bias"],
-                },
-                "linear_2": {
-                    "weight": diffusers_weights["time_text_embed.timestep_embedder.linear_2.weight"],
-                    "bias": diffusers_weights["time_text_embed.timestep_embedder.linear_2.bias"],
-                },
-            }
-        }
-
-        # 3. Output head (exact MLX structure)
-        weights["norm_out"] = {
-            "linear": {
-                "weight": diffusers_weights["norm_out.linear.weight"],
-                "bias": diffusers_weights["norm_out.linear.bias"],
-            }
-        }
-        weights["proj_out"] = {
-            "weight": diffusers_weights["proj_out.weight"],
-            "bias": diffusers_weights["proj_out.bias"],
-        }
-
-        # 4. Transformer blocks (exact MLX parameter names - no applier needed!)
-        transformer_blocks = []
-        for block_idx in range(60):  # 60 blocks based on debug output
-            block = {}
-
-            # Stage 1: Normalization + modulation modules (QwenLayerNorm)
-            block["img_norm1"] = {
-                "mod_linear": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.img_mod.1.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.img_mod.1.bias"],
-                },
-                "norm1": {},  # LayerNorm with affine=False has no weights
-                "norm2": {},  # LayerNorm with affine=False has no weights
-            }
-            block["txt_norm1"] = {
-                "mod_linear": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.txt_mod.1.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.txt_mod.1.bias"],
-                },
-                "norm1": {},  # LayerNorm with affine=False has no weights
-                "norm2": {},  # LayerNorm with affine=False has no weights
-            }
-
-            # Stage 2: Separate normalization modules (Flux-style)
-            block["img_norm2"] = {}  # LayerNorm with affine=False has no weights
-            block["txt_norm2"] = {}  # LayerNorm with affine=False has no weights
-
-            # Attention module (nested under "attn" like Flux)
-            block["attn"] = {
-                "to_q": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_q.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_q.bias"],
-                },
-                "to_k": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_k.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_k.bias"],
-                },
-                "to_v": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_v.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_v.bias"],
-                },
-                "add_q_proj": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.add_q_proj.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.add_q_proj.bias"],
-                },
-                "add_k_proj": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.add_k_proj.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.add_k_proj.bias"],
-                },
-                "add_v_proj": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.add_v_proj.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.add_v_proj.bias"],
-                },
-                "norm_q": {"weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.norm_q.weight"]},
-                "norm_k": {"weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.norm_k.weight"]},
-                "norm_added_q": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.norm_added_q.weight"]
-                },
-                "norm_added_k": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.norm_added_k.weight"]
-                },
-                "attn_to_out": [  # type: ignore[dict-item]
-                    {
-                        "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_out.0.weight"],
-                        "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_out.0.bias"],
-                    }
-                ],
-                "to_add_out": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_add_out.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.attn.to_add_out.bias"],
-                },
-            }
-
-            # Feed Forward modules (nested under img_ff/txt_ff)
-            block["img_ff"] = {
-                "mlp_in": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.img_mlp.net.0.proj.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.img_mlp.net.0.proj.bias"],
-                },
-                "mlp_out": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.img_mlp.net.2.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.img_mlp.net.2.bias"],
-                },
-            }
-            block["txt_ff"] = {
-                "mlp_in": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.txt_mlp.net.0.proj.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.txt_mlp.net.0.proj.bias"],
-                },
-                "mlp_out": {
-                    "weight": diffusers_weights[f"transformer_blocks.{block_idx}.txt_mlp.net.2.weight"],
-                    "bias": diffusers_weights[f"transformer_blocks.{block_idx}.txt_mlp.net.2.bias"],
-                },
-            }
-
-            transformer_blocks.append(block)
-
-        weights["transformer_blocks"] = transformer_blocks
-        return weights
-
     @staticmethod
     def _load_safetensors_shards(path: Path, loading_mode: str = "multi_glob") -> dict[str, mx.array]:
         all_weights = {}
@@ -696,9 +140,6 @@ class QwenWeightHandler:
                     file_weights = mlx_load_file(str(file_path))
                 except Exception:  # noqa: BLE001
                     # If MLX can't load directly, try with torch and convert
-                    import torch
-                    from safetensors.torch import load_file as torch_load_file
-
                     torch_weights = torch_load_file(str(file_path))
                     file_weights = {}
                     for name, tensor in torch_weights.items():
@@ -724,79 +165,26 @@ class QwenWeightHandler:
         return all_weights
 
     @staticmethod
-    def _manual_text_encoder_mapping(hf_weights: dict[str, mx.array]) -> dict:
-        weights = {}
-        converted_count = 0
-        skipped_count = 0
+    def _load_saved_model_weights(
+        weights: dict[str, mx.array] | None,
+        path: Path | None,
+        quantization_level: int | None,
+        mflux_version: str | None,
+    ) -> tuple[dict, int | None, str | None]:
+        # If weights already loaded, use them; otherwise load from path
+        if weights is None:
+            # For saved models, always use multi_glob (no index.json needed)
+            weights = QwenWeightHandler._load_safetensors_shards(path, loading_mode="multi_glob")
 
-        # Skip LM head and vision encoder weights - we only need the text encoder
-        filtered_weights = {}
-        for hf_name, weight in hf_weights.items():
-            if hf_name.startswith("lm_head") or hf_name.startswith("visual."):
-                skipped_count += 1
-                continue
-            filtered_weights[hf_name] = weight
+        return tree_unflatten(list(weights.items())), quantization_level, mflux_version
 
-        # 1. Top-level embeddings (exact MLX parameter names)
-        weights["encoder"] = {}
-        weights["encoder"]["embed_tokens"] = {"weight": filtered_weights["model.embed_tokens.weight"]}
-        converted_count += 1
-
-        # 2. Final norm (exact MLX parameter names)
-        weights["encoder"]["norm"] = {"weight": filtered_weights["model.norm.weight"]}
-        converted_count += 1
-
-        # 3. Encoder layers (exact MLX parameter names - 28 layers)
-        layers = []
-        for layer_idx in range(28):  # 28 text encoder layers
-            layer = {}
-
-            # Layer norms
-            layer["input_layernorm"] = {"weight": filtered_weights[f"model.layers.{layer_idx}.input_layernorm.weight"]}
-            layer["post_attention_layernorm"] = {
-                "weight": filtered_weights[f"model.layers.{layer_idx}.post_attention_layernorm.weight"]
-            }
-            converted_count += 2
-
-            # Self attention (exact MLX parameter names)
-            layer["self_attn"] = {
-                "q_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.self_attn.q_proj.weight"],
-                    "bias": filtered_weights[f"model.layers.{layer_idx}.self_attn.q_proj.bias"],
-                },
-                "k_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.self_attn.k_proj.weight"],
-                    "bias": filtered_weights[f"model.layers.{layer_idx}.self_attn.k_proj.bias"],
-                },
-                "v_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.self_attn.v_proj.weight"],
-                    "bias": filtered_weights[f"model.layers.{layer_idx}.self_attn.v_proj.bias"],
-                },
-                "o_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.self_attn.o_proj.weight"]
-                    # Note: o_proj has no bias in MLX structure
-                },
-            }
-            converted_count += 7
-
-            # MLP (exact MLX parameter names)
-            layer["mlp"] = {
-                "gate_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.mlp.gate_proj.weight"]
-                    # Note: gate_proj has no bias in MLX structure
-                },
-                "up_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.mlp.up_proj.weight"]
-                    # Note: up_proj has no bias in MLX structure
-                },
-                "down_proj": {
-                    "weight": filtered_weights[f"model.layers.{layer_idx}.mlp.down_proj.weight"]
-                    # Note: down_proj has no bias in MLX structure
-                },
-            }
-            converted_count += 3
-
-            layers.append(layer)
-
-        weights["encoder"]["layers"] = layers
-        return weights
+    @staticmethod
+    def _detect_metadata(path: Path) -> tuple[int | None, str | None]:
+        file_glob = sorted(path.glob("*.safetensors"))
+        if file_glob:
+            data = mx.load(str(file_glob[0]), return_metadata=True)
+            if len(data) > 1:
+                quantization_level = data[1].get("quantization_level")
+                mflux_version = data[1].get("mflux_version")
+                return quantization_level, mflux_version
+        return None, None

src/mflux/models/qwen/weights/qwen_weight_mapping.py

@@ -0,0 +1,919 @@
+from typing import List
+
+import mlx.core as mx
+
+from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget
+
+
+def reshape_gamma_to_1d(tensor: mx.array) -> mx.array:
+    if len(tensor.shape) > 1:
+        return mx.reshape(tensor, (tensor.shape[0],))
+    return tensor
+
+
+def transpose_patch_embed(tensor: mx.array) -> mx.array:
+    if len(tensor.shape) == 5:
+        return tensor.transpose(0, 2, 3, 4, 1)
+    return tensor
+
+
+def transpose_conv3d_weight(tensor: mx.array) -> mx.array:
+    if len(tensor.shape) == 5:
+        return tensor.transpose(0, 2, 3, 4, 1)
+    return tensor
+
+
+def transpose_conv2d_weight(tensor: mx.array) -> mx.array:
+    if len(tensor.shape) == 4:
+        return tensor.transpose(0, 2, 3, 1)
+    return tensor
+
+
+class QwenWeightMapping(WeightMapping):
+    @staticmethod
+    def get_transformer_mapping() -> List[WeightTarget]:
+        return [
+            # Top-level mappings
+            WeightTarget(
+                mlx_path="img_in.weight",
+                hf_patterns=["img_in.weight"],
+            ),
+            WeightTarget(
+                mlx_path="img_in.bias",
+                hf_patterns=["img_in.bias"],
+            ),
+            WeightTarget(
+                mlx_path="txt_norm.weight",
+                hf_patterns=["txt_norm.weight"],
+            ),
+            WeightTarget(
+                mlx_path="txt_in.weight",
+                hf_patterns=["txt_in.weight"],
+            ),
+            WeightTarget(
+                mlx_path="txt_in.bias",
+                hf_patterns=["txt_in.bias"],
+            ),
+            # Time text embedder
+            WeightTarget(
+                mlx_path="time_text_embed.timestep_embedder.linear_1.weight",
+                hf_patterns=["time_text_embed.timestep_embedder.linear_1.weight"],
+            ),
+            WeightTarget(
+                mlx_path="time_text_embed.timestep_embedder.linear_1.bias",
+                hf_patterns=["time_text_embed.timestep_embedder.linear_1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="time_text_embed.timestep_embedder.linear_2.weight",
+                hf_patterns=["time_text_embed.timestep_embedder.linear_2.weight"],
+            ),
+            WeightTarget(
+                mlx_path="time_text_embed.timestep_embedder.linear_2.bias",
+                hf_patterns=["time_text_embed.timestep_embedder.linear_2.bias"],
+            ),
+            # Output head
+            WeightTarget(
+                mlx_path="norm_out.linear.weight",
+                hf_patterns=["norm_out.linear.weight"],
+            ),
+            WeightTarget(
+                mlx_path="norm_out.linear.bias",
+                hf_patterns=["norm_out.linear.bias"],
+            ),
+            WeightTarget(
+                mlx_path="proj_out.weight",
+                hf_patterns=["proj_out.weight"],
+            ),
+            WeightTarget(
+                mlx_path="proj_out.bias",
+                hf_patterns=["proj_out.bias"],
+            ),
+            # Transformer blocks - Attention
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_q.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.to_q.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_q.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.to_q.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_k.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.to_k.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_k.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.to_k.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_v.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.to_v.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_v.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.to_v.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.add_q_proj.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.add_q_proj.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.add_k_proj.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.add_k_proj.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.add_v_proj.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.add_v_proj.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.norm_q.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.norm_q.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.norm_k.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.norm_k.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.norm_added_q.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.norm_added_q.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.norm_added_k.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.norm_added_k.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.attn_to_out.0.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.to_out.0.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.attn_to_out.0.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.to_out.0.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_add_out.weight",
+                hf_patterns=["transformer_blocks.{block}.attn.to_add_out.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.attn.to_add_out.bias",
+                hf_patterns=["transformer_blocks.{block}.attn.to_add_out.bias"],
+            ),
+            # Transformer blocks - Modulation
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.img_mod_linear.weight",
+                hf_patterns=["transformer_blocks.{block}.img_mod.1.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.img_mod_linear.bias",
+                hf_patterns=["transformer_blocks.{block}.img_mod.1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.txt_mod_linear.weight",
+                hf_patterns=["transformer_blocks.{block}.txt_mod.1.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.txt_mod_linear.bias",
+                hf_patterns=["transformer_blocks.{block}.txt_mod.1.bias"],
+            ),
+            # Transformer blocks - Feed Forward
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.img_ff.mlp_in.weight",
+                hf_patterns=["transformer_blocks.{block}.img_mlp.net.0.proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.img_ff.mlp_in.bias",
+                hf_patterns=["transformer_blocks.{block}.img_mlp.net.0.proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.img_ff.mlp_out.weight",
+                hf_patterns=["transformer_blocks.{block}.img_mlp.net.2.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.img_ff.mlp_out.bias",
+                hf_patterns=["transformer_blocks.{block}.img_mlp.net.2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.txt_ff.mlp_in.weight",
+                hf_patterns=["transformer_blocks.{block}.txt_mlp.net.0.proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.txt_ff.mlp_in.bias",
+                hf_patterns=["transformer_blocks.{block}.txt_mlp.net.0.proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.txt_ff.mlp_out.weight",
+                hf_patterns=["transformer_blocks.{block}.txt_mlp.net.2.weight"],
+            ),
+            WeightTarget(
+                mlx_path="transformer_blocks.{block}.txt_ff.mlp_out.bias",
+                hf_patterns=["transformer_blocks.{block}.txt_mlp.net.2.bias"],
+            ),
+        ]
+
+    @staticmethod
+    def get_vae_mapping() -> List[WeightTarget]:
+        return [
+            # ========== Decoder ==========
+            # Decoder conv_in
+            WeightTarget(
+                mlx_path="decoder.conv_in.conv3d.weight",
+                hf_patterns=["decoder.conv_in.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.conv_in.conv3d.bias",
+                hf_patterns=["decoder.conv_in.bias"],
+            ),
+            # Decoder conv_out
+            WeightTarget(
+                mlx_path="decoder.conv_out.conv3d.weight",
+                hf_patterns=["decoder.conv_out.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.conv_out.conv3d.bias",
+                hf_patterns=["decoder.conv_out.bias"],
+            ),
+            # Decoder norm_out (gamma -> weight with reshape)
+            WeightTarget(
+                mlx_path="decoder.norm_out.weight",
+                hf_patterns=["decoder.norm_out.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            # Post quant conv
+            WeightTarget(
+                mlx_path="post_quant_conv.conv3d.weight",
+                hf_patterns=["post_quant_conv.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="post_quant_conv.conv3d.bias",
+                hf_patterns=["post_quant_conv.bias"],
+            ),
+            # Decoder mid_block resnets
+            WeightTarget(
+                mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.weight",
+                hf_patterns=["decoder.mid_block.resnets.{i}.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.bias",
+                hf_patterns=["decoder.mid_block.resnets.{i}.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.weight",
+                hf_patterns=["decoder.mid_block.resnets.{i}.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.bias",
+                hf_patterns=["decoder.mid_block.resnets.{i}.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.resnets.{i}.norm1.weight",
+                hf_patterns=["decoder.mid_block.resnets.{i}.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.resnets.{i}.norm2.weight",
+                hf_patterns=["decoder.mid_block.resnets.{i}.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            # Decoder mid_block attention
+            WeightTarget(
+                mlx_path="decoder.mid_block.attentions.0.norm.weight",
+                hf_patterns=["decoder.mid_block.attentions.0.norm.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.attentions.0.to_qkv.weight",
+                hf_patterns=["decoder.mid_block.attentions.0.to_qkv.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.attentions.0.to_qkv.bias",
+                hf_patterns=["decoder.mid_block.attentions.0.to_qkv.bias"],
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.attentions.0.proj.weight",
+                hf_patterns=["decoder.mid_block.attentions.0.proj.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.mid_block.attentions.0.proj.bias",
+                hf_patterns=["decoder.mid_block.attentions.0.proj.bias"],
+            ),
+            # Decoder up_blocks resnets
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.conv1.conv3d.weight",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.conv1.conv3d.bias",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.conv2.conv3d.weight",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.conv2.conv3d.bias",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.norm1.weight",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.norm2.weight",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            # Decoder up_blocks optional conv_shortcut -> skip_conv
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.weight",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
+                transform=transpose_conv3d_weight,
+                required=False,  # Optional weight
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.resnets.{res}.skip_conv.conv3d.bias",
+                hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
+                required=False,  # Optional weight
+            ),
+            # Decoder up_blocks upsamplers (blocks 0, 1, 2)
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.upsamplers.0.resample_conv.weight",
+                hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.resample.1.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.upsamplers.0.resample_conv.bias",
+                hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.resample.1.bias"],
+            ),
+            # Decoder up_blocks time_conv (blocks 0, 1)
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.weight",
+                hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.time_conv.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="decoder.up_block{block}.upsamplers.0.time_conv.conv3d.bias",
+                hf_patterns=["decoder.up_blocks.{block}.upsamplers.0.time_conv.bias"],
+            ),
+            # ========== Encoder ==========
+            # Encoder conv_in
+            WeightTarget(
+                mlx_path="encoder.conv_in.conv3d.weight",
+                hf_patterns=["encoder.conv_in.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.conv_in.conv3d.bias",
+                hf_patterns=["encoder.conv_in.bias"],
+            ),
+            # Encoder conv_out
+            WeightTarget(
+                mlx_path="encoder.conv_out.conv3d.weight",
+                hf_patterns=["encoder.conv_out.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.conv_out.conv3d.bias",
+                hf_patterns=["encoder.conv_out.bias"],
+            ),
+            # Encoder norm_out
+            WeightTarget(
+                mlx_path="encoder.norm_out.weight",
+                hf_patterns=["encoder.norm_out.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            # Encoder mid_block attention
+            WeightTarget(
+                mlx_path="encoder.mid_block.attentions.0.norm.weight",
+                hf_patterns=["encoder.mid_block.attentions.0.norm.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.attentions.0.to_qkv.weight",
+                hf_patterns=["encoder.mid_block.attentions.0.to_qkv.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.attentions.0.to_qkv.bias",
+                hf_patterns=["encoder.mid_block.attentions.0.to_qkv.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.attentions.0.proj.weight",
+                hf_patterns=["encoder.mid_block.attentions.0.proj.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.attentions.0.proj.bias",
+                hf_patterns=["encoder.mid_block.attentions.0.proj.bias"],
+            ),
+            # Encoder mid_block resnets
+            WeightTarget(
+                mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.weight",
+                hf_patterns=["encoder.mid_block.resnets.{i}.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.bias",
+                hf_patterns=["encoder.mid_block.resnets.{i}.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.weight",
+                hf_patterns=["encoder.mid_block.resnets.{i}.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.bias",
+                hf_patterns=["encoder.mid_block.resnets.{i}.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.resnets.{i}.norm1.weight",
+                hf_patterns=["encoder.mid_block.resnets.{i}.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.mid_block.resnets.{i}.norm2.weight",
+                hf_patterns=["encoder.mid_block.resnets.{i}.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            # Encoder down_blocks - Stage 0 (flat_idx 0,1 -> resnets[0,1]; flat_idx 2 -> downsampler)
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.0.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.0.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.0.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.0.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.0.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.0.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.0.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.0.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.0.norm1.weight",
+                hf_patterns=["encoder.down_blocks.0.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.0.norm2.weight",
+                hf_patterns=["encoder.down_blocks.0.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.1.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.1.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.1.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.1.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.1.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.1.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.1.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.1.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.1.norm1.weight",
+                hf_patterns=["encoder.down_blocks.1.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.resnets.1.norm2.weight",
+                hf_patterns=["encoder.down_blocks.1.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.downsamplers.0.resample_conv.weight",
+                hf_patterns=["encoder.down_blocks.2.resample.1.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.0.downsamplers.0.resample_conv.bias",
+                hf_patterns=["encoder.down_blocks.2.resample.1.bias"],
+            ),
+            # Encoder down_blocks - Stage 1 (flat_idx 3,4 -> resnets[0,1]; flat_idx 5 -> downsampler)
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.3.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.3.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.3.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.3.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.norm1.weight",
+                hf_patterns=["encoder.down_blocks.3.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.norm2.weight",
+                hf_patterns=["encoder.down_blocks.3.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.3.conv_shortcut.weight"],
+                transform=transpose_conv3d_weight,
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.0.skip_conv.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.3.conv_shortcut.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.1.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.4.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.1.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.4.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.1.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.4.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.1.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.4.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.1.norm1.weight",
+                hf_patterns=["encoder.down_blocks.4.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.resnets.1.norm2.weight",
+                hf_patterns=["encoder.down_blocks.4.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.downsamplers.0.resample_conv.weight",
+                hf_patterns=["encoder.down_blocks.5.resample.1.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.1.downsamplers.0.resample_conv.bias",
+                hf_patterns=["encoder.down_blocks.5.resample.1.bias"],
+            ),
+            # Encoder down_blocks - Stage 2 (flat_idx 6,7 -> resnets[0,1]; flat_idx 8 -> downsampler)
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.6.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.6.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.6.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.6.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.norm1.weight",
+                hf_patterns=["encoder.down_blocks.6.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.norm2.weight",
+                hf_patterns=["encoder.down_blocks.6.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.6.conv_shortcut.weight"],
+                transform=transpose_conv3d_weight,
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.0.skip_conv.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.6.conv_shortcut.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.1.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.7.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.1.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.7.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.1.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.7.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.1.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.7.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.1.norm1.weight",
+                hf_patterns=["encoder.down_blocks.7.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.resnets.1.norm2.weight",
+                hf_patterns=["encoder.down_blocks.7.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.downsamplers.0.resample_conv.weight",
+                hf_patterns=["encoder.down_blocks.8.resample.1.weight"],
+                transform=transpose_conv2d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.downsamplers.0.resample_conv.bias",
+                hf_patterns=["encoder.down_blocks.8.resample.1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.8.time_conv.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.2.downsamplers.0.time_conv.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.8.time_conv.bias"],
+            ),
+            # Encoder down_blocks - Stage 3 (flat_idx 9,10 -> resnets[0,1])
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.0.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.9.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.0.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.9.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.0.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.9.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.0.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.9.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.0.norm1.weight",
+                hf_patterns=["encoder.down_blocks.9.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.0.norm2.weight",
+                hf_patterns=["encoder.down_blocks.9.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.1.conv1.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.10.conv1.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.1.conv1.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.10.conv1.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.1.conv2.conv3d.weight",
+                hf_patterns=["encoder.down_blocks.10.conv2.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.1.conv2.conv3d.bias",
+                hf_patterns=["encoder.down_blocks.10.conv2.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.1.norm1.weight",
+                hf_patterns=["encoder.down_blocks.10.norm1.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            WeightTarget(
+                mlx_path="encoder.down_blocks.3.resnets.1.norm2.weight",
+                hf_patterns=["encoder.down_blocks.10.norm2.gamma"],
+                transform=reshape_gamma_to_1d,
+            ),
+            # Quant conv
+            WeightTarget(
+                mlx_path="quant_conv.conv3d.weight",
+                hf_patterns=["quant_conv.weight"],
+                transform=transpose_conv3d_weight,
+            ),
+            WeightTarget(
+                mlx_path="quant_conv.conv3d.bias",
+                hf_patterns=["quant_conv.bias"],
+            ),
+        ]
+
+    @staticmethod
+    def get_text_encoder_mapping() -> List[WeightTarget]:
+        return [
+            # Top-level embeddings
+            WeightTarget(
+                mlx_path="encoder.embed_tokens.weight",
+                hf_patterns=["model.embed_tokens.weight"],
+            ),
+            # Final norm
+            WeightTarget(
+                mlx_path="encoder.norm.weight",
+                hf_patterns=["model.norm.weight"],
+            ),
+            # Encoder layers (28 layers)
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.input_layernorm.weight",
+                hf_patterns=["model.layers.{layer}.input_layernorm.weight"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.post_attention_layernorm.weight",
+                hf_patterns=["model.layers.{layer}.post_attention_layernorm.weight"],
+            ),
+            # Self attention
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.q_proj.weight",
+                hf_patterns=["model.layers.{layer}.self_attn.q_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.q_proj.bias",
+                hf_patterns=["model.layers.{layer}.self_attn.q_proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.k_proj.weight",
+                hf_patterns=["model.layers.{layer}.self_attn.k_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.k_proj.bias",
+                hf_patterns=["model.layers.{layer}.self_attn.k_proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.v_proj.weight",
+                hf_patterns=["model.layers.{layer}.self_attn.v_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.v_proj.bias",
+                hf_patterns=["model.layers.{layer}.self_attn.v_proj.bias"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.self_attn.o_proj.weight",
+                hf_patterns=["model.layers.{layer}.self_attn.o_proj.weight"],
+            ),
+            # MLP
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.mlp.gate_proj.weight",
+                hf_patterns=["model.layers.{layer}.mlp.gate_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.mlp.up_proj.weight",
+                hf_patterns=["model.layers.{layer}.mlp.up_proj.weight"],
+            ),
+            WeightTarget(
+                mlx_path="encoder.layers.{layer}.mlp.down_proj.weight",
+                hf_patterns=["model.layers.{layer}.mlp.down_proj.weight"],
+            ),
+            # Visual weights (optional, only present in Edit models)
+            # Patch embedding (with transpose transform)
+            WeightTarget(
+                mlx_path="encoder.visual.patch_embed.proj.weight",
+                hf_patterns=["visual.patch_embed.proj.weight"],
+                transform=transpose_patch_embed,
+                required=False,
+            ),
+            # Vision transformer blocks (32 blocks) - use {block} placeholder
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.attn.qkv.weight",
+                hf_patterns=["visual.blocks.{block}.attn.qkv.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.attn.qkv.bias",
+                hf_patterns=["visual.blocks.{block}.attn.qkv.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.attn.proj.weight",
+                hf_patterns=["visual.blocks.{block}.attn.proj.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.attn.proj.bias",
+                hf_patterns=["visual.blocks.{block}.attn.proj.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.mlp.gate_proj.weight",
+                hf_patterns=["visual.blocks.{block}.mlp.gate_proj.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.mlp.gate_proj.bias",
+                hf_patterns=["visual.blocks.{block}.mlp.gate_proj.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.mlp.up_proj.weight",
+                hf_patterns=["visual.blocks.{block}.mlp.up_proj.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.mlp.up_proj.bias",
+                hf_patterns=["visual.blocks.{block}.mlp.up_proj.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.mlp.down_proj.weight",
+                hf_patterns=["visual.blocks.{block}.mlp.down_proj.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.mlp.down_proj.bias",
+                hf_patterns=["visual.blocks.{block}.mlp.down_proj.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.norm1.weight",
+                hf_patterns=["visual.blocks.{block}.norm1.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.blocks.{block}.norm2.weight",
+                hf_patterns=["visual.blocks.{block}.norm2.weight"],
+                required=False,
+            ),
+            # Patch merger
+            WeightTarget(
+                mlx_path="encoder.visual.merger.ln_q.weight",
+                hf_patterns=["visual.merger.ln_q.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.merger.mlp_0.weight",
+                hf_patterns=["visual.merger.mlp.0.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.merger.mlp_0.bias",
+                hf_patterns=["visual.merger.mlp.0.bias"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.merger.mlp_1.weight",
+                hf_patterns=["visual.merger.mlp.2.weight"],
+                required=False,
+            ),
+            WeightTarget(
+                mlx_path="encoder.visual.merger.mlp_1.bias",
+                hf_patterns=["visual.merger.mlp.2.bias"],
+                required=False,
+            ),
+        ]
+
+    @staticmethod
+    def get_mapping() -> List[WeightTarget]:
+        return QwenWeightMapping.get_transformer_mapping()

src/mflux/models/qwen/weights/qwen_weight_util.py

@@ -1,9 +1,11 @@
 from typing import TYPE_CHECKING
 
+import mlx.core as mx
 import mlx.nn as nn
 from mlx.utils import tree_flatten  # noqa: F401
 
 from mflux.config.config import Config
+from mflux.models.qwen.model.qwen_text_encoder.qwen_vision_transformer import VisionTransformer
 from mflux.utils.quantization_util import QuantizationUtil
 
 if TYPE_CHECKING:
@@ -50,6 +52,25 @@ class QwenWeightUtil:
 
         raise Exception("Error setting weights")
 
+    @staticmethod
+    def _convert_weights_to_bf16(weights_dict: dict):
+        """
+        Recursively convert all weight tensors to BF16.
+        This matches PyTorch's behavior where the text encoder model is loaded in BF16.
+        """
+
+        def convert_recursive(obj):
+            if isinstance(obj, mx.array):
+                return obj.astype(mx.bfloat16)
+            elif isinstance(obj, dict):
+                return {k: convert_recursive(v) for k, v in obj.items()}
+            elif isinstance(obj, list):
+                return [convert_recursive(item) for item in obj]
+            else:
+                return obj
+
+        return convert_recursive(weights_dict)
+
     @staticmethod
     def _set_model_weights(
         weights: "QwenWeightHandler",
@@ -59,4 +80,28 @@ class QwenWeightUtil:
     ):
         vae.update(weights.vae, strict=False)
         transformer.update(weights.transformer, strict=False)
-        text_encoder.update(weights.qwen_text_encoder, strict=False)
+
+        # Check if visual weights are present and create visual transformer if needed
+        if text_encoder is not None:
+            has_visual_weights = (
+                "encoder" in weights.qwen_text_encoder and "visual" in weights.qwen_text_encoder["encoder"]
+            )
+            if has_visual_weights and text_encoder.encoder.visual is None:
+                text_encoder.encoder.visual = VisionTransformer(
+                    patch_size=14,
+                    temporal_patch_size=2,
+                    in_channels=3,
+                    embed_dim=1280,
+                    depth=32,
+                    num_heads=16,
+                    mlp_ratio=2.671875,
+                    hidden_size=text_encoder.encoder.hidden_size,
+                    spatial_merge_size=2,  # Match HF's spatial merging
+                )
+
+            # Convert all text encoder weights to BF16 to match PyTorch's behavior
+            # PyTorch loads the model in BF16 and performs all computations in BF16
+            # This ensures numerical consistency between MLX and PyTorch
+            weights.qwen_text_encoder = QwenWeightUtil._convert_weights_to_bf16(weights.qwen_text_encoder)
+
+            text_encoder.update(weights.qwen_text_encoder, strict=False)

src/mflux/schedulers/__init__.py

@@ -1,7 +1,9 @@
+from .flow_match_euler_discrete_scheduler import FlowMatchEulerDiscreteScheduler
 from .linear_scheduler import LinearScheduler
 
 __all__ = [
     "LinearScheduler",
+    "FlowMatchEulerDiscreteScheduler",
 ]
 
 
@@ -17,6 +19,8 @@ class InvalidSchedulerType(TypeError): ...
 SCHEDULER_REGISTRY = {
     "linear": LinearScheduler,
     "LinearScheduler": LinearScheduler,
+    "flow_match_euler_discrete": FlowMatchEulerDiscreteScheduler,
+    "FlowMatchEulerDiscreteScheduler": FlowMatchEulerDiscreteScheduler,
 }
 
 
@@ -56,7 +60,7 @@ def try_import_external_scheduler(scheduler_object_path: str):
     # Step 3: Validate that it's a class and inherits from BaseScheduler
     if not inspect.isclass(SchedulerClass):
         raise InvalidSchedulerType(
-            f"{scheduler_object_path!r} is not a class. " f"Schedulers must be classes inheriting from BaseScheduler."
+            f"{scheduler_object_path!r} is not a class. Schedulers must be classes inheriting from BaseScheduler."
         )
 
     if not issubclass(SchedulerClass, BaseScheduler):

src/mflux/schedulers/flow_match_euler_discrete_scheduler.py

@@ -0,0 +1,75 @@
+import math
+from typing import TYPE_CHECKING
+
+import mlx.core as mx
+
+if TYPE_CHECKING:
+    from mflux.config.runtime_config import RuntimeConfig
+
+from mflux.schedulers.base_scheduler import BaseScheduler
+
+
+class FlowMatchEulerDiscreteScheduler(BaseScheduler):
+    def __init__(self, runtime_config: "RuntimeConfig"):
+        self.runtime_config = runtime_config
+        self.model_config = runtime_config.model_config
+        self.num_train_timesteps = 1000
+        self.shift_terminal = 0.02
+        self.base_shift = 0.5
+        self.max_shift = 0.9
+        self.base_image_seq_len = 256
+        self.max_image_seq_len = 8192
+        self._sigmas, self._timesteps = self._compute_timesteps_and_sigmas()
+
+    @property
+    def sigmas(self) -> mx.array:
+        return self._sigmas
+
+    @property
+    def timesteps(self) -> mx.array:
+        return self._timesteps
+
+    def _compute_mu(self) -> float:
+        h_patches = self.runtime_config.height // 16
+        w_patches = self.runtime_config.width // 16
+        seq_len = h_patches * w_patches
+        m = (self.max_shift - self.base_shift) / (self.max_image_seq_len - self.base_image_seq_len)
+        b = self.base_shift - m * self.base_image_seq_len
+        mu = m * seq_len + b
+        return mu
+
+    @staticmethod
+    def _time_shift_exponential(mu: float, sigma_power: float, t: float) -> float:
+        return math.exp(mu) / (math.exp(mu) + ((1.0 / t - 1.0) ** sigma_power))
+
+    def _stretch_to_terminal(self, sigmas: list[float]) -> list[float]:
+        one_minus_sigmas = [1.0 - s for s in sigmas]
+        scale_factor = one_minus_sigmas[-1] / (1.0 - self.shift_terminal)
+        stretched = [1.0 - (oms / scale_factor) for oms in one_minus_sigmas]
+        return stretched
+
+    def _compute_timesteps_and_sigmas(self) -> tuple[mx.array, mx.array]:
+        num_steps = self.runtime_config.num_inference_steps
+        sigma_min = 1.0 / self.num_train_timesteps
+        sigma_max = 1.0
+        timesteps_linear = [
+            sigma_max * self.num_train_timesteps
+            - i * (sigma_max - sigma_min) * self.num_train_timesteps / (num_steps - 1)
+            for i in range(num_steps)
+        ]
+        sigmas_linear = [t / self.num_train_timesteps for t in timesteps_linear]
+        sigmas_shifted = [FlowMatchEulerDiscreteScheduler._time_shift_exponential(1.0, 1.0, s) for s in sigmas_linear]
+        sigmas_final = self._stretch_to_terminal(sigmas_shifted)
+        timesteps = [s * self.num_train_timesteps for s in sigmas_final]
+        sigmas_with_zero = sigmas_final + [0.0]
+        sigmas_arr = mx.array(sigmas_with_zero, dtype=mx.float32)
+        timesteps_arr = mx.array(timesteps, dtype=mx.float32)
+        return sigmas_arr, timesteps_arr
+
+    def step(self, model_output: mx.array, timestep: int, sample: mx.array, **kwargs) -> mx.array:
+        dt = self._sigmas[timestep + 1] - self._sigmas[timestep]
+        prev_sample = sample + dt * model_output
+        return prev_sample
+
+    def scale_model_input(self, latents: mx.array, t: int) -> mx.array:
+        return latents

src/mflux/schedulers/linear_scheduler.py

@@ -9,19 +9,19 @@ from mflux.schedulers.base_scheduler import BaseScheduler
 
 
 class LinearScheduler(BaseScheduler):
-    """
-    Linear scheduler - the default/classic scheduler used in mflux.
-    Creates a linear schedule from 1.0 to 1/num_steps.
-    """
-
     def __init__(self, runtime_config: "RuntimeConfig"):
         self.runtime_config = runtime_config
         self._sigmas = self._get_sigmas()
+        self._timesteps = self._get_timesteps()
 
     @property
     def sigmas(self) -> mx.array:
         return self._sigmas
 
+    @property
+    def timesteps(self) -> mx.array:
+        return self._timesteps
+
     def _get_sigmas(self) -> mx.array:
         model_config = self.runtime_config.model_config
         sigmas = mx.linspace(
@@ -44,6 +44,12 @@ class LinearScheduler(BaseScheduler):
         else:
             return sigmas
 
+    def _get_timesteps(self) -> mx.array:
+        num_steps = self.runtime_config.num_inference_steps
+        timesteps = mx.arange(num_steps, dtype=mx.float32)
+
+        return timesteps
+
     def step(self, model_output: mx.array, timestep: int, sample: mx.array, **kwargs) -> mx.array:
         dt = self._sigmas[timestep + 1] - self._sigmas[timestep]
         return sample + model_output * dt

src/mflux/utils/image_compare.py

@@ -0,0 +1,94 @@
+import os
+from pathlib import Path
+
+import cv2
+import numpy as np
+from PIL import Image
+
+
+class ImageCompare:
+    # How we determined DEFAULT_MISMATCH_THRESHOLD value: by eye test
+    # successive mlx/mflux versions have generated images that are visually close enough
+    # to consider as a valid upgrade. Minor visual differences are attributable to mlx updates
+    DEFAULT_MISMATCH_THRESHOLD = 0.15
+    ENV_MISMATCH_THRESHOLD = float(os.environ.get("MFLUX_IMAGE_MISMATCH_THRESHOLD", DEFAULT_MISMATCH_THRESHOLD))
+
+    @staticmethod
+    def check_images_close_enough(
+        image1_path: str | Path,
+        image2_path: str | Path,
+        error_message_prefix: str,
+        mismatch_threshold: float | None = None,
+    ) -> float:
+        if mismatch_threshold is None:
+            mismatch_threshold = ImageCompare.ENV_MISMATCH_THRESHOLD
+
+        image1_path = Path(image1_path)
+        image2_path = Path(image2_path)
+        image1_data = np.array(Image.open(image1_path))
+        image2_data = np.array(Image.open(image2_path))
+        closeness_array = np.isclose(
+            image1_data,
+            image2_data,
+            rtol=float(os.environ.get("MFLUX_IMAGE_ALLCLOSE_RTOL", 0.1)),
+            atol=0,  # ignore absolute tolerance
+        )
+        num_mismatched = np.count_nonzero(~closeness_array)
+        total_elements = image1_data.size
+        mismatch_ratio = num_mismatched / total_elements
+        if mismatch_ratio > mismatch_threshold:
+            diff = image1_data - image2_data
+            # Scale the difference so it's visible. A difference of 50 will become bright.
+            # We'll scale it so that a difference of 50 or more becomes pure white (255).
+            diff_visual = (np.abs(diff) / 50 * 255).clip(0, 255).astype(np.uint8)
+            cv2.imwrite(image2_path.with_stem(image1_path.stem + "_diff"), diff_visual)
+            raise AssertionError(
+                f"{error_message_prefix} Check {image1_path} vs {image2_path} :: their elements are {mismatch_ratio:.1%} different. Fails assertion for {mismatch_threshold=}"
+            )
+        return mismatch_ratio
+
+    @staticmethod
+    def compare_images(
+        image1_path: str | Path,
+        image2_path: str | Path,
+        mismatch_threshold: float | None = None,
+        create_diff: bool = True,
+    ) -> dict:
+        image1_path = Path(image1_path)
+        image2_path = Path(image2_path)
+        image1_data = np.array(Image.open(image1_path))
+        image2_data = np.array(Image.open(image2_path))
+
+        rtol = float(os.environ.get("MFLUX_IMAGE_ALLCLOSE_RTOL", 0.1))
+        closeness_array = np.isclose(
+            image1_data,
+            image2_data,
+            rtol=rtol,
+            atol=0,
+        )
+
+        num_mismatched = np.count_nonzero(~closeness_array)
+        total_elements = image1_data.size
+        mismatch_ratio = num_mismatched / total_elements
+
+        diff_path = None
+        if create_diff:
+            diff = image1_data - image2_data
+            diff_visual = (np.abs(diff) / 50 * 255).clip(0, 255).astype(np.uint8)
+            diff_path = image2_path.with_stem(image1_path.stem + "_diff")
+            cv2.imwrite(str(diff_path), diff_visual)
+
+        result = {
+            "mismatch_ratio": mismatch_ratio,
+            "total_pixels": total_elements,
+            "mismatched_pixels": num_mismatched,
+            "image1_shape": image1_data.shape,
+            "image2_shape": image2_data.shape,
+            "diff_path": str(diff_path) if diff_path else None,
+        }
+
+        if mismatch_threshold is not None:
+            result["passes_threshold"] = mismatch_ratio <= mismatch_threshold
+            result["threshold"] = mismatch_threshold
+
+        return result

src/mflux/utils/quantization_util.py

@@ -91,4 +91,4 @@ class QuantizationUtil:
             bits = int(q_level) if q_level is not None else quantize
             nn.quantize(vae, bits=bits)
             nn.quantize(transformer, bits=bits)
-            nn.quantize(text_encoder, bits=bits)
+            # nn.quantize(text_encoder, bits=bits)  # Quantization of text encoder causes significant semantic degradation

tests/depth/test_depth_pro.py

@@ -2,7 +2,7 @@ import os
 from pathlib import Path
 
 from mflux.models.depth_pro.depth_pro import DepthPro
-from tests.image_generation.helpers.image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class TestDepthPro:
@@ -24,7 +24,7 @@ class TestDepthPro:
             depth_result.depth_image.save(output_image_path)
 
             # Assert that the generated depth image matches the reference image
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated image doesn't match reference image.",

tests/image_generation/helpers/image_compare.py

@@ -1,46 +0,0 @@
-import os
-from pathlib import Path
-
-import cv2
-import numpy as np
-from PIL import Image
-
-
-class ReferenceVsOutputImageError(AssertionError): ...
-
-
-# How we determined DEFAULT_MISMATCH_THRESHOLD value: by eye test
-# successive mlx/mflux versions have generated images that are visually close enough
-# to consider as a valid upgrade. Minor visual differences are attributable to mlx updates
-DEFAULT_MISMATCH_THRESHOLD = 0.15
-ENV_MISMATCH_THRESHOLD = float(os.environ.get("MFLUX_IMAGE_MISMATCH_THRESHOLD", DEFAULT_MISMATCH_THRESHOLD))
-
-
-def check_images_close_enough(
-    image1_path: str | Path,
-    image2_path: str | Path,
-    error_message_prefix: str,
-    mismatch_threshold: float = ENV_MISMATCH_THRESHOLD,
-):
-    image1_path = Path(image1_path)
-    image2_path = Path(image2_path)
-    image1_data = np.array(Image.open(image1_path))
-    image2_data = np.array(Image.open(image2_path))
-    closeness_array = np.isclose(
-        image1_data,
-        image2_data,
-        rtol=float(os.environ.get("MFLUX_IMAGE_ALLCLOSE_RTOL", 0.1)),
-        atol=0,  # ignore absolute tolerance
-    )
-    num_mismatched = np.count_nonzero(~closeness_array)
-    total_elements = image1_data.size
-    mismatch_ratio = num_mismatched / total_elements
-    if mismatch_ratio > mismatch_threshold:
-        diff = image1_data - image2_data
-        # Scale the difference so it's visible. A difference of 50 will become bright.
-        # We'll scale it so that a difference of 50 or more becomes pure white (255).
-        diff_visual = (diff / 50 * 255).clip(0, 255).astype(np.uint8)
-        cv2.imwrite(image2_path.with_stem(image1_path.stem + "_diff"), diff_visual)
-        raise AssertionError(
-            f"{error_message_prefix} Check {image1_path} vs {image2_path} :: their elements are {mismatch_ratio:.1%} different. Fails assertion for {mismatch_threshold=}"
-        )

tests/image_generation/helpers/image_generation_concept_test_helper.py

@@ -5,8 +5,7 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.concept_attention.flux_concept import Flux1Concept
 from mflux.models.flux.variants.concept_attention.flux_concept_from_image import Flux1ConceptFromImage
-
-from .image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class ImageGenerationConceptTestHelper:
@@ -56,7 +55,7 @@ class ImageGenerationConceptTestHelper:
             image.save_concept_heatmap(path=output_heatmap_path, overwrite=True)
 
             # then - verify the heatmap matches reference
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_heatmap_path,
                 reference_heatmap_path,
                 "Generated concept heatmap doesn't match reference heatmap.",
@@ -117,7 +116,7 @@ class ImageGenerationConceptTestHelper:
             image.save_concept_heatmap(path=output_heatmap_path, overwrite=True)
 
             # then - verify the heatmap matches reference
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_heatmap_path,
                 reference_heatmap_path,
                 "Generated concept from image heatmap doesn't match reference heatmap.",

tests/image_generation/helpers/image_generation_controlnet_test_helper.py

@@ -3,10 +3,9 @@ import os
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet
+from mflux.utils.image_compare import ImageCompare
 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
 
-from .image_compare import check_images_close_enough
-
 
 class ImageGeneratorControlnetTestHelper:
     @staticmethod
@@ -54,7 +53,7 @@ class ImageGeneratorControlnetTestHelper:
             image.save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated controlnet image doesn't match reference image",

tests/image_generation/helpers/image_generation_depth_test_helper.py

@@ -4,8 +4,7 @@ from pathlib import Path
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.depth.flux_depth import Flux1Depth
-
-from .image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class ImageGeneratorDepthTestHelper:
@@ -46,7 +45,7 @@ class ImageGeneratorDepthTestHelper:
             image.save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated depth image doesn't match reference image.",

tests/image_generation/helpers/image_generation_edit_test_helper.py

@@ -0,0 +1,103 @@
+import os
+from pathlib import Path
+from typing import Any, Type
+
+from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.config.config import Config
+from mflux.config.model_config import ModelConfig
+from mflux.utils.image_compare import ImageCompare
+
+
+class ImageGeneratorEditTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_class: Type[Any],
+        model_config: ModelConfig,
+        steps: int,
+        seed: int,
+        height: int,
+        width: int,
+        prompt: str,
+        image_path: str | None = None,
+        guidance: float = 2.5,
+        negative_prompt: str | None = None,
+        quantize: int = 8,
+        image_paths: list[str] | None = None,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorEditTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorEditTestHelper.resolve_path(output_image_path)
+
+        # For Edit Plus, if image_paths is provided but image_path is not, use first element of image_paths
+        is_edit_plus = "Plus" in model_class.__name__
+        if is_edit_plus and image_paths and not image_path:
+            image_path = image_paths[0]
+
+        if image_path:
+            image_path = ImageGeneratorEditTestHelper.resolve_path(image_path)
+        if image_paths:
+            image_paths = [str(ImageGeneratorEditTestHelper.resolve_path(p)) for p in image_paths]
+
+        try:
+            # given
+            model_kwargs = {
+                "quantize": quantize,
+            }
+            model = model_class(**model_kwargs)
+
+            # when
+            config_kwargs = {
+                "num_inference_steps": steps,
+                "height": height,
+                "width": width,
+                "guidance": guidance,
+                "image_path": image_path,
+                "scheduler": "flow_match_euler_discrete",  # Match debug script
+            }
+            generate_kwargs = {
+                "seed": seed,
+                "prompt": prompt,
+                "negative_prompt": negative_prompt,
+                "config": Config(**config_kwargs),
+            }
+
+            # Edit Plus uses image_paths instead of image_path in config
+            # Check if it's Edit Plus by checking the class name
+            if "Plus" in model_class.__name__:
+                if image_paths:
+                    generate_kwargs["image_paths"] = image_paths
+                else:
+                    generate_kwargs["image_paths"] = [str(image_path)]
+
+            # Use a temporary directory for stepwise handler output
+            import tempfile
+
+            temp_dir = tempfile.mkdtemp()
+            handler = StepwiseHandler(model=model, output_dir=temp_dir)
+            CallbackRegistry.register_before_loop(handler)
+            CallbackRegistry.register_in_loop(handler)
+            CallbackRegistry.register_interrupt(handler)
+
+            image = model.generate_image(**generate_kwargs)
+            image.save(path=output_image_path, overwrite=True)
+
+            # then
+            model_name = "qwen edit"
+            ImageCompare.check_images_close_enough(
+                output_image_path,
+                reference_image_path,
+                f"Generated {model_name} 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

tests/image_generation/helpers/image_generation_fill_test_helper.py

@@ -4,10 +4,9 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.fill.flux_fill import Flux1Fill
 from mflux.ui import defaults as ui_defaults
+from mflux.utils.image_compare import ImageCompare
 from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
 
-from .image_compare import check_images_close_enough
-
 
 class ImageGeneratorFillTestHelper:
     @staticmethod
@@ -56,7 +55,7 @@ class ImageGeneratorFillTestHelper:
             image.save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated fill image doesn't match reference image.",

tests/image_generation/helpers/image_generation_ic_edit_test_helper.py

@@ -5,8 +5,7 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill
 from mflux.models.flux.variants.in_context.utils.in_context_loras import prepare_ic_edit_loras
-
-from .image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class ImageGeneratorICEditTestHelper:
@@ -55,7 +54,7 @@ class ImageGeneratorICEditTestHelper:
             image.get_right_half().save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated ic-edit image doesn't match reference image.",

tests/image_generation/helpers/image_generation_in_context_test_helper.py

@@ -5,8 +5,7 @@ from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.in_context.flux_in_context_dev import Flux1InContextDev
 from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename
-
-from .image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class ImageGeneratorInContextTestHelper:
@@ -58,7 +57,7 @@ class ImageGeneratorInContextTestHelper:
             image.get_right_half().save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated in-context image doesn't match reference image.",

tests/image_generation/helpers/image_generation_kontext_test_helper.py

@@ -1,65 +0,0 @@
-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 .image_compare import check_images_close_enough
-
-
-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
-            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

tests/image_generation/helpers/image_generation_redux_test_helper.py

@@ -4,8 +4,7 @@ from pathlib import Path
 from mflux.config.config import Config
 from mflux.config.model_config import ModelConfig
 from mflux.models.flux.variants.redux.flux_redux import Flux1Redux
-
-from .image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class ImageGeneratorReduxTestHelper:
@@ -47,7 +46,7 @@ class ImageGeneratorReduxTestHelper:
             image.save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated redux image doesn't match reference image.",

tests/image_generation/helpers/image_generation_test_helper.py

@@ -6,8 +6,7 @@ from mflux.config.config import Config
 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 .image_compare import check_images_close_enough
+from mflux.utils.image_compare import ImageCompare
 
 
 class ImageGeneratorTestHelper:
@@ -80,7 +79,7 @@ class ImageGeneratorTestHelper:
             image.save(path=output_image_path, overwrite=True)
 
             # then
-            check_images_close_enough(
+            ImageCompare.check_images_close_enough(
                 output_image_path,
                 reference_image_path,
                 "Generated image doesn't match reference image.",

tests/image_generation/test_generate_image_edit.py

@@ -0,0 +1,69 @@
+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, QwenImageEditPlus
+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",
+        )
+
+    def test_image_generation_qwen_edit(self):
+        ImageGeneratorEditTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_qwen_edit.png",
+            output_image_path="output_qwen_edit.png",
+            model_class=QwenImageEdit,
+            model_config=ModelConfig.qwen_image_edit(),
+            steps=20,
+            seed=4869845,
+            height=384,
+            width=640,
+            guidance=2.5,
+            quantize=6,
+            prompt="Make the hand fistbump the camera instead of showing a flat palm",
+            image_path="reference_upscaled.png",
+        )
+
+    def test_image_generation_qwen_edit_plus(self):
+        ImageGeneratorEditTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_qwen_edit_plus.png",
+            output_image_path="output_qwen_edit_plus.png",
+            model_class=QwenImageEditPlus,
+            model_config=ModelConfig.qwen_image_edit_plus(),
+            steps=20,
+            seed=4869845,
+            height=384,
+            width=640,
+            guidance=2.5,
+            quantize=6,
+            prompt="Make the hand fistbump the camera instead of showing a flat palm",
+            image_path="reference_upscaled.png",
+        )
+
+    def test_image_generation_qwen_edit_plus_multiple_images(self):
+        ImageGeneratorEditTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_qwen_edit_plus_multiple_images.png",
+            output_image_path="output_qwen_edit_plus_multiple_images.png",
+            model_class=QwenImageEditPlus,
+            model_config=ModelConfig.qwen_image_edit_plus(),
+            steps=20,
+            seed=4869845,
+            height=384,
+            width=640,
+            guidance=2.5,
+            quantize=8,
+            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"],
+        )

tests/image_generation/test_generate_image_kontext.py

@@ -1,18 +0,0 @@
-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",
-        )