← 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
M .gitignoreA .mlx_typings/mflux/models/flux/variants/kontext/__init__.pyiA .mlx_typings/mflux/models/flux/variants/kontext/flux_kontext.pyiA .mlx_typings/mflux/models/flux/variants/kontext/kontext_util.pyiM pyproject.tomlA resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-4bit.tomlA resources/image_model_cards/exolabs--FLUX.1-Kontext-dev-8bit.tomlA resources/image_model_cards/exolabs--FLUX.1-Kontext-dev.tomlM src/exo/worker/engines/image/models/__init__.pyM src/exo/worker/engines/image/models/base.pyM src/exo/worker/engines/image/models/flux/__init__.pyM src/exo/worker/engines/image/models/flux/adapter.pyM src/exo/worker/engines/image/models/flux/config.pyA src/exo/worker/engines/image/models/flux/kontext_adapter.pyM src/exo/worker/engines/image/models/qwen/adapter.pyM src/exo/worker/engines/image/models/qwen/edit_adapter.pyM src/exo/worker/engines/image/pipeline/runner.pyA tmp/quantize_and_upload.pyM uv.lock
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 →