[object Object]

← back to Exo

Ciaran/flux1 kontext (#1394)

6f0cb99919b710a57f73f30d357119721c82ccb0 · 2026-02-06 16:20:31 +0000 · ciaranbor

## Motivation

Add support for FLUX.1-Kontext-dev, an image editing variant of
FLUX.1-dev

## Changes

- New FluxKontextModelAdapter: Handles Kontext's image-to-image workflow
- encodes input image as conditioning latents with special position IDs,
generates from pure noise
- Model config: 57 transformer blocks (19 joint + 38 single), guidance
scale 4.0, ImageToImage task
- Pipeline updates: Added kontext_image_ids property to PromptData
interface, passed through diffusion runner
  - Model cards: Added TOML configs for base, 4-bit, and 8-bit variants
  - Dependency: mflux 0.15.4 → 0.15.5
- Utility: tmp/quantize_and_upload.py for quantizing and uploading
models to HuggingFace

## Test Plan

### Manual Testing

Works better than Qwen-Image-Edit

Files touched

Diff

commit 6f0cb99919b710a57f73f30d357119721c82ccb0
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date:   Fri Feb 6 16:20:31 2026 +0000

    Ciaran/flux1 kontext (#1394)
    
    ## Motivation
    
    Add support for FLUX.1-Kontext-dev, an image editing variant of
    FLUX.1-dev
    
    ## Changes
    
    - New FluxKontextModelAdapter: Handles Kontext's image-to-image workflow
    - encodes input image as conditioning latents with special position IDs,
    generates from pure noise
    - Model config: 57 transformer blocks (19 joint + 38 single), guidance
    scale 4.0, ImageToImage task
    - Pipeline updates: Added kontext_image_ids property to PromptData
    interface, passed through diffusion runner
      - Model cards: Added TOML configs for base, 4-bit, and 8-bit variants
      - Dependency: mflux 0.15.4 → 0.15.5
    - Utility: tmp/quantize_and_upload.py for quantizing and uploading
    models to HuggingFace
    
    ## Test Plan
    
    ### Manual Testing
    
    Works better than Qwen-Image-Edit
---
 .gitignore                                         |   3 +
 .../models/flux/variants/kontext/__init__.pyi      |   7 +
 .../models/flux/variants/kontext/flux_kontext.pyi  |  49 +++
 .../models/flux/variants/kontext/kontext_util.pyi  |  16 +
 pyproject.toml                                     |   2 +-
 .../exolabs--FLUX.1-Kontext-dev-4bit.toml          |  45 +++
 .../exolabs--FLUX.1-Kontext-dev-8bit.toml          |  45 +++
 .../exolabs--FLUX.1-Kontext-dev.toml               |  45 +++
 src/exo/worker/engines/image/models/__init__.py    |   5 +
 src/exo/worker/engines/image/models/base.py        |  13 +
 .../worker/engines/image/models/flux/__init__.py   |   6 +
 .../worker/engines/image/models/flux/adapter.py    |   4 +
 src/exo/worker/engines/image/models/flux/config.py |  16 +
 .../engines/image/models/flux/kontext_adapter.py   | 348 +++++++++++++++++++
 .../worker/engines/image/models/qwen/adapter.py    |   4 +
 .../engines/image/models/qwen/edit_adapter.py      |   4 +
 src/exo/worker/engines/image/pipeline/runner.py    |   9 +-
 tmp/quantize_and_upload.py                         | 377 +++++++++++++++++++++
 uv.lock                                            |  20 +-
 19 files changed, 999 insertions(+), 19 deletions(-)

diff --git a/.gitignore b/.gitignore
index 15979672..3afebe2e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -35,3 +35,6 @@ hosts_*.json
 
 # bench files
 bench/**/*.json
+
+# tmp
+tmp/models
diff --git a/.mlx_typings/mflux/models/flux/variants/kontext/__init__.pyi b/.mlx_typings/mflux/models/flux/variants/kontext/__init__.pyi
new file mode 100644
index 00000000..36c5443a
--- /dev/null
+++ b/.mlx_typings/mflux/models/flux/variants/kontext/__init__.pyi
@@ -0,0 +1,7 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
+
+__all__ = ["Flux1Kontext"]
diff --git a/.mlx_typings/mflux/models/flux/variants/kontext/flux_kontext.pyi b/.mlx_typings/mflux/models/flux/variants/kontext/flux_kontext.pyi
new file mode 100644
index 00000000..8050e68b
--- /dev/null
+++ b/.mlx_typings/mflux/models/flux/variants/kontext/flux_kontext.pyi
@@ -0,0 +1,49 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+from pathlib import Path
+from typing import Any
+
+from mlx import nn
+
+from mflux.models.common.config.model_config import ModelConfig
+from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import (
+    CLIPEncoder,
+)
+from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.flux.model.flux_transformer.transformer import Transformer
+from mflux.models.flux.model.flux_vae.vae import VAE
+from mflux.utils.generated_image import GeneratedImage
+
+class Flux1Kontext(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+    bits: int | None
+    lora_paths: list[str] | None
+    lora_scales: list[float] | None
+    prompt_cache: dict[str, Any]
+    tokenizers: dict[str, Any]
+
+    def __init__(
+        self,
+        quantize: int | None = ...,
+        model_path: str | None = ...,
+        lora_paths: list[str] | None = ...,
+        lora_scales: list[float] | None = ...,
+        model_config: ModelConfig = ...,
+    ) -> None: ...
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        num_inference_steps: int = ...,
+        height: int = ...,
+        width: int = ...,
+        guidance: float = ...,
+        image_path: Path | str | None = ...,
+        image_strength: float | None = ...,
+        scheduler: str = ...,
+    ) -> GeneratedImage: ...
diff --git a/.mlx_typings/mflux/models/flux/variants/kontext/kontext_util.pyi b/.mlx_typings/mflux/models/flux/variants/kontext/kontext_util.pyi
new file mode 100644
index 00000000..c7588ec6
--- /dev/null
+++ b/.mlx_typings/mflux/models/flux/variants/kontext/kontext_util.pyi
@@ -0,0 +1,16 @@
+"""
+This type stub file was generated by pyright.
+"""
+
+import mlx.core as mx
+
+from mflux.models.flux.model.flux_vae.vae import VAE
+
+class KontextUtil:
+    @staticmethod
+    def create_image_conditioning_latents(
+        vae: VAE,
+        height: int,
+        width: int,
+        image_path: str,
+    ) -> tuple[mx.array, mx.array]: ...
diff --git a/pyproject.toml b/pyproject.toml
index dd84cf17..ee219b74 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,7 @@ dependencies = [
     "httpx>=0.28.1",
     "tomlkit>=0.14.0",
     "pillow>=11.0,<12.0", # compatibility with mflux
-    "mflux==0.15.4",
+    "mflux==0.15.5",
     "python-multipart>=0.0.21",
 ]
 
diff --git a/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-4bit.toml b/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-4bit.toml
new file mode 100644
index 00000000..baf305ca
--- /dev/null
+++ b/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-4bit.toml
@@ -0,0 +1,45 @@
+model_id = "exolabs/FLUX.1-Kontext-dev-4bit"
+n_layers = 57
+hidden_size = 1
+supports_tensor = false
+tasks = ["ImageToImage"]
+
+[storage_size]
+in_bytes = 15475325472
+
+[[components]]
+component_name = "text_encoder"
+component_path = "text_encoder/"
+n_layers = 12
+can_shard = false
+
+[components.storage_size]
+in_bytes = 0
+
+[[components]]
+component_name = "text_encoder_2"
+component_path = "text_encoder_2/"
+n_layers = 24
+can_shard = false
+safetensors_index_filename = "model.safetensors.index.json"
+
+[components.storage_size]
+in_bytes = 9524621312
+
+[[components]]
+component_name = "transformer"
+component_path = "transformer/"
+n_layers = 57
+can_shard = true
+safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json"
+
+[components.storage_size]
+in_bytes = 5950704160
+
+[[components]]
+component_name = "vae"
+component_path = "vae/"
+can_shard = false
+
+[components.storage_size]
+in_bytes = 0
diff --git a/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-8bit.toml b/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-8bit.toml
new file mode 100644
index 00000000..ce0809c2
--- /dev/null
+++ b/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-8bit.toml
@@ -0,0 +1,45 @@
+model_id = "exolabs/FLUX.1-Kontext-dev-8bit"
+n_layers = 57
+hidden_size = 1
+supports_tensor = false
+tasks = ["ImageToImage"]
+
+[storage_size]
+in_bytes = 21426029632
+
+[[components]]
+component_name = "text_encoder"
+component_path = "text_encoder/"
+n_layers = 12
+can_shard = false
+
+[components.storage_size]
+in_bytes = 0
+
+[[components]]
+component_name = "text_encoder_2"
+component_path = "text_encoder_2/"
+n_layers = 24
+can_shard = false
+safetensors_index_filename = "model.safetensors.index.json"
+
+[components.storage_size]
+in_bytes = 9524621312
+
+[[components]]
+component_name = "transformer"
+component_path = "transformer/"
+n_layers = 57
+can_shard = true
+safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json"
+
+[components.storage_size]
+in_bytes = 11901408320
+
+[[components]]
+component_name = "vae"
+component_path = "vae/"
+can_shard = false
+
+[components.storage_size]
+in_bytes = 0
diff --git a/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev.toml b/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev.toml
new file mode 100644
index 00000000..2ebb0c43
--- /dev/null
+++ b/resources/image_model_cards/exolabs--FLUX.1-Kontext-dev.toml
@@ -0,0 +1,45 @@
+model_id = "exolabs/FLUX.1-Kontext-dev"
+n_layers = 57
+hidden_size = 1
+supports_tensor = false
+tasks = ["ImageToImage"]
+
+[storage_size]
+in_bytes = 33327437952
+
+[[components]]
+component_name = "text_encoder"
+component_path = "text_encoder/"
+n_layers = 12
+can_shard = false
+
+[components.storage_size]
+in_bytes = 0
+
+[[components]]
+component_name = "text_encoder_2"
+component_path = "text_encoder_2/"
+n_layers = 24
+can_shard = false
+safetensors_index_filename = "model.safetensors.index.json"
+
+[components.storage_size]
+in_bytes = 9524621312
+
+[[components]]
+component_name = "transformer"
+component_path = "transformer/"
+n_layers = 57
+can_shard = true
+safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json"
+
+[components.storage_size]
+in_bytes = 23802816640
+
+[[components]]
+component_name = "vae"
+component_path = "vae/"
+can_shard = false
+
+[components.storage_size]
+in_bytes = 0
diff --git a/src/exo/worker/engines/image/models/__init__.py b/src/exo/worker/engines/image/models/__init__.py
index dc0a9d8c..9f58cf9d 100644
--- a/src/exo/worker/engines/image/models/__init__.py
+++ b/src/exo/worker/engines/image/models/__init__.py
@@ -5,7 +5,9 @@ from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models.base import ModelAdapter
 from exo.worker.engines.image.models.flux import (
     FLUX_DEV_CONFIG,
+    FLUX_KONTEXT_CONFIG,
     FLUX_SCHNELL_CONFIG,
+    FluxKontextModelAdapter,
     FluxModelAdapter,
 )
 from exo.worker.engines.image.models.qwen import (
@@ -26,13 +28,16 @@ AdapterFactory = Callable[
 # Registry maps model_family string to adapter factory
 _ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
     "flux": FluxModelAdapter,
+    "flux-kontext": FluxKontextModelAdapter,
     "qwen-edit": QwenEditModelAdapter,
     "qwen": QwenModelAdapter,
 }
 
 # Config registry: maps model ID patterns to configs
+# Order matters: longer/more-specific patterns must come before shorter ones
 _CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
     "flux.1-schnell": FLUX_SCHNELL_CONFIG,
+    "flux.1-kontext": FLUX_KONTEXT_CONFIG,  # Must come before "flux.1-dev" for pattern matching
     "flux.1-krea-dev": FLUX_DEV_CONFIG,  # Must come before "flux.1-dev" for pattern matching
     "flux.1-dev": FLUX_DEV_CONFIG,
     "qwen-image-edit": QWEN_IMAGE_EDIT_CONFIG,  # Must come before "qwen-image" for pattern matching
diff --git a/src/exo/worker/engines/image/models/base.py b/src/exo/worker/engines/image/models/base.py
index f77ea882..0aa660e2 100644
--- a/src/exo/worker/engines/image/models/base.py
+++ b/src/exo/worker/engines/image/models/base.py
@@ -66,6 +66,19 @@ class PromptData(ABC):
         """
         ...
 
+    @property
+    @abstractmethod
+    def kontext_image_ids(self) -> mx.array | None:
+        """Kontext-style position IDs for image conditioning.
+
+        For FLUX.1-Kontext models, returns position IDs with first_coord=1
+        to distinguish conditioning tokens from generation tokens (first_coord=0).
+
+        Returns:
+            Position IDs array [1, seq_len, 3] for Kontext, None for other models.
+        """
+        ...
+
     @abstractmethod
     def get_batched_cfg_data(
         self,
diff --git a/src/exo/worker/engines/image/models/flux/__init__.py b/src/exo/worker/engines/image/models/flux/__init__.py
index 3adc2626..ac3e9335 100644
--- a/src/exo/worker/engines/image/models/flux/__init__.py
+++ b/src/exo/worker/engines/image/models/flux/__init__.py
@@ -1,11 +1,17 @@
 from exo.worker.engines.image.models.flux.adapter import FluxModelAdapter
 from exo.worker.engines.image.models.flux.config import (
     FLUX_DEV_CONFIG,
+    FLUX_KONTEXT_CONFIG,
     FLUX_SCHNELL_CONFIG,
 )
+from exo.worker.engines.image.models.flux.kontext_adapter import (
+    FluxKontextModelAdapter,
+)
 
 __all__ = [
     "FluxModelAdapter",
+    "FluxKontextModelAdapter",
     "FLUX_DEV_CONFIG",
+    "FLUX_KONTEXT_CONFIG",
     "FLUX_SCHNELL_CONFIG",
 ]
diff --git a/src/exo/worker/engines/image/models/flux/adapter.py b/src/exo/worker/engines/image/models/flux/adapter.py
index 1aa510da..90b7bacd 100644
--- a/src/exo/worker/engines/image/models/flux/adapter.py
+++ b/src/exo/worker/engines/image/models/flux/adapter.py
@@ -59,6 +59,10 @@ class FluxPromptData(PromptData):
     def conditioning_latents(self) -> mx.array | None:
         return None
 
+    @property
+    def kontext_image_ids(self) -> mx.array | None:
+        return None
+
     def get_batched_cfg_data(
         self,
     ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
diff --git a/src/exo/worker/engines/image/models/flux/config.py b/src/exo/worker/engines/image/models/flux/config.py
index d37161b0..a85bd0c8 100644
--- a/src/exo/worker/engines/image/models/flux/config.py
+++ b/src/exo/worker/engines/image/models/flux/config.py
@@ -32,3 +32,19 @@ FLUX_DEV_CONFIG = ImageModelConfig(
     default_steps={"low": 10, "medium": 25, "high": 50},
     num_sync_steps=4,
 )
+
+
+FLUX_KONTEXT_CONFIG = ImageModelConfig(
+    model_family="flux-kontext",
+    block_configs=(
+        TransformerBlockConfig(
+            block_type=BlockType.JOINT, count=19, has_separate_text_output=True
+        ),
+        TransformerBlockConfig(
+            block_type=BlockType.SINGLE, count=38, has_separate_text_output=False
+        ),
+    ),
+    default_steps={"low": 10, "medium": 25, "high": 50},
+    num_sync_steps=4,
+    guidance_scale=4.0,
+)
diff --git a/src/exo/worker/engines/image/models/flux/kontext_adapter.py b/src/exo/worker/engines/image/models/flux/kontext_adapter.py
new file mode 100644
index 00000000..19d0be56
--- /dev/null
+++ b/src/exo/worker/engines/image/models/flux/kontext_adapter.py
@@ -0,0 +1,348 @@
+import math
+from pathlib import Path
+from typing import Any, final
+
+import mlx.core as mx
+from mflux.models.common.config.config import Config
+from mflux.models.common.config.model_config import ModelConfig
+from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
+from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder
+from mflux.models.flux.model.flux_transformer.transformer import Transformer
+from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
+from mflux.models.flux.variants.kontext.kontext_util import KontextUtil
+
+from exo.worker.engines.image.config import ImageModelConfig
+from exo.worker.engines.image.models.base import (
+    ModelAdapter,
+    PromptData,
+    RotaryEmbeddings,
+)
+from exo.worker.engines.image.models.flux.wrappers import (
+    FluxJointBlockWrapper,
+    FluxSingleBlockWrapper,
+)
+from exo.worker.engines.image.pipeline.block_wrapper import (
+    JointBlockWrapper,
+    SingleBlockWrapper,
+)
+
+
+@final
+class FluxKontextPromptData(PromptData):
+    """Prompt data for FLUX.1-Kontext image editing.
+
+    Stores text embeddings along with conditioning latents and position IDs
+    for the input image.
+    """
+
+    def __init__(
+        self,
+        prompt_embeds: mx.array,
+        pooled_prompt_embeds: mx.array,
+        conditioning_latents: mx.array,
+        kontext_image_ids: mx.array,
+    ):
+        self._prompt_embeds = prompt_embeds
+        self._pooled_prompt_embeds = pooled_prompt_embeds
+        self._conditioning_latents = conditioning_latents
+        self._kontext_image_ids = kontext_image_ids
+
+    @property
+    def prompt_embeds(self) -> mx.array:
+        return self._prompt_embeds
+
+    @property
+    def pooled_prompt_embeds(self) -> mx.array:
+        return self._pooled_prompt_embeds
+
+    @property
+    def negative_prompt_embeds(self) -> mx.array | None:
+        return None
+
+    @property
+    def negative_pooled_prompt_embeds(self) -> mx.array | None:
+        return None
+
+    def get_encoder_hidden_states_mask(self, positive: bool = True) -> mx.array | None:
+        return None
+
+    @property
+    def cond_image_grid(
+        self,
+    ) -> tuple[int, int, int] | list[tuple[int, int, int]] | None:
+        return None
+
+    @property
+    def conditioning_latents(self) -> mx.array | None:
+        """VAE-encoded input image latents for Kontext conditioning."""
+        return self._conditioning_latents
+
+    @property
+    def kontext_image_ids(self) -> mx.array | None:
+        """Position IDs for Kontext conditioning (first_coord=1)."""
+        return self._kontext_image_ids
+
+    def get_cfg_branch_data(
+        self, positive: bool
+    ) -> tuple[mx.array, mx.array | None, mx.array | None, mx.array | None]:
+        """Kontext doesn't use CFG, but we return positive data for compatibility."""
+        return (
+            self._prompt_embeds,
+            None,
+            self._pooled_prompt_embeds,
+            self._conditioning_latents,
+        )
+
+    def get_batched_cfg_data(
+        self,
+    ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
+        # Kontext doesn't use CFG
+        return None
+
+
+@final
+class FluxKontextModelAdapter(ModelAdapter[Flux1Kontext, Transformer]):
+    """Adapter for FLUX.1-Kontext image editing model.
+
+    Key differences from standard FluxModelAdapter:
+    - Takes an input image and computes output dimensions from it
+    - Creates conditioning latents from the input image via VAE
+    - Creates special position IDs (kontext_image_ids) for conditioning tokens
+    - Creates pure noise latents (not img2img blending)
+    """
+
+    def __init__(
+        self,
+        config: ImageModelConfig,
+        model_id: str,
+        local_path: Path,
+        quantize: int | None = None,
+    ):
+        self._config = config
+        self._model = Flux1Kontext(
+            model_config=ModelConfig.from_name(model_name=model_id, base_model=None),
+            model_path=str(local_path),
+            quantize=quantize,
+        )
+        self._transformer = self._model.transformer
+
+        # Stores image path and computed dimensions after set_image_dimensions
+        self._image_path: str | None = None
+        self._output_height: int | None = None
+        self._output_width: int | None = None
+
+    @property
+    def hidden_dim(self) -> int:
+        return self._transformer.x_embedder.weight.shape[0]  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
+
+    @property
+    def needs_cfg(self) -> bool:
+        return False
+
+    def _get_latent_creator(self) -> type:
+        return FluxLatentCreator
+
+    def get_joint_block_wrappers(
+        self,
+        text_seq_len: int,
+        encoder_hidden_states_mask: mx.array | None = None,
+    ) -> list[JointBlockWrapper[Any]]:
+        """Create wrapped joint blocks for Flux Kontext."""
+        return [
+            FluxJointBlockWrapper(block, text_seq_len)
+            for block in self._transformer.transformer_blocks
+        ]
+
+    def get_single_block_wrappers(
+        self,
+        text_seq_len: int,
+    ) -> list[SingleBlockWrapper[Any]]:
+        """Create wrapped single blocks for Flux Kontext."""
+        return [
+            FluxSingleBlockWrapper(block, text_seq_len)
+            for block in self._transformer.single_transformer_blocks
+        ]
+
+    def slice_transformer_blocks(
+        self,
+        start_layer: int,
+        end_layer: int,
+    ):
+        all_joint = list(self._transformer.transformer_blocks)
+        all_single = list(self._transformer.single_transformer_blocks)
+        total_joint_blocks = len(all_joint)
+        if end_layer <= total_joint_blocks:
+            # All assigned are joint blocks
+            joint_start, joint_end = start_layer, end_layer
+            single_start, single_end = 0, 0
+        elif start_layer >= total_joint_blocks:
+            # All assigned are single blocks
+            joint_start, joint_end = 0, 0
+            single_start = start_layer - total_joint_blocks
+            single_end = end_layer - total_joint_blocks
+        else:
+            # Spans both joint and single
+            joint_start, joint_end = start_layer, total_joint_blocks
+            single_start = 0
+            single_end = end_layer - total_joint_blocks
+
+        self._transformer.transformer_blocks = all_joint[joint_start:joint_end]
+        self._transformer.single_transformer_blocks = all_single[
+            single_start:single_end
+        ]
+
+    def set_image_dimensions(self, image_path: Path) -> tuple[int, int]:
+        """Compute and store dimensions from input image.
+
+        Also stores image_path for use in encode_prompt().
+
+        Args:
+            image_path: Path to the input image
+
+        Returns:
+            (output_width, output_height) for runtime config
+        """
+        from mflux.utils.image_util import ImageUtil
+
+        pil_image = ImageUtil.load_image(str(image_path)).convert("RGB")
+        image_size = pil_image.size
+
+        # Compute output dimensions from input image aspect ratio
+        # Target area of 1024x1024 = ~1M pixels
+        target_area = 1024 * 1024
+        ratio = image_size[0] / image_size[1]
+        output_width = math.sqrt(target_area * ratio)
+        output_height = output_width / ratio
+        output_width = round(output_width / 32) * 32
+        output_height = round(output_height / 32) * 32
+
+        # Ensure multiple of 16 for VAE
+        vae_scale_factor = 8
+        multiple_of = vae_scale_factor * 2
+        output_width = output_width // multiple_of * multiple_of
+        output_height = output_height // multiple_of * multiple_of
+
+        self._image_path = str(image_path)
+        self._output_width = int(output_width)
+        self._output_height = int(output_height)
+
+        return self._output_width, self._output_height
+
+    def create_latents(self, seed: int, runtime_config: Config) -> mx.array:
+        """Create initial noise latents for Kontext.
+
+        Unlike standard img2img which blends noise with encoded input,
+        Kontext uses pure noise latents. The input image is provided
+        separately as conditioning.
+        """
+        return FluxLatentCreator.create_noise(
+            seed=seed,
+            height=runtime_config.height,
+            width=runtime_config.width,
+        )
+
+    def encode_prompt(
+        self, prompt: str, negative_prompt: str | None = None
+    ) -> FluxKontextPromptData:
+        """Encode prompt and create conditioning from stored input image.
+
+        Must call set_image_dimensions() before this method.
+
+        Args:
+            prompt: Text prompt for editing
+            negative_prompt: Ignored (Kontext doesn't use CFG)
+
+        Returns:
+            FluxKontextPromptData with text embeddings and image conditioning
+        """
+        del negative_prompt  # Kontext doesn't support negative prompts or CFG
+
+        if (
+            self._image_path is None
+            or self._output_height is None
+            or self._output_width is None
+        ):
+            raise RuntimeError(
+                "set_image_dimensions() must be called before encode_prompt() "
+                "for FluxKontextModelAdapter"
+            )
+
+        assert isinstance(self.model.prompt_cache, dict)
+        assert isinstance(self.model.tokenizers, dict)
+
+        # Encode text prompt
+        prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=self.model.prompt_cache,
+            t5_tokenizer=self.model.tokenizers["t5"],  # pyright: ignore[reportAny]
+            clip_tokenizer=self.model.tokenizers["clip"],  # pyright: ignore[reportAny]
+            t5_text_encoder=self.model.t5_text_encoder,
+            clip_text_encoder=self.model.clip_text_encoder,
+        )
+
+        # Create conditioning latents from input image
+        conditioning_latents, kontext_image_ids = (
+            KontextUtil.create_image_conditioning_latents(
+                vae=self.model.vae,
+                height=self._output_height,
+                width=self._output_width,
+                image_path=self._image_path,
+            )
+        )
+
+        return FluxKontextPromptData(
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            conditioning_latents=conditioning_latents,
+            kontext_image_ids=kontext_image_ids,
+        )
+
+    def compute_embeddings(
+        self,
+        hidden_states: mx.array,
+        prompt_embeds: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        embedded_hidden = self._transformer.x_embedder(hidden_states)
+        embedded_encoder = self._transformer.context_embedder(prompt_embeds)
+        return embedded_hidden, embedded_encoder
+
+    def compute_text_embeddings(
+        self,
+        t: int,
+        runtime_config: Config,
+        pooled_prompt_embeds: mx.array | None = None,
+        hidden_states: mx.array | None = None,
+    ) -> mx.array:
+        if pooled_prompt_embeds is None:
+            raise ValueError(
+                "pooled_prompt_embeds is required for Flux Kontext text embeddings"
+            )
+
+        return Transformer.compute_text_embeddings(
+            t, pooled_prompt_embeds, self._transformer.time_text_embed, runtime_config
+        )
+
+    def compute_rotary_embeddings(
+        self,
+        prompt_embeds: mx.array,
+        runtime_config: Config,
+        encoder_hidden_states_mask: mx.array | None = None,
+        cond_image_grid: tuple[int, int, int]
+        | list[tuple[int, int, int]]
+        | None = None,
+        kontext_image_ids: mx.array | None = None,
+    ) -> RotaryEmbeddings:
+        return Transformer.compute_rotary_embeddings(
+            prompt_embeds,
+            self._transformer.pos_embed,
+            runtime_config,
+            kontext_image_ids,
+        )
+
+    def apply_guidance(
+        self,
+        noise_positive: mx.array,
+        noise_negative: mx.array,
+        guidance_scale: float,
+    ) -> mx.array:
+        raise NotImplementedError("Flux Kontext does not use classifier-free guidance")
diff --git a/src/exo/worker/engines/image/models/qwen/adapter.py b/src/exo/worker/engines/image/models/qwen/adapter.py
index e88d2a75..be1edb0c 100644
--- a/src/exo/worker/engines/image/models/qwen/adapter.py
+++ b/src/exo/worker/engines/image/models/qwen/adapter.py
@@ -69,6 +69,10 @@ class QwenPromptData(PromptData):
     def conditioning_latents(self) -> mx.array | None:
         return None
 
+    @property
+    def kontext_image_ids(self) -> mx.array | None:
+        return None
+
     def get_batched_cfg_data(
         self,
     ) -> tuple[mx.array, mx.array, mx.array | None, mx.array | None] | None:
diff --git a/src/exo/worker/engines/image/models/qwen/edit_adapter.py b/src/exo/worker/engines/image/models/qwen/edit_adapter.py
index 4a88a4e3..fee79738 100644
--- a/src/exo/worker/engines/image/models/qwen/edit_adapter.py
+++ b/src/exo/worker/engines/image/models/qwen/edit_adapter.py
@@ -85,6 +85,10 @@ class QwenEditPromptData(PromptData):
     def qwen_image_ids(self) -> mx.array:
         return self._qwen_image_ids
 
+    @property
+    def kontext_image_ids(self) -> mx.array | None:
+        return None
+
     @property
     def is_edit_mode(self) -> bool:
         return True
diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py
index f7054763..01ad7b05 100644
--- a/src/exo/worker/engines/image/pipeline/runner.py
+++ b/src/exo/worker/engines/image/pipeline/runner.py
@@ -567,6 +567,7 @@ class DiffusionRunner:
         | list[tuple[int, int, int]]
         | None = None,
         conditioning_latents: mx.array | None = None,
+        kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         """Run a single forward pass through the transformer.
         Args:
@@ -578,6 +579,7 @@ class DiffusionRunner:
             encoder_hidden_states_mask: Attention mask for text (Qwen)
             cond_image_grid: Conditioning image grid dimensions (Qwen edit)
             conditioning_latents: Conditioning latents for edit mode
+            kontext_image_ids: Position IDs for Kontext conditioning (Flux Kontext)
 
         Returns:
             Noise prediction tensor
@@ -610,6 +612,7 @@ class DiffusionRunner:
             config,
             encoder_hidden_states_mask=encoder_hidden_states_mask,
             cond_image_grid=cond_image_grid,
+            kontext_image_ids=kontext_image_ids,
         )
 
         assert self.joint_block_wrappers is not None
@@ -681,6 +684,7 @@ class DiffusionRunner:
         prompt_data: PromptData,
     ) -> mx.array:
         cond_image_grid = prompt_data.cond_image_grid
+        kontext_image_ids = prompt_data.kontext_image_ids
         results: list[tuple[bool, mx.array]] = []
 
         for branch in self._get_cfg_branches(prompt_data):
@@ -700,6 +704,7 @@ class DiffusionRunner:
                 encoder_hidden_states_mask=branch.mask,
                 cond_image_grid=cond_image_grid,
                 conditioning_latents=branch.cond_latents,
+                kontext_image_ids=kontext_image_ids,
             )
             results.append((branch.positive, noise))
 
@@ -902,10 +907,10 @@ class DiffusionRunner:
         config: Config,
         hidden_states: mx.array,
         prompt_data: PromptData,
-        kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         prev_latents = hidden_states
         cond_image_grid = prompt_data.cond_image_grid
+        kontext_image_ids = prompt_data.kontext_image_ids
 
         scaled_hidden_states = config.scheduler.scale_model_input(hidden_states, t)  # pyright: ignore[reportAny]
         original_latent_tokens: int = scaled_hidden_states.shape[1]  # pyright: ignore[reportAny]
@@ -979,10 +984,10 @@ class DiffusionRunner:
         latents: mx.array,
         prompt_data: PromptData,
         is_first_async_step: bool,
-        kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
         patch_latents, token_indices = self._create_patches(latents, config)
         cond_image_grid = prompt_data.cond_image_grid
+        kontext_image_ids = prompt_data.kontext_image_ids
 
         prev_patch_latents = [p for p in patch_latents]
 
diff --git a/tmp/quantize_and_upload.py b/tmp/quantize_and_upload.py
new file mode 100755
index 00000000..ee421a14
--- /dev/null
+++ b/tmp/quantize_and_upload.py
@@ -0,0 +1,377 @@
+#!/usr/bin/env python3
+"""
+Download an mflux model, quantize it, and upload to HuggingFace.
+
+Usage (run from mflux project directory):
+    cd /path/to/mflux
+    uv run python /path/to/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev
+    uv run python /path/to/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --skip-base --skip-8bit
+    uv run python /path/to/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --dry-run
+
+Requires:
+    - Must be run from mflux project directory using `uv run`
+    - huggingface_hub installed (add to mflux deps or install separately)
+    - HuggingFace authentication: run `huggingface-cli login` or set HF_TOKEN
+"""
+
+from __future__ import annotations
+
+import argparse
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from mflux.models.flux.variants.txt2img.flux import Flux1
+
+
+HF_ORG = "exolabs"
+
+
+def get_model_class(model_name: str) -> type:
+    """Get the appropriate model class based on model name."""
+    from mflux.models.fibo.variants.txt2img.fibo import FIBO
+    from mflux.models.flux.variants.txt2img.flux import Flux1
+    from mflux.models.flux2.variants.txt2img.flux2_klein import Flux2Klein
+    from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
+    from mflux.models.z_image.variants.turbo.z_image_turbo import ZImageTurbo
+
+    model_name_lower = model_name.lower()
+    if "qwen" in model_name_lower:
+        return QwenImage
+    elif "fibo" in model_name_lower:
+        return FIBO
+    elif "z-image" in model_name_lower or "zimage" in model_name_lower:
+        return ZImageTurbo
+    elif "flux2" in model_name_lower or "flux.2" in model_name_lower:
+        return Flux2Klein
+    else:
+        return Flux1
+
+
+def get_repo_name(model_name: str, bits: int | None) -> str:
+    """Get the HuggingFace repo name for a model variant."""
+    # Extract repo name from HF path (e.g., "black-forest-labs/FLUX.1-Kontext-dev" -> "FLUX.1-Kontext-dev")
+    base_name = model_name.split("/")[-1] if "/" in model_name else model_name
+    suffix = f"-{bits}bit" if bits else ""
+    return f"{HF_ORG}/{base_name}{suffix}"
+
+
+def get_local_path(output_dir: Path, model_name: str, bits: int | None) -> Path:
+    """Get the local save path for a model variant."""
+    # Extract repo name from HF path (e.g., "black-forest-labs/FLUX.1-Kontext-dev" -> "FLUX.1-Kontext-dev")
+    base_name = model_name.split("/")[-1] if "/" in model_name else model_name
+    suffix = f"-{bits}bit" if bits else ""
+    return output_dir / f"{base_name}{suffix}"
+
+
+def copy_source_repo(
+    source_repo: str,
+    local_path: Path,
+    dry_run: bool = False,
+) -> None:
+    """Copy all files from source repo (replicating original HF structure)."""
+    print(f"\n{'=' * 60}")
+    print(f"Copying full repo from source: {source_repo}")
+    print(f"Output path: {local_path}")
+    print(f"{'=' * 60}")
+
+    if dry_run:
+        print("[DRY RUN] Would download all files from source repo")
+        return
+
+    from huggingface_hub import snapshot_download
+
+    # Download all files to our local path
+    snapshot_download(
+        repo_id=source_repo,
+        local_dir=local_path,
+    )
+
+    # Remove root-level safetensors files (flux.1-dev.safetensors, etc.)
+    # These are redundant with the component directories
+    for f in local_path.glob("*.safetensors"):
+        print(f"Removing root-level safetensors: {f.name}")
+        if not dry_run:
+            f.unlink()
+
+    print(f"Source repo copied to {local_path}")
+
+
+def load_and_save_quantized_model(
+    model_name: str,
+    bits: int,
+    output_path: Path,
+    dry_run: bool = False,
+) -> None:
+    """Load a model with quantization and save it in mflux format."""
+    print(f"\n{'=' * 60}")
+    print(f"Loading {model_name} with {bits}-bit quantization...")
+    print(f"Output path: {output_path}")
+    print(f"{'=' * 60}")
+
+    if dry_run:
+        print("[DRY RUN] Would load and save quantized model")
+        return
+
+    from mflux.models.common.config.model_config import ModelConfig
+
+    model_class = get_model_class(model_name)
+    model_config = ModelConfig.from_name(model_name=model_name, base_model=None)
+
+    model: Flux1 = model_class(
+        quantize=bits,
+        model_config=model_config,
+    )
+
+    print(f"Saving model to {output_path}...")
+    model.save_model(str(output_path))
+    print(f"Model saved successfully to {output_path}")
+
+
+def copy_source_metadata(
+    source_repo: str,
+    local_path: Path,
+    dry_run: bool = False,
+) -> None:
+    """Copy metadata files (LICENSE, README, etc.) from source repo, excluding safetensors."""
+    print(f"\n{'=' * 60}")
+    print(f"Copying metadata from source repo: {source_repo}")
+    print(f"{'=' * 60}")
+
+    if dry_run:
+        print("[DRY RUN] Would download metadata files (excluding *.safetensors)")
+        return
+
+    from huggingface_hub import snapshot_download
+
+    # Download all files except safetensors to our local path
+    snapshot_download(
+        repo_id=source_repo,
+        local_dir=local_path,
+        ignore_patterns=["*.safetensors"],
+    )
+    print(f"Metadata files copied to {local_path}")
+
+
+def upload_to_huggingface(
+    local_path: Path,
+    repo_id: str,
+    dry_run: bool = False,
+    clean_remote: bool = False,
+) -> None:
+    """Upload a saved model to HuggingFace."""
+    print(f"\n{'=' * 60}")
+    print(f"Uploading to HuggingFace: {repo_id}")
+    print(f"Local path: {local_path}")
+    print(f"Clean remote first: {clean_remote}")
+    print(f"{'=' * 60}")
+
+    if dry_run:
+        print("[DRY RUN] Would upload to HuggingFace")
+        return
+
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+
+    # Create the repo if it doesn't exist
+    print(f"Creating/verifying repo: {repo_id}")
+    api.create_repo(repo_id=repo_id, repo_type="model", exist_ok=True)
+
+    # Clean remote repo if requested (delete old mflux-format files)
+    if clean_remote:
+        print("Cleaning old mflux-format files from remote...")
+        try:
+            # Pattern for mflux numbered shards: <dir>/<number>.safetensors
+            numbered_pattern = re.compile(r".*/\d+\.safetensors$")
+
+            repo_files = api.list_repo_files(repo_id=repo_id, repo_type="model")
+            for file_path in repo_files:
+                # Delete numbered safetensors (mflux format) and mflux index files
+                if numbered_pattern.match(file_path) or file_path.endswith(
+                    "/model.safetensors.index.json"
+                ):
+                    print(f"  Deleting: {file_path}")
+                    api.delete_file(
+                        path_in_repo=file_path, repo_id=repo_id, repo_type="model"
+                    )
+        except Exception as e:
+            print(f"Warning: Could not clean remote files: {e}")
+
+    # Upload the folder
+    print("Uploading folder contents...")
+    api.upload_folder(
+        folder_path=str(local_path),
+        repo_id=repo_id,
+        repo_type="model",
+    )
+    print(f"Upload complete: https://huggingface.co/{repo_id}")
+
+
+def clean_local_files(local_path: Path, dry_run: bool = False) -> None:
+    """Remove local model files after upload."""
+    print(f"\nCleaning up: {local_path}")
+    if dry_run:
+        print("[DRY RUN] Would remove local files")
+        return
+
+    if local_path.exists():
+        shutil.rmtree(local_path)
+        print(f"Removed {local_path}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Download an mflux model, quantize it, and upload to HuggingFace.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    # Process all variants (base, 4-bit, 8-bit) for FLUX.1-Kontext-dev
+    python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev
+
+    # Only process 4-bit variant
+    python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --skip-base --skip-8bit
+
+    # Save locally without uploading
+    python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --skip-upload
+
+    # Preview what would happen
+    python tmp/quantize_and_upload.py --model black-forest-labs/FLUX.1-Kontext-dev --dry-run
+        """,
+    )
+
+    parser.add_argument(
+        "--model",
+        "-m",
+        required=True,
+        help="HuggingFace model path (e.g., black-forest-labs/FLUX.1-Kontext-dev)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("./tmp/models"),
+        help="Local directory to save models (default: ./tmp/models)",
+    )
+    parser.add_argument(
+        "--skip-base",
+        action="store_true",
+        help="Skip base model (no quantization)",
+    )
+    parser.add_argument(
+        "--skip-4bit",
+        action="store_true",
+        help="Skip 4-bit quantized model",
+    )
+    parser.add_argument(
+        "--skip-8bit",
+        action="store_true",
+        help="Skip 8-bit quantized model",
+    )
+    parser.add_argument(
+        "--skip-download",
+        action="store_true",
+        help="Skip downloading/processing, only do upload/clean operations",
+    )
+    parser.add_argument(
+        "--skip-upload",
+        action="store_true",
+        help="Only save locally, don't upload to HuggingFace",
+    )
+    parser.add_argument(
+        "--clean",
+        action="store_true",
+        help="Remove local files after upload",
+    )
+    parser.add_argument(
+        "--clean-remote",
+        action="store_true",
+        help="Delete old mflux-format files from remote repo before uploading",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print actions without executing",
+    )
+
+    args = parser.parse_args()
+
+    # Determine which variants to process
+    variants: list[int | None] = []
+    if not args.skip_base:
+        variants.append(None)  # Base model (no quantization)
+    if not args.skip_4bit:
+        variants.append(4)
+    if not args.skip_8bit:
+        variants.append(8)
+
+    if not variants:
+        print("Error: All variants skipped. Nothing to do.")
+        return 1
+
+    # Create output directory
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Model: {args.model}")
+    print(f"Output directory: {args.output_dir}")
+    print(
+        f"Variants to process: {['base' if v is None else f'{v}-bit' for v in variants]}"
+    )
+    print(f"Upload to HuggingFace: {not args.skip_upload}")
+    print(f"Clean after upload: {args.clean}")
+    if args.dry_run:
+        print("\n*** DRY RUN MODE - No actual changes will be made ***")
+
+    # Process each variant
+    for bits in variants:
+        local_path = get_local_path(args.output_dir, args.model, bits)
+        repo_id = get_repo_name(args.model, bits)
+
+        if not args.skip_download:
+            if bits is None:
+                # Base model: copy original HF repo structure (no mflux conversion)
+                copy_source_repo(
+                    source_repo=args.model,
+                    local_path=local_path,
+                    dry_run=args.dry_run,
+                )
+            else:
+                # Quantized model: load, quantize, and save with mflux
+                load_and_save_quantized_model(
+                    model_name=args.model,
+                    bits=bits,
+                    output_path=local_path,
+                    dry_run=args.dry_run,
+                )
+
+                # Copy metadata from source repo (LICENSE, README, etc.)
+                copy_source_metadata(
+                    source_repo=args.model,
+                    local_path=local_path,
+                    dry_run=args.dry_run,
+                )
+
+        # Upload
+        if not args.skip_upload:
+            upload_to_huggingface(
+                local_path=local_path,
+                repo_id=repo_id,
+                dry_run=args.dry_run,
+                clean_remote=args.clean_remote,
+            )
+
+            # Clean up if requested
+            if args.clean:
+                clean_local_files(local_path, dry_run=args.dry_run)
+
+    print("\n" + "=" * 60)
+    print("All done!")
+    print("=" * 60)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/uv.lock b/uv.lock
index e89fbd34..637d9e42 100644
--- a/uv.lock
+++ b/uv.lock
@@ -192,20 +192,14 @@ sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8
 wheels = [
     { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" },
     { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" },
-    { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" },
-    { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" },
     { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" },
     { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" },
     { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" },
     { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" },
-    { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" },
-    { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" },
     { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" },
     { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" },
     { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" },
     { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" },
-    { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" },
     { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" },
     { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" },
     { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" },
@@ -311,10 +305,8 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" },
     { url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" },
     { url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" },
-    { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" },
     { url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" },
     { url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" },
-    { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" },
     { url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" },
     { url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" },
     { url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" },
@@ -322,10 +314,8 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c5/fd/bc1daf8230eaa075184cbbf5f8cd00ba9db4fd32d63fb83da4671b72ed8a/cryptography-46.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39b6755623145ad5eff1dab323f4eae2a32a77a7abef2c5089a04a3d04366715", size = 4435078, upload-time = "2025-10-15T23:17:23.042Z" },
     { url = "https://files.pythonhosted.org/packages/82/98/d3bd5407ce4c60017f8ff9e63ffee4200ab3e23fe05b765cab805a7db008/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:db391fa7c66df6762ee3f00c95a89e6d428f4d60e7abc8328f4fe155b5ac6e54", size = 4293460, upload-time = "2025-10-15T23:17:24.885Z" },
     { url = "https://files.pythonhosted.org/packages/26/e9/e23e7900983c2b8af7a08098db406cf989d7f09caea7897e347598d4cd5b/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:78a97cf6a8839a48c49271cdcbd5cf37ca2c1d6b7fdd86cc864f302b5e9bf459", size = 3995237, upload-time = "2025-10-15T23:17:26.449Z" },
-    { url = "https://files.pythonhosted.org/packages/91/15/af68c509d4a138cfe299d0d7ddb14afba15233223ebd933b4bbdbc7155d3/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:dfb781ff7eaa91a6f7fd41776ec37c5853c795d3b358d4896fdbb5df168af422", size = 4967344, upload-time = "2025-10-15T23:17:28.06Z" },
     { url = "https://files.pythonhosted.org/packages/ca/e3/8643d077c53868b681af077edf6b3cb58288b5423610f21c62aadcbe99f4/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6f61efb26e76c45c4a227835ddeae96d83624fb0d29eb5df5b96e14ed1a0afb7", size = 4466564, upload-time = "2025-10-15T23:17:29.665Z" },
     { url = "https://files.pythonhosted.org/packages/0e/43/c1e8726fa59c236ff477ff2b5dc071e54b21e5a1e51aa2cee1676f1c986f/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:23b1a8f26e43f47ceb6d6a43115f33a5a37d57df4ea0ca295b780ae8546e8044", size = 4292415, upload-time = "2025-10-15T23:17:31.686Z" },
-    { url = "https://files.pythonhosted.org/packages/42/f9/2f8fefdb1aee8a8e3256a0568cffc4e6d517b256a2fe97a029b3f1b9fe7e/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:b419ae593c86b87014b9be7396b385491ad7f320bde96826d0dd174459e54665", size = 4931457, upload-time = "2025-10-15T23:17:33.478Z" },
     { url = "https://files.pythonhosted.org/packages/79/30/9b54127a9a778ccd6d27c3da7563e9f2d341826075ceab89ae3b41bf5be2/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:50fc3343ac490c6b08c0cf0d704e881d0d660be923fd3076db3e932007e726e3", size = 4466074, upload-time = "2025-10-15T23:17:35.158Z" },
     { url = "https://files.pythonhosted.org/packages/ac/68/b4f4a10928e26c941b1b6a179143af9f4d27d88fe84a6a3c53592d2e76bf/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:22d7e97932f511d6b0b04f2bfd818d73dcd5928db509460aaf48384778eb6d20", size = 4420569, upload-time = "2025-10-15T23:17:37.188Z" },
     { url = "https://files.pythonhosted.org/packages/a3/49/3746dab4c0d1979888f125226357d3262a6dd40e114ac29e3d2abdf1ec55/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d55f3dffadd674514ad19451161118fd010988540cee43d8bc20675e775925de", size = 4681941, upload-time = "2025-10-15T23:17:39.236Z" },
@@ -333,10 +323,8 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" },
     { url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" },
     { url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" },
-    { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" },
     { url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" },
     { url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" },
-    { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" },
     { url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" },
     { url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" },
     { url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" },
@@ -412,7 +400,7 @@ requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
-    { name = "mflux", specifier = "==0.15.4" },
+    { name = "mflux", specifier = "==0.15.5" },
     { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.5" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.5" },
     { name = "mlx-lm", specifier = "==0.30.6" },
@@ -987,7 +975,7 @@ wheels = [
 
 [[package]]
 name = "mflux"
-version = "0.15.4"
+version = "0.15.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1013,9 +1001,9 @@ dependencies = [
     { name = "twine", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "urllib3", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a6/f8/95322db7a865e4df6bad108b1c99aa7fbe211aac3f298f3ad696c2744a39/mflux-0.15.4.tar.gz", hash = "sha256:138e1aedae86e13eafeb8faec017945fcdcca42c3234daabcd81a83c9a202ace", size = 741228, upload-time = "2026-01-20T15:39:26.807Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/35/8e/f20de51bf9dc0a986535d9a825db4ae314163421b3d3ddaa90a2b959b9fd/mflux-0.15.5.tar.gz", hash = "sha256:9a3372bd64d51c4caff4ff9e7d7d698bea5833242fd849c59cbb0c92f7d7aa3b", size = 743700, upload-time = "2026-01-26T12:41:45.272Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/8e/be/81cf4ce2d1933b9b210c028a05ac95e958008c0d43e377a5f2757b7f2d4d/mflux-0.15.4-py3-none-any.whl", hash = "sha256:f04d9b1d7c5cd67880f483ab29fb2097648a25459eef9c5ee6480fad46de5e82", size = 987644, upload-time = "2026-01-20T15:39:24.817Z" },
+    { url = "https://files.pythonhosted.org/packages/ac/bb/ef936eae2ae78a47cd92ddffc18fc06ad3fd5f438a0915fb62d8bb9508ec/mflux-0.15.5-py3-none-any.whl", hash = "sha256:c94891d4a518047a818863bb099c755e93af90c524ced358baf5b31502c09e82", size = 990939, upload-time = "2026-01-26T12:41:42.898Z" },
 ]
 
 [[package]]

← c8d3154f More image dimensions (#1395)  ·  back to Exo  ·  Fix GLM4Moe Tensor Sharding (#1411) 5455a97a →