monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Bria FIBO support (#279)

Browse Source 
Co-authored-by: Filip Strand <filip@host-022.local>
This commit is contained in:
Filip Strand 2025-11-27 13:01:55 +01:00 committed by GitHub
parent c3e72bd73c
commit c2c860f4de
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
187 changed files with 8364 additions and 1756 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

343
README.md
View File

@ -4,7 +4,7 @@
### About
Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) and [Qwen Image](https://github.com/QwenLM/Qwen-Image) models locally on your Mac!
Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux), [Qwen Image](https://github.com/QwenLM/Qwen-Image) and [FIBO](https://huggingface.co/briaai/FIBO) models locally on your Mac!
### Table of contents
@ -23,6 +23,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) and [Qwen Image](h
- [🦙 Qwen Models](#-qwen-models)
* [🖼️ Qwen Image](#%EF%B8%8F-qwen-image)
* [✏️ Qwen Image Edit](#%EF%B8%8F-qwen-image-edit)
- [🌀 FIBO](#-fibo)
- [🔌 LoRA](#-lora)
- [🎭 In-Context Generation](#-in-context-generation)
* [📸 Kontext](#-kontext)
@ -50,7 +51,7 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) and [Qwen Image](h
### Philosophy
MFLUX is a line-by-line port of the FLUX and Qwen implementations in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) and [Huggingface Transformers](https://github.com/huggingface/transformers) libraries to [Apple MLX](https://github.com/ml-explore/mlx).
MFLUX is a line-by-line port of the FLUX, Qwen and Bria models implementations in the [Huggingface Diffusers](https://github.com/huggingface/diffusers) and [Huggingface Transformers](https://github.com/huggingface/transformers) libraries to [Apple MLX](https://github.com/ml-explore/mlx).
MFLUX is purposefully kept minimal and explicit - Network architectures are hardcoded and no config files are used
except for the tokenizers. The aim is to have a tiny codebase with the single purpose of expressing these models
(thereby avoiding too many abstractions). While MFLUX priorities readability over generality and performance, [it can still be quite fast](#%EF%B8%8F-image-generation-speed-updated), [and even faster quantized](#%EF%B8%8F-quantization).
@ -1087,6 +1088,344 @@ mflux-generate-qwen-edit \
---
### 🌀 FIBO
MFLUX supports [FIBO](https://huggingface.co/briaai/FIBO) from [Bria.ai](https://bria.ai), the first open-source JSON-native text-to-image model trained on long structured captions. FIBO delivers high image quality, strong prompt adherence, and professional-grade control—trained exclusively on licensed data. ([Technical Paper](https://arxiv.org/abs/2511.06876))
![FIBO Example](src/mflux/assets/fibo_example.jpg)
FIBO is an 8B-parameter DiT-based, flow-matching model using **SmolLM3-3B** as the text encoder with a novel **DimFusion** conditioning architecture for efficient long-caption training, and **Wan 2.2** as the VAE. The VLM-assisted prompting uses a fine-tuned **Qwen3-VL** to expand short user intents, fill in missing details, and extract/edit structured prompts from images.
Most text-to-image models excel at imagination—but not control. FIBO is trained on structured JSON captions up to 1,000+ words, enabling precise, reproducible control over lighting, composition, color, and camera settings. The structured captions foster native disentanglement, allowing targeted, iterative refinement without prompt drift.
#### Key Features
- **VLM-guided JSON-native prompting**: Transform short prompts into structured schemas with 1,000+ words (lighting, camera, composition, DoF)
- **Disentangled control**: Tweak a single attribute (e.g., camera angle) without breaking the scene
- **Strong prompt adherence**: High alignment on PRISM-style evaluations
- **Enterprise-grade**: 100% licensed data with governance, repeatability, and legal clarity
#### The three modes: ✨ Generate, 🔧 Refine, and 💡 Inspire
**✨ Generate**: While the actual prompt input to FIBO is a structured JSON file, the generate command provides an interface to input pure text prompts. These are then expanded into structured JSON prompts using FIBO's Vision-Language Model (VLM) before being passed to the diffusion model for image generation. For example, the following prompt produces one of the images above:
```sh
mflux-generate-fibo \
--prompt "Three cartoon animal chefs in a colorful bakery kitchen, Pixar style: a bunny with floppy ears wearing a tall white chef hat and pink apron holding a chocolate cake on the left, a raccoon with a striped tail wearing blue oven mitts and a yellow bandana frosting cupcakes in the center, a penguin wearing a red bowtie and checkered apron carrying a tray of golden croissants on the right, warm kitchen lighting with flour dust in air" \
--width 1200 \
--height 540 \
--steps 20 \
--guidance 4.0 \
--seed 42 \
--output animal_bakers.png
```
This command will output both the generated image (`animal_bakers.png`) and a JSON prompt file (`animal_bakers.json`) containing the expanded structured prompt used for generation.
When the input prompt is pure text, it will be processed through FIBO's VLM to create the structured JSON prompt automatically.
Conversely, if a JSON prompt file is provided, it will be used directly for image generation, thus bypassing the VLM step and giving you full control over the prompt structure.
Another way to call the model directly is to provide a JSON prompt file as input:
```sh
mflux-generate-fibo \
--prompt-file animal_bakers.json \
--width 1200 \
--height 540 \
--steps 20 \
--guidance 4.0 \
--seed 42 \
--output animal_bakers.png
```
A point worth emphasizing is that when working with a JSON prompt file, the user can use whatever tool they prefer to edit it and is not forced to use the built in FIBO-VLM. Other good alternatives are [coding](https://cursor.com/agents) [agents](https://www.claude.com/product/claude-code), other [LLMs](https://github.com/ml-explore/mlx-lm)/[VLMs](https://github.com/Blaizzy/mlx-vlm) etc.
<details>
<summary><strong>Click to expand the JSON prompt file used (animal_bakers.json)</strong></summary>
```json
{
"short_description": "Three cartoon animal chefs are in a bakery kitchen, each holding a culinary creation. A bunny chef on the left presents a chocolate cake, a raccoon chef in the center is frosting cupcakes, and a penguin chef on the right carries a tray of croissants. The kitchen is brightly lit with warm tones, and flour dusts the air, creating a lively and cheerful baking atmosphere.",
"objects": [
{
"description": "A cartoon bunny wearing a white chef's hat and a pink apron, holding a chocolate cake with white frosting and cherries.",
"location": "left foreground",
"relationship": "The bunny chef is presenting the chocolate cake.",
"relative_size": "medium",
"shape_and_color": "Rounded bunny shape, white hat, pink apron, brown cake, white frosting, red cherries.",
"texture": "smooth",
"appearance_details": "Floppy ears, rosy cheeks, smiling expression.",
"pose": "Standing upright, holding the cake with both hands.",
"expression": "Joyful and proud.",
"clothing": "White chef's hat, pink apron.",
"action": "Holding and presenting a cake.",
"gender": "female",
"skin_tone_and_texture": "White fur, smooth texture.",
"orientation": "Upright, facing forward."
},
{
"description": "A cartoon raccoon wearing blue oven mitts and a yellow bandana, actively frosting cupcakes with pink and yellow frosting.",
"location": "center midground",
"relationship": "The raccoon chef is in the process of frosting cupcakes.",
"relative_size": "medium",
"shape_and_color": "Distinct raccoon shape, blue mitts, yellow bandana, pink and yellow frosting, brown cupcakes.",
"texture": "smooth",
"appearance_details": "Striped tail, bushy fur, focused expression.",
"pose": "Leaning forward slightly, hands busy frosting.",
"expression": "Concentrated and happy.",
"clothing": "Blue oven mitts, yellow bandana.",
"action": "Frosting cupcakes.",
"gender": "male",
"skin_tone_and_texture": "Brown and black fur, smooth texture.",
"orientation": "Upright, slightly angled."
},
{
"description": "A cartoon penguin wearing a red bowtie and a checkered apron, carrying a tray of golden croissants.",
"location": "right foreground",
"relationship": "The penguin chef is carrying a tray of freshly baked croissants.",
"relative_size": "medium",
"shape_and_color": "Classic penguin shape, red bowtie, red and white apron, golden croissants.",
"texture": "smooth",
"appearance_details": "Black and white body, orange beak and feet, smiling expression.",
"pose": "Standing upright, holding the tray with both hands.",
"expression": "Cheerful and friendly.",
"clothing": "Red bowtie, checkered apron.",
"action": "Carrying a tray of croissants.",
"gender": "male",
"skin_tone_and_texture": "Feathered texture, smooth appearance.",
"orientation": "Upright, facing forward."
},
{
"description": "A chocolate cake with white frosting and two red cherries on top.",
"location": "left foreground",
"relationship": "Held by the bunny chef.",
"relative_size": "medium",
"shape_and_color": "Round cake, dark brown, white frosting, red cherries.",
"texture": "smooth frosting, slightly textured cake",
"appearance_details": "Decorated with cherries.",
"number_of_objects": 1,
"orientation": "Horizontal"
},
{
"description": "A tray filled with golden-brown croissants.",
"location": "right foreground",
"relationship": "Carried by the penguin chef.",
"relative_size": "medium",
"shape_and_color": "Elongated, crescent-shaped croissants, golden brown.",
"texture": "flaky, slightly crisp exterior",
"appearance_details": "Arranged neatly on a metal tray.",
"number_of_objects": 1,
"orientation": "Horizontal"
}
],
"background_setting": "A brightly lit bakery kitchen with wooden countertops, shelves stocked with baking ingredients and equipment, and a window in the background. Flour is lightly dusted in the air.",
"lighting": {
"conditions": "warm indoor lighting",
"direction": "front-lit and side-lit",
"shadows": "soft, diffused shadows"
},
"aesthetics": {
"composition": "centered composition with the three chefs forming a horizontal line",
"color_scheme": "warm and cheerful, with dominant browns, yellows, pinks, and whites",
"mood_atmosphere": "joyful, friendly, and inviting",
"aesthetic_score": "very high",
"preference_score": "very high"
},
"photographic_characteristics": {
"depth_of_field": "shallow",
"focus": "sharp focus on the chefs and their creations",
"camera_angle": "eye-level",
"lens_focal_length": "standard lens"
},
"style_medium": "digital illustration",
"text_render": [],
"context": "This image is a charming illustration, likely for a children's book, educational material, or a bakery-themed advertisement, designed to evoke feelings of happiness and the joy of baking.",
"artistic_style": "Pixar-style animation"
}
```
Note: This JSON prompt was generated using FIBO's "Generate" mode from a short text description. Note the strong alignment between the JSON prompt and the image!
</details>
**🔧 Refine**: While the JSON prompt can be edited manually, it can be quite complex and inconvenient to modify directly. The refinement mode helps to solve this issue by also expanding a simple user instruction in order to tweak specific attributes. The VLM processes these instructions and updates the JSON prompt accordingly before generating new images.
![FIBO Refine Example](src/mflux/assets/fibo_refine_example.jpg)
Assuming we already have a previous prompt file, like `owl_brown.json`, we can refine this prompt to change the owl's color and add some accessories:
```sh
mflux-refine-fibo \
--prompt-file owl_brown.json \
--instructions "Make the owl white instead of brown, and add round glasses and a black scarf. Keep everything else exactly the same - the same forest background, moonlight lighting, composition, and overall whimsical atmosphere." \
--output owl_white.json
```
<details>
<summary><strong>Click to expand the refined JSON prompt file (owl_white.json)</strong></summary>
```json
{
"short_description": "A hyper-detailed, ultra-fluffy owl sitting in the trees at night, looking directly at the camera with wide, adorable, expressive eyes. Its feathers are soft and voluminous, catching the cool moonlight with subtle silver highlights. The owl's gaze is curious and full of charm, giving it a whimsical, storybook-like personality. It is wearing round glasses and a black scarf.",
"objects": [
{
"description": "An adorable, fluffy owl with large, expressive eyes and soft, voluminous feathers. Its plumage is white, with subtle silver highlights from the moonlight. It is wearing round glasses and a black scarf.",
"location": "center",
"relationship": "The owl is the sole subject, perched comfortably within its environment.",
"relative_size": "large within frame",
"shape_and_color": "Round head, large eyes, bulky body, predominantly white with silver accents.",
"texture": "Extremely soft, fluffy, and detailed feathers, giving a plush toy-like appearance.",
"appearance_details": "The eyes are wide, dark, and reflective, conveying a sense of wonder and curiosity. The beak is small and light-colored, almost hidden by the feathers. Subtle silver highlights catch the moonlight on its feathers. It has round glasses on its nose and a black scarf around its neck.",
"orientation": "upright, facing forward"
}
],
"background_setting": "A dark, nocturnal forest setting with blurred trees and foliage, illuminated by a soft, cool moonlight. The background is out of focus, emphasizing the owl.",
"lighting": {
"conditions": "moonlight",
"direction": "backlit and side-lit from the left",
"shadows": "soft, diffused shadows on the right side of the owl and within the background foliage, indicating a single light source."
},
"aesthetics": {
"composition": "centered, portrait composition",
"color_scheme": "cool blues and silvers from the moonlight contrasting with white of the owl and forest.",
"mood_atmosphere": "mysterious, enchanting, whimsical, and serene.",
"aesthetic_score": "very high",
"preference_score": "very high"
},
"photographic_characteristics": {
"depth_of_field": "shallow",
"focus": "sharp focus on the owl's face and eyes, with a soft blur in the background.",
"camera_angle": "eye-level",
"lens_focal_length": "portrait lens (e.g., 50mm-85mm)"
},
"style_medium": "digital illustration",
"text_render": [],
"context": "A whimsical character illustration, possibly for a children's book, animated film, or fantasy art collection.",
"artistic_style": "fantasy, illustrative, detailed"
}
```
Note: This JSON prompt was refined from the original `owl_brown.json` by changing the owl's color to white and adding round glasses and a black scarf, while preserving the forest background, moonlight lighting, and whimsical atmosphere.
</details>
Finally, generate the refined white owl image using the updated JSON prompt:
```sh
mflux-generate-fibo \
--prompt-file owl_white.json \
--width 1024 \
--height 560 \
--steps 20 \
--guidance 4.0 \
--seed 42 \
--quantize 4 \
--output owl_white.png
```
The refine command reads the existing JSON prompt, applies the refinement instructions to change the owl's color and add accessories, and outputs a new refined JSON file. This refined prompt is then used to generate a new image that reflects the requested changes while maintaining the overall scene and composition.
It is worth noting that refine does not work the same way as other editing techniques like Flux Kontext or Qwen Image Edit. Instead of modifying an existing image, it modifies the underlying **structured prompt** to produce a new image that reflects the requested changes while maintaining the overall scene and composition.
**💡 Inspire**: Provide an image instead of text. FIBO's vision-language model extracts a detailed, structured prompt, blends it with your creative intent, and produces related images—ideal for inspiration without overreliance on the original.
![FIBO Inspire Example](src/mflux/assets/fibo_inspire_example.jpg)
Starting from an image, you can extract a structured JSON prompt that captures its visual characteristics. For example, using a [blue and brown bird on brown tree trunk](https://unsplash.com/photos/blue-and-brown-bird-on-brown-tree-trunk-DPXytK8Z59Y), we can extract a detailed prompt:
```sh
mflux-inspire-fibo \
--image-path bird.jpg \
--prompt "blue and brown bird on brown tree trunk" \
--output bird_inspired.json \
--seed 42
```
This command analyzes the image and generates a structured JSON prompt file (`bird_inspired.json`) that describes the visual elements, composition, lighting, and style. You can then use this JSON prompt to generate new images with similar characteristics:
```sh
mflux-generate-fibo \
--prompt-file bird_inspired.json \
--width 1024 \
--height 672 \
--steps 20 \
--guidance 4.0 \
--seed 42 \
-q 8 \
--output bird_inspired.png
```
<details>
<summary><strong>Click to expand the JSON prompt file used (bird_inspired.json)</strong></summary>
```json
{
"short_description": "A vibrant blue and brown kingfisher is perched on a weathered tree trunk, facing left. The bird's iridescent plumage is detailed, with striking orange and white markings on its chest and throat. Its long, sharp black beak is prominent. The background is a soft, out-of-focus gradient of warm yellow and green, creating a natural and serene environment. The lighting highlights the textures of the bird's feathers and the rough bark of the trunk.",
"objects": [
{
"description": "A male kingfisher with striking iridescent blue plumage on its back and wings, a rich orange and white chest, and a black beak. It has a small, dark eye and a red-orange patch on its face.",
"location": "center",
"relationship": "perched on the tree trunk",
"relative_size": "medium within frame",
"shape_and_color": "Bird shape, predominantly blue, orange, and white.",
"texture": "Feathers appear smooth and slightly glossy.",
"appearance_details": "The beak is long, thin, and black. The orange patch on its face is distinct.",
"number_of_objects": 1,
"pose": "Standing upright on its legs, head turned to the left.",
"expression": "Alert and focused.",
"action": "Perched, observing its surroundings.",
"gender": "male",
"orientation": "Facing left, upright"
},
{
"description": "A section of a weathered, rough tree trunk, providing a perch for the kingfisher. It has a natural, organic shape with visible bark texture.",
"location": "bottom-center foreground",
"relationship": "supports the kingfisher",
"relative_size": "medium",
"shape_and_color": "Irregular cylindrical shape, brown and grey tones.",
"texture": "Rough, gnarled bark texture.",
"appearance_details": "Some smaller branches or knots are visible on the trunk.",
"number_of_objects": 1,
"orientation": "Horizontal, lying on its side"
}
],
"background_setting": "A soft, blurred background with a warm gradient transitioning from a light yellow at the top to a muted green at the bottom. This creates a natural, out-of-focus environment, likely foliage or sky.",
"lighting": {
"conditions": "natural daylight",
"direction": "side-lit from the right",
"shadows": "soft shadows, particularly on the left side of the bird and the trunk"
},
"aesthetics": {
"composition": "rule of thirds, with the bird positioned slightly off-center",
"color_scheme": "vibrant blues and oranges contrasting with the soft yellow and green background",
"mood_atmosphere": "serene, natural, captivating",
"aesthetic_score": "very high",
"preference_score": "very high"
},
"photographic_characteristics": {
"depth_of_field": "shallow, with a strong bokeh effect in the background",
"focus": "sharp focus on the kingfisher",
"camera_angle": "eye-level",
"lens_focal_length": "telephoto lens"
},
"style_medium": "photograph",
"text_render": [],
"context": "This image is a wildlife photograph, likely intended for nature magazines, educational materials, or as a decorative print for nature enthusiasts.",
"artistic_style": "photorealistic, detailed"
}
```
Note: This JSON prompt was extracted from the input image using FIBO's VLM, capturing the visual characteristics of a kingfisher bird perched on a tree trunk with a bokeh background.
</details>
The inspire command is particularly useful when you want to:
- Extract the visual style and composition from a reference image
- Create variations of an existing image while maintaining its core characteristics
- Understand how FIBO interprets and structures visual information
- Blend your creative intent (via the optional `--prompt` parameter) with the visual content of the image
Note: The optional `--prompt` parameter allows you to guide the VLM's interpretation of the image. For example, you might use `--prompt "futuristic cityscape"` to influence how the image is analyzed and structured.
⚠️ *Note: FIBO requires downloading the `briaai/FIBO` model weights (~24GB) and the `briaai/FIBO-vlm` vision-language model (~8GB), totaling ~32GB for the full model, or use quantization for smaller sizes.*
---
### 🔌 LoRA
MFLUX support loading trained [LoRA](https://huggingface.co/docs/diffusers/en/training/lora) adapters (actual training support is coming).

7
pyproject.toml
View File

@ -17,7 +17,7 @@ source-exclude = [
[project]
name = "mflux"
version = "0.11.1"
version = "0.12.0.dev0"
description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
readme = "README.md"
keywords = ["diffusers", "flux", "mlx"]
@ -50,7 +50,7 @@ dependencies = [
"torch>=2.3.1,<3.0; python_version<'3.13'",
"torch>=2.8.0,<3.0; python_version>='3.13'",
"tqdm>=4.66.5,<5.0",
"transformers>=4.55.0,<5.0",
"transformers>=4.57,<5.0",
"twine>=6.1.0,<7.0",
]
classifiers = [
@ -86,6 +86,9 @@ mflux-generate-redux = "mflux.generate_redux:main"
mflux-generate-kontext = "mflux.generate_kontext:main"
mflux-generate-qwen = "mflux.generate_qwen:main"
mflux-generate-qwen-edit = "mflux.generate_qwen_edit:main"
mflux-generate-fibo = "mflux.generate_fibo:main"
mflux-refine-fibo = "mflux.refine_fibo:main"
mflux-inspire-fibo = "mflux.inspire_fibo:main"
mflux-concept = "mflux.concept:main"
mflux-concept-from-image = "mflux.concept_from_image:main"
mflux-save = "mflux.save:main"

6
src/mflux/__init__.py Normal file
View File

@ -0,0 +1,6 @@
import os
# Set TOKENIZERS_PARALLELISM to avoid fork warning
# This must be set before any tokenizers are imported/used
if "TOKENIZERS_PARALLELISM" not in os.environ:
os.environ["TOKENIZERS_PARALLELISM"] = "false"

BIN
src/mflux/assets/fibo_example.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.7 MiB

BIN
src/mflux/assets/fibo_inspire_example.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 712 KiB

BIN
src/mflux/assets/fibo_refine_example.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 768 KiB

2
src/mflux/callbacks/instances/battery_saver.py
View File

@ -4,7 +4,7 @@ import re
import subprocess
from mflux.callbacks.callback import BeforeLoopCallback
from mflux.error.exceptions import StopImageGenerationException
from mflux.utils.exceptions import StopImageGenerationException
PMSET_AC_POWER_STATUS = "Now drawing from 'AC Power'"
PMSET_BATT_STATUS_PATTERN = r"InternalBattery-.+?(\d+)%"

2
src/mflux/callbacks/instances/canny_saver.py
View File

@ -6,7 +6,7 @@ import PIL.Image
from mflux.callbacks.callback import BeforeLoopCallback
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.image_util import ImageUtil
from mflux.utils.image_util import ImageUtil
class CannyImageSaver(BeforeLoopCallback):

2
src/mflux/callbacks/instances/depth_saver.py
View File

@ -6,7 +6,7 @@ import PIL.Image
from mflux.callbacks.callback import BeforeLoopCallback
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.image_util import ImageUtil
from mflux.utils.image_util import ImageUtil
class DepthImageSaver(BeforeLoopCallback):

4
src/mflux/callbacks/instances/stepwise_handler.py
View File

@ -6,8 +6,8 @@ import tqdm
from mflux.callbacks.callback import BeforeLoopCallback, InLoopCallback, InterruptCallback
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.image_util import ImageUtil
from mflux.utils.array_util import ArrayUtil
from mflux.utils.image_util import ImageUtil
class StepwiseHandler(BeforeLoopCallback, InLoopCallback, InterruptCallback):

6
src/mflux/concept.py
View File

@ -1,11 +1,11 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.concept_attention.flux_concept import Flux1Concept
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -41,7 +41,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
concept=args.concept,
heatmap_timesteps=args.heatmap_timesteps,
heatmap_layer_indices=args.heatmap_layer_indices,

6
src/mflux/concept_from_image.py
View File

@ -1,11 +1,11 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.concept_attention.flux_concept_from_image import Flux1ConceptFromImage
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -41,7 +41,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
concept=args.concept,
image_path=str(args.input_image_path),
heatmap_timesteps=args.heatmap_timesteps,

19
src/mflux/config/model_config.py
View File

@ -1,7 +1,7 @@
from functools import lru_cache
from typing import Literal
from mflux.error.error import InvalidBaseModel, ModelConfigError
from mflux.utils.exceptions import InvalidBaseModel, ModelConfigError
class ModelConfig:
@ -94,6 +94,11 @@ class ModelConfig:
def qwen_image_edit() -> "ModelConfig":
return AVAILABLE_MODELS["qwen-image-edit"]
@staticmethod
@lru_cache
def fibo() -> "ModelConfig":
return AVAILABLE_MODELS["fibo"]
def x_embedder_input_dim(self) -> int:
if "Fill" in self.model_name:
return 384
@ -319,4 +324,16 @@ AVAILABLE_MODELS = {
requires_sigma_shift=None,
priority=12,
),
"fibo": ModelConfig(
aliases=["fibo"],
model_name="briaai/FIBO",
base_model=None,
controlnet_model=None,
custom_transformer_model=None,
num_train_steps=1000,
max_sequence_length=512,
supports_guidance=True,
requires_sigma_shift=False,
priority=13,
),
}

4
src/mflux/config/runtime_config.py
View File

@ -5,8 +5,8 @@ import mlx.core as mx
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.schedulers import SCHEDULER_REGISTRY, try_import_external_scheduler
from mflux.schedulers.linear_scheduler import LinearScheduler
from mflux.models.common.schedulers import SCHEDULER_REGISTRY, try_import_external_scheduler
from mflux.models.common.schedulers.linear_scheduler import LinearScheduler
logger = logging.getLogger(__name__)

6
src/mflux/error/error.py
View File

@ -1,6 +0,0 @@
class ModelConfigError(ValueError):
"""User error in model config."""
class InvalidBaseModel(ModelConfigError):
"""Invalid base model, cannot infer model properties."""

8
src/mflux/generate.py
View File

@ -1,11 +1,11 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.txt2img.flux import Flux1
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_negative_prompt, get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -40,8 +40,8 @@ def main():
# 3. Generate an image for each seed value
image = model.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
negative_prompt=get_effective_negative_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
negative_prompt=PromptUtils.get_effective_negative_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

6
src/mflux/generate_controlnet.py
View File

@ -1,11 +1,11 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.controlnet.flux_controlnet import Flux1Controlnet
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -40,7 +40,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
controlnet_image_path=args.controlnet_image_path,
config=Config(
num_inference_steps=args.steps,

6
src/mflux/generate_depth.py
View File

@ -1,10 +1,10 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.depth.flux_depth import Flux1Depth
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -38,7 +38,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

87
src/mflux/generate_fibo.py Normal file
View File

@ -0,0 +1,87 @@
import gc
import json
import mlx.core as mx
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.models.fibo.variants.txt2img.fibo import FIBO
from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
# 0. Parse command line arguments
parser = CommandLineParser(description="Generate an image using FIBO model.")
parser.add_general_arguments()
parser.add_model_arguments(require_model_arg=False)
parser.add_lora_arguments()
parser.add_image_generator_arguments(supports_metadata_config=True)
parser.add_image_to_image_arguments(required=False)
parser.add_output_arguments()
args = parser.parse_args()
# 0. Set default guidance value if not provided by user
if args.guidance is None:
args.guidance = ui_defaults.GUIDANCE_SCALE
json_prompt = _get_json_prompt(args)
# 1. Load the FIBO model
fibo = FIBO(
model_config=ModelConfig.fibo(),
quantize=args.quantize,
local_path=args.path,
)
# 2. Register callbacks
memory_saver = CallbackManager.register_callbacks(args=args, model=fibo)
try:
for seed in args.seed:
# 3. Generate an image for each seed value
image = fibo.generate_image(
seed=seed,
prompt=json_prompt,
negative_prompt=PromptUtils.get_effective_negative_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,
width=args.width,
guidance=args.guidance,
image_path=args.image_path,
image_strength=args.image_strength,
scheduler="flow_match_euler_discrete",
),
)
# 4. Save the image
image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
except (StopImageGenerationException, PromptFileReadError) as exc:
print(exc)
finally:
if memory_saver:
print(memory_saver.memory_stats())
def _get_json_prompt(args):
prompt = PromptUtils.get_effective_prompt(args)
try:
json.loads(prompt)
json_prompt = prompt
except json.JSONDecodeError:
vlm = FiboVLM()
json_prompt = vlm.generate(prompt=prompt, seed=42)
del vlm
gc.collect()
mx.clear_cache()
return json_prompt
if __name__ == "__main__":
main()

6
src/mflux/generate_fill.py
View File

@ -1,10 +1,10 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.fill.flux_fill import Flux1Fill
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -38,7 +38,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

6
src/mflux/generate_in_context_catvton.py
View File

@ -3,11 +3,11 @@ from pathlib import Path
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -51,7 +51,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
left_image_path=args.garment_image,
right_image_path=args.person_image,
config=Config(

6
src/mflux/generate_in_context_dev.py
View File

@ -3,12 +3,12 @@ from pathlib import Path
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.in_context.flux_in_context_dev import Flux1InContextDev
from mflux.models.flux.variants.in_context.utils.in_context_loras import LORA_REPO_ID, get_lora_filename
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -49,7 +49,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

6
src/mflux/generate_in_context_edit.py
View File

@ -5,12 +5,12 @@ from PIL import Image
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.in_context.flux_in_context_fill import Flux1InContextFill
from mflux.models.flux.variants.in_context.utils.in_context_loras import prepare_ic_edit_loras
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -82,7 +82,7 @@ def _get_effective_ic_edit_prompt(args):
if hasattr(args, "instruction") and args.instruction:
return f"A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but {args.instruction}"
else:
return get_effective_prompt(args)
return PromptUtils.get_effective_prompt(args)
def _resize_for_ic_edit_optimal_width(args):

6
src/mflux/generate_kontext.py
View File

@ -2,11 +2,11 @@ from pathlib import Path
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -40,7 +40,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

8
src/mflux/generate_qwen.py
View File

@ -1,11 +1,11 @@
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.qwen.variants.txt2img.qwen_image import QwenImage
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_negative_prompt, get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -40,8 +40,8 @@ def main():
# 3. Generate an image for each seed value
image = qwen.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
negative_prompt=get_effective_negative_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
negative_prompt=PromptUtils.get_effective_negative_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

8
src/mflux/generate_qwen_edit.py
View File

@ -2,11 +2,11 @@ from pathlib import Path
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.qwen.variants.edit.qwen_image_edit import QwenImageEdit
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_negative_prompt, get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -53,7 +53,7 @@ def main():
# 4. Generate an image for each seed value
image = qwen.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,
@ -61,7 +61,7 @@ def main():
guidance=args.guidance,
image_path=config_image_path,
),
negative_prompt=get_effective_negative_prompt(args),
negative_prompt=PromptUtils.get_effective_negative_prompt(args),
image_paths=image_paths,
)

6
src/mflux/generate_redux.py
View File

@ -3,11 +3,11 @@ from pathlib import Path
from mflux.callbacks.callback_manager import CallbackManager
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
from mflux.models.flux.variants.redux.flux_redux import Flux1Redux
from mflux.ui import defaults as ui_defaults
from mflux.ui.cli.parsers import CommandLineParser
from mflux.ui.prompt_utils import get_effective_prompt
from mflux.ui.prompt_utils import PromptUtils
from mflux.utils.exceptions import PromptFileReadError, StopImageGenerationException
def main():
@ -48,7 +48,7 @@ def main():
# 3. Generate an image for each seed value
image = flux.generate_image(
seed=seed,
prompt=get_effective_prompt(args),
prompt=PromptUtils.get_effective_prompt(args),
config=Config(
num_inference_steps=args.steps,
height=args.height,

3
src/mflux/info.py
View File

@ -4,8 +4,8 @@ import sys
from datetime import datetime
from pathlib import Path
from mflux.post_processing.metadata_reader import MetadataReader
from mflux.ui.cli.parsers import CommandLineParser
from mflux.utils.metadata_reader import MetadataReader
def format_metadata(metadata: dict) -> str:
@ -121,4 +121,3 @@ def main():
if __name__ == "__main__":
main()

68
src/mflux/inspire_fibo.py Normal file
View File

@ -0,0 +1,68 @@
import json
from pathlib import Path
from PIL import Image
from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM
from mflux.ui.cli.parsers import CommandLineParser
from mflux.utils.exceptions import PromptFileReadError
def main():
# 0. Parse command line arguments
parser = CommandLineParser(description="Generate FIBO JSON prompts from images using VLM.")
# fmt: off
parser.add_argument("--image-path", type=Path, required=True, help="Path to image file to inspire from")
parser.add_argument("--prompt", type=str, default=None, help="Optional text prompt to blend with the image (e.g., 'Make futuristic')")
parser.add_argument("--output", type=Path, default=Path("inspired.json"), help="Output path for generated JSON prompt (default: inspired.json)")
parser.add_argument("--path", type=str, default=None, help="Local path for loading the VLM model from disk")
parser.add_argument("--top-p", type=float, default=0.9, help="Top-p sampling for VLM (default: 0.9)")
parser.add_argument("--temperature", type=float, default=0.2, help="Temperature for VLM (default: 0.2)")
parser.add_argument("--max-tokens", type=int, default=4096, help="Max tokens for VLM generation (default: 4096)")
parser.add_argument("--seed", type=int, default=None, help="Seed for VLM generation")
# fmt: on
args = parser.parse_args()
try:
# 1. Load the image
image = _load_image(args.image_path)
# 2. Generate JSON prompt from image
vlm = FiboVLM(local_path=args.path)
inspired_json = vlm.inspire(
image=image,
prompt=args.prompt,
top_p=args.top_p,
temperature=args.temperature,
max_tokens=args.max_tokens,
seed=args.seed,
)
# 3. Parse and save generated JSON prompt
_save_prompt(args, inspired_json)
except (PromptFileReadError, ValueError) as exc:
print(exc)
def _save_prompt(args, inspired_json):
try:
inspired_json_parsed = json.loads(inspired_json)
except json.JSONDecodeError as e:
raise ValueError(f"VLM did not return valid JSON: {e}")
args.output.parent.mkdir(parents=True, exist_ok=True)
with open(args.output, "w") as f:
json.dump(inspired_json_parsed, f, indent=2, ensure_ascii=False)
def _load_image(image_path: Path):
if not image_path.exists():
raise PromptFileReadError(f"Image file does not exist: {image_path}")
try:
image = Image.open(image_path).convert("RGB")
except (OSError, IOError, ValueError) as e:
raise PromptFileReadError(f"Failed to load image: {e}")
return image
if __name__ == "__main__":
main()

85
src/mflux/latent_creator/latent_creator.py
View File

@ -1,85 +0,0 @@
from pathlib import Path
import mlx.core as mx
from mlx import nn
from mflux.post_processing.array_util import ArrayUtil
from mflux.post_processing.image_util import ImageUtil
class Img2Img:
def __init__(
self,
vae: nn.Module,
sigmas: mx.array,
init_time_step: int,
image_path: str | Path | None,
):
self.vae = vae
self.sigmas = sigmas
self.init_time_step = init_time_step
self.image_path = image_path
class LatentCreator:
@staticmethod
def create(
seed: int,
height: int,
width: int,
) -> mx.array:
return mx.random.normal(
shape=[1, (height // 16) * (width // 16), 64],
key=mx.random.key(seed),
)
@staticmethod
def create_for_txt2img_or_img2img(
seed: int,
height: int,
width: int,
img2img: Img2Img,
) -> mx.array:
# 0. Determine type of image generation
if img2img.image_path is None:
# 1. Create the pure noise
return LatentCreator.create(
seed=seed,
height=height,
width=width,
)
else:
# 1. Create the pure noise
pure_noise = LatentCreator.create(
seed=seed,
height=height,
width=width,
)
# 2. Encode the image
encoded = LatentCreator.encode_image(vae=img2img.vae, image_path=img2img.image_path, height=height, width=width) # fmt: off
latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width)
# 3. Find the appropriate sigma value
sigma = img2img.sigmas[img2img.init_time_step]
# 4. Blend the appropriate amount of noise based on linear interpolation
return LatentCreator.add_noise_by_interpolation(
clean=latents,
noise=pure_noise,
sigma=sigma,
)
@staticmethod
def encode_image(vae: nn.Module, image_path: str | Path, height: int, width: int):
scaled_user_image = ImageUtil.scale_to_dimensions(
image=ImageUtil.load_image(image_path).convert("RGB"),
target_width=width,
target_height=height,
)
encoded = vae.encode(ImageUtil.to_array(scaled_user_image))
return encoded
@staticmethod
def add_noise_by_interpolation(clean: mx.array, noise: mx.array, sigma: float) -> mx.array:
return (1 - sigma) * clean + sigma * noise

2
src/mflux/lora_library.py
View File

@ -7,7 +7,7 @@ import sys
from collections import defaultdict
from pathlib import Path
from mflux.utils.lora_library import _discover_lora_files
from mflux.models.common.lora.download.lora_library import _discover_lora_files
def list_loras(paths: list[str] | None = None) -> int:

0
src/mflux/error/__init__.py → src/mflux/models/common/latent_creator/__init__.py
View File

63
src/mflux/models/common/latent_creator/latent_creator.py Normal file
View File

@ -0,0 +1,63 @@
from pathlib import Path
from typing import TYPE_CHECKING
import mlx.core as mx
from mlx import nn
from mflux.utils.image_util import ImageUtil
if TYPE_CHECKING:
from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
from mflux.models.qwen.latent_creator.qwen_latent_creator import QwenLatentCreator
class Img2Img:
def __init__(
self,
vae: nn.Module,
latent_creator: type["FiboLatentCreator"] | type["FluxLatentCreator"] | type["QwenLatentCreator"],
sigmas: mx.array,
init_time_step: int,
image_path: str | Path | None,
):
self.vae = vae
self.sigmas = sigmas
self.init_time_step = init_time_step
self.image_path = image_path
self.latent_creator = latent_creator
class LatentCreator:
@staticmethod
def create_for_txt2img_or_img2img(
seed: int,
height: int,
width: int,
img2img: Img2Img,
) -> mx.array:
latent_creator = img2img.latent_creator
if img2img.image_path is None:
# txt2img: just create noise
return latent_creator.create_noise(seed, height, width)
else:
# img2img: blend encoded image with noise
pure_noise = latent_creator.create_noise(seed, height, width)
encoded = LatentCreator.encode_image(vae=img2img.vae, image_path=img2img.image_path, height=height, width=width) # fmt: off
latents = latent_creator.pack_latents(encoded, height, width)
sigma = img2img.sigmas[img2img.init_time_step]
return LatentCreator.add_noise_by_interpolation(clean=latents, noise=pure_noise, sigma=sigma)
@staticmethod
def encode_image(vae: nn.Module, image_path: str | Path, height: int, width: int) -> mx.array:
scaled_user_image = ImageUtil.scale_to_dimensions(
image=ImageUtil.load_image(image_path).convert("RGB"),
target_width=width,
target_height=height,
)
return vae.encode(ImageUtil.to_array(scaled_user_image))
@staticmethod
def add_noise_by_interpolation(clean: mx.array, noise: mx.array, sigma: float) -> mx.array:
return (1 - sigma) * clean + sigma * noise

0
src/mflux/utils/lora_library.py → src/mflux/models/common/lora/download/lora_library.py
View File

0
src/mflux/latent_creator/__init__.py → src/mflux/models/common/quantization/__init__.py
View File

24
src/mflux/utils/quantization_util.py → src/mflux/models/common/quantization/quantization_util.py
View File

@ -3,6 +3,7 @@ from typing import TYPE_CHECKING
import mlx.nn as nn
if TYPE_CHECKING:
from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler
from mflux.models.flux.variants.controlnet.weight_handler_controlnet import WeightHandlerControlnet
from mflux.models.flux.weights.weight_handler import WeightHandler
from mflux.models.qwen.weights.qwen_weight_handler import QwenWeightHandler
@ -19,9 +20,6 @@ class QuantizationUtil:
weights: "WeightHandler",
) -> None:
q_level = weights.meta_data.quantization_level
# mx.save_tensors saves metadata dict kv 'quantization_level': 'None' as a str: str mapping
# we coerce both configs to NoneType to help users use non-quantized saved model files
if q_level == "None":
q_level = None
if quantize == "None":
@ -92,3 +90,23 @@ class QuantizationUtil:
nn.quantize(vae, bits=bits)
nn.quantize(transformer, bits=bits)
# nn.quantize(text_encoder, bits=bits) # Quantization of text encoder causes significant semantic degradation
@staticmethod
def quantize_fibo_models(
vae: nn.Module,
transformer: nn.Module,
text_encoder: nn.Module | None,
quantize: int,
weights: "FIBOWeightHandler",
) -> None:
q_level = weights.meta_data.quantization_level
if q_level == "None":
q_level = None
if quantize == "None":
quantize = None
if quantize is not None or q_level is not None:
bits = int(q_level) if q_level is not None else quantize
nn.quantize(vae, class_predicate=QuantizationUtil.quantization_predicate, bits=bits)
nn.quantize(transformer, class_predicate=QuantizationUtil.quantization_predicate, bits=bits)
nn.quantize(text_encoder, class_predicate=QuantizationUtil.quantization_predicate, bits=bits)

2
src/mflux/schedulers/__init__.py → src/mflux/models/common/schedulers/__init__.py
View File

@ -66,7 +66,7 @@ def try_import_external_scheduler(scheduler_object_path: str):
if not issubclass(SchedulerClass, BaseScheduler):
raise InvalidSchedulerType(
f"{scheduler_object_path!r} does not inherit from BaseScheduler. "
f"All schedulers must inherit from mflux.schedulers.BaseScheduler."
f"All schedulers must inherit from mflux.models.common.schedulers.BaseScheduler."
)
return SchedulerClass

0
src/mflux/schedulers/base_scheduler.py → src/mflux/models/common/schedulers/base_scheduler.py
View File

2
src/mflux/schedulers/flow_match_euler_discrete_scheduler.py → src/mflux/models/common/schedulers/flow_match_euler_discrete_scheduler.py
View File

@ -6,7 +6,7 @@ import mlx.core as mx
if TYPE_CHECKING:
from mflux.config.runtime_config import RuntimeConfig
from mflux.schedulers.base_scheduler import BaseScheduler
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
class FlowMatchEulerDiscreteScheduler(BaseScheduler):

2
src/mflux/schedulers/linear_scheduler.py → src/mflux/models/common/schedulers/linear_scheduler.py
View File

@ -5,7 +5,7 @@ import mlx.core as mx
if TYPE_CHECKING:
from mflux.config.runtime_config import RuntimeConfig
from mflux.schedulers.base_scheduler import BaseScheduler
from mflux.models.common.schedulers.base_scheduler import BaseScheduler
class LinearScheduler(BaseScheduler):

4
src/mflux/models/common/weights/__init__.py
View File

@ -1 +1,5 @@
"""Common weight mapping utilities."""
from mflux.models.common.weights.model_saver import ModelSaver
__all__ = ["ModelSaver"]

5
src/mflux/models/common/weights/mapping/weight_mapper.py
View File

@ -133,8 +133,11 @@ class WeightMapper:
flat[concrete_hf] = (concrete_mlx, target.transform)
elif has_block:
# Expand {block} only (for transformer blocks or visual blocks)
# Check if target has max_blocks override
if target.max_blocks is not None:
max_blocks = target.max_blocks
# Check if this is for visual blocks (32 blocks) or transformer blocks
if "visual.blocks" in hf_pattern or "visual.blocks" in target.mlx_path:
elif "visual.blocks" in hf_pattern or "visual.blocks" in target.mlx_path:
max_blocks = 32 # Visual blocks are always 32
else:
max_blocks = num_blocks if num_blocks > 0 else 4 # Default 4 for up_blocks

1
src/mflux/models/common/weights/mapping/weight_mapping.py
View File

@ -22,6 +22,7 @@ class WeightTarget:
hf_patterns: List[str] # HuggingFace naming patterns, e.g., ["transformer_blocks.{block}.attn.to_q.weight"]
transform: Optional[Callable[[mx.array], mx.array]] = None # Optional transform (reshape, transpose, etc.)
required: bool = True # If False, weight is optional (may not exist in all models)
max_blocks: Optional[int] = None # Override num_blocks for this target (e.g., for fixed-size lists)
class WeightMapping(Protocol):

114
src/mflux/models/common/weights/model_saver.py Normal file
View File

@ -0,0 +1,114 @@
from pathlib import Path
from typing import Any
import mlx.core as mx
from mlx import nn
from mlx.utils import tree_flatten
from transformers import PreTrainedTokenizer
from mflux.utils.version_util import VersionUtil
class ModelSaver:
@staticmethod
def save_model(
model: Any,
bits: int,
base_path: str,
tokenizers: list[tuple[str, str]] | None = None,
components: list[tuple[str, str]] | None = None,
) -> None:
# Default tokenizers: try common patterns
if tokenizers is None:
tokenizers = ModelSaver._detect_tokenizers(model)
# Default components: try common patterns
if components is None:
components = ModelSaver._detect_components(model)
# Save tokenizers
for attr_path, subdir in tokenizers:
tokenizer = ModelSaver._get_nested_attr(model, attr_path)
if tokenizer is not None:
ModelSaver._save_tokenizer(base_path, tokenizer, subdir)
# Save model components
for attr_name, subdir in components:
component = getattr(model, attr_name, None)
if component is not None:
ModelSaver._save_weights(base_path, bits, component, subdir)
@staticmethod
def _detect_tokenizers(model: Any) -> list[tuple[str, str]]:
tokenizers = []
if hasattr(model, "clip_tokenizer") and hasattr(model.clip_tokenizer, "tokenizer"):
tokenizers.append(("clip_tokenizer.tokenizer", "tokenizer"))
if hasattr(model, "t5_tokenizer") and hasattr(model.t5_tokenizer, "tokenizer"):
tokenizers.append(("t5_tokenizer.tokenizer", "tokenizer_2"))
if hasattr(model, "qwen_tokenizer") and hasattr(model.qwen_tokenizer, "tokenizer"):
tokenizers.append(("qwen_tokenizer.tokenizer", "tokenizer"))
if hasattr(model, "fibo_tokenizer") and hasattr(model.fibo_tokenizer, "tokenizer"):
tokenizers.append(("fibo_tokenizer.tokenizer", "tokenizer"))
return tokenizers
@staticmethod
def _detect_components(model: Any) -> list[tuple[str, str]]:
components = []
if hasattr(model, "vae"):
components.append(("vae", "vae"))
if hasattr(model, "transformer"):
components.append(("transformer", "transformer"))
if hasattr(model, "clip_text_encoder"):
components.append(("clip_text_encoder", "text_encoder"))
if hasattr(model, "t5_text_encoder"):
components.append(("t5_text_encoder", "text_encoder_2"))
if hasattr(model, "text_encoder"):
# Only add if we haven't already added clip_text_encoder or t5_text_encoder
if not any(c[1] == "text_encoder" for c in components):
components.append(("text_encoder", "text_encoder"))
return components
@staticmethod
def _get_nested_attr(obj: Any, attr_path: str) -> Any:
attrs = attr_path.split(".")
result = obj
for attr in attrs:
if not hasattr(result, attr):
return None
result = getattr(result, attr)
return result
@staticmethod
def _save_tokenizer(base_path: str, tokenizer: PreTrainedTokenizer, subdir: str) -> None:
path = Path(base_path) / subdir
path.mkdir(parents=True, exist_ok=True)
tokenizer.save_pretrained(path)
@staticmethod
def _save_weights(base_path: str, bits: int, model: nn.Module, subdir: str) -> None:
path = Path(base_path) / subdir
path.mkdir(parents=True, exist_ok=True)
weights = ModelSaver._split_weights(base_path, dict(tree_flatten(model.parameters())))
for i, weight in enumerate(weights):
mx.save_safetensors(
str(path / f"{i}.safetensors"),
weight,
{
"quantization_level": str(bits),
"mflux_version": VersionUtil.get_mflux_version(),
},
)
@staticmethod
def _split_weights(base_path: str, weights: dict, max_file_size_gb: int = 2) -> list:
max_file_size_bytes = max_file_size_gb << 30
shards = []
shard, shard_size = {}, 0
for k, v in weights.items():
if shard_size + v.nbytes > max_file_size_bytes:
shards.append(shard)
shard, shard_size = {}, 0
shard[k] = v
shard_size += v.nbytes
shards.append(shard)
return shards

2
src/mflux/models/depth_pro/depth_pro.py
View File

@ -9,7 +9,7 @@ from PIL import Image
from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer
from mflux.models.depth_pro.model.depth_pro_model import DepthProModel
from mflux.models.depth_pro.model.depth_pro_util import DepthProUtil
from mflux.post_processing.image_util import ImageUtil
from mflux.utils.image_util import ImageUtil
@dataclass

0
src/mflux/post_processing/__init__.py → src/mflux/models/fibo/__init__.py
View File

47
src/mflux/models/fibo/fibo_initializer.py Normal file
View File

@ -0,0 +1,47 @@
from mflux.config.model_config import ModelConfig
from mflux.models.fibo.model.fibo_text_encoder import SmolLM3_3B_TextEncoder
from mflux.models.fibo.model.fibo_transformer import FiboTransformer
from mflux.models.fibo.model.fibo_vae.wan_2_2_vae import Wan2_2_VAE
from mflux.models.fibo.tokenizer import FiboTokenizerHandler
from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler
from mflux.models.fibo.weights.fibo_weight_util import FIBOWeightUtil
class FIBOInitializer:
@staticmethod
def init(
fibo_model,
model_config: ModelConfig | None = None,
quantize: int | None = None,
local_path: str | None = None,
) -> None:
# 1. Load VAE weights
weights = FIBOWeightHandler.load_regular_weights(
repo_id=model_config.model_name,
local_path=local_path,
)
# 2. Initialize tokenizers
tokenizer_handler = FiboTokenizerHandler(
repo_id=model_config.model_name,
local_path=local_path,
)
fibo_model.fibo_tokenizer = tokenizer_handler.fibo
# 3. Initialize all models
fibo_model.vae = Wan2_2_VAE()
fibo_model.text_encoder = SmolLM3_3B_TextEncoder()
fibo_model.transformer = FiboTransformer(
in_channels=48,
num_layers=8,
num_single_layers=38,
)
# 4. Apply weights and quantize VAE, transformer, and text encoder
fibo_model.bits = FIBOWeightUtil.set_weights_and_quantize(
quantize_arg=quantize,
weights=weights,
vae=fibo_model.vae,
transformer=fibo_model.transformer,
text_encoder=fibo_model.text_encoder,
)

0
src/mflux/models/fibo/latent_creator/__init__.py Normal file
View File

19
src/mflux/models/fibo/latent_creator/fibo_latent_creator.py Normal file
View File

@ -0,0 +1,19 @@
import mlx.core as mx
class FiboLatentCreator:
@staticmethod
def create_noise(seed: int, height: int, width: int) -> mx.array:
latents = mx.random.normal(
shape=(1, 48, (height // 16), (width // 16)),
key=mx.random.key(seed),
)
return FiboLatentCreator.pack_latents(latents, height, width)
@staticmethod
def pack_latents(latents: mx.array, height: int, width: int) -> mx.array:
if latents.ndim == 5:
latents = latents[:, :, 0, :, :]
batch_size, channels, latent_height, latent_width = latents.shape
latents = mx.transpose(latents, (0, 2, 3, 1))
return mx.reshape(latents, (batch_size, latent_height * latent_width, channels))

1
src/mflux/models/fibo/model/__init__.py Normal file
View File

@ -0,0 +1 @@
"""FIBO model components."""

7
src/mflux/models/fibo/model/fibo_text_encoder/__init__.py Normal file
View File

@ -0,0 +1,7 @@
from .prompt_encoder import PromptEncoder
from .smol_lm3_3b_text_encoder import SmolLM3_3B_TextEncoder
__all__ = [
"PromptEncoder",
"SmolLM3_3B_TextEncoder",
]

166
src/mflux/models/fibo/model/fibo_text_encoder/prompt_encoder.py Normal file
View File

@ -0,0 +1,166 @@
import json
from typing import List, Union
import mlx.core as mx
from mflux.models.fibo.model.fibo_text_encoder.smol_lm3_3b_text_encoder import SmolLM3_3B_TextEncoder
from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo
class PromptEncoder:
@staticmethod
def encode_prompt(
prompt: str,
negative_prompt: str | None,
tokenizer: TokenizerFibo,
text_encoder: SmolLM3_3B_TextEncoder,
) -> tuple[str, mx.array, List[mx.array]]:
# 0. Set default negative prompt if not provided
if negative_prompt is None or negative_prompt == "":
negative_prompt = "ugly, blurry, low quality"
# 1. Convert prompt to JSON format
json.loads(prompt)
json_prompt = prompt
# 2. Get prompt embeddings for positive and negative prompt
prompt_embeds, prompt_layers, prompt_attention_mask = PromptEncoder._get_prompt_embeds(
prompt=json_prompt,
tokenizer=tokenizer,
text_encoder=text_encoder,
num_images_per_prompt=1,
max_sequence_length=2048,
tokenization_prefix="positive",
)
neg_prompt_embeds, neg_prompt_layers, neg_prompt_attention_mask = PromptEncoder._get_prompt_embeds(
prompt=negative_prompt,
tokenizer=tokenizer,
text_encoder=text_encoder,
num_images_per_prompt=1,
max_sequence_length=2048,
tokenization_prefix="negative",
)
encoder_hidden_states, max_tokens = PromptEncoder._get_encoder_hidden_states(
prompt_embeds=prompt_embeds,
neg_prompt_embeds=neg_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
neg_prompt_attention_mask=neg_prompt_attention_mask,
)
prompt_layers = PromptEncoder._get_prompt_layers(
max_tokens=max_tokens,
prompt_layers=prompt_layers,
neg_prompt_layers=neg_prompt_layers,
)
return json_prompt, encoder_hidden_states, prompt_layers
@staticmethod
def _get_encoder_hidden_states(neg_prompt_attention_mask, neg_prompt_embeds, prompt_attention_mask, prompt_embeds):
max_tokens = max(neg_prompt_embeds.shape[1], prompt_embeds.shape[1])
prompt_embeds, prompt_attention_mask = PromptEncoder._pad_embedding(
max_tokens=max_tokens,
prompt_embeds=prompt_embeds,
attention_mask=prompt_attention_mask,
)
neg_prompt_embeds, neg_prompt_attention_mask = PromptEncoder._pad_embedding(
max_tokens=max_tokens,
prompt_embeds=neg_prompt_embeds,
attention_mask=neg_prompt_attention_mask,
)
encoder_hidden_states = mx.concatenate([neg_prompt_embeds, prompt_embeds], axis=0)
return encoder_hidden_states, max_tokens
@staticmethod
def _get_prompt_layers(max_tokens, neg_prompt_layers, prompt_layers):
prompt_layers = [PromptEncoder._pad_embedding(layer, max_tokens)[0] for layer in prompt_layers]
neg_prompt_layers = [PromptEncoder._pad_embedding(layer, max_tokens)[0] for layer in neg_prompt_layers]
prompt_layers = [mx.concatenate([neg_prompt_layers[i], prompt_layers[i]], axis=0) for i in range(len(prompt_layers))] # fmt: off
total_num_layers_transformer = 46 # FIBO transformer has 8 joint + 38 single blocks = 46 total layers
if len(prompt_layers) >= total_num_layers_transformer:
prompt_layers = prompt_layers[len(prompt_layers) - total_num_layers_transformer :]
else:
prompt_layers = prompt_layers + [prompt_layers[-1]] * (total_num_layers_transformer - len(prompt_layers))
return prompt_layers
@staticmethod
def _get_prompt_embeds(
prompt: Union[str, List[str]],
text_encoder: SmolLM3_3B_TextEncoder,
tokenizer: TokenizerFibo,
num_images_per_prompt: int = 1,
max_sequence_length: int = 2048,
tokenization_prefix: str | None = None,
) -> tuple[mx.array, List[mx.array], mx.array]:
prompts = [prompt] if isinstance(prompt, str) else list(prompt)
if not prompts:
raise ValueError("`prompt` must be a non-empty string or list of strings.")
# 1) Tokenize and convert to MX
input_ids_mx, attention_mask_mx = tokenizer.tokenize(
prompts=prompts,
max_length=max_sequence_length,
padding="longest",
truncation=True,
add_special_tokens=True,
)
# 2) Run MLX text encoder and collect hidden states
hidden_states_list = text_encoder(
input_ids=input_ids_mx,
attention_mask=attention_mask_mx,
output_hidden_states=True,
)
# 3) Build prompt_embeds from last two layers
last_hidden = hidden_states_list[-1]
second_last_hidden = hidden_states_list[-2]
prompt_embeds = mx.concatenate(
[last_hidden, second_last_hidden],
axis=-1,
)
# 4) Repeat along batch dimension for num_images_per_prompt
prompt_embeds = PromptEncoder._repeat_interleave_batch(prompt_embeds, num_images_per_prompt)
attention_mask = PromptEncoder._repeat_interleave_batch(attention_mask_mx, num_images_per_prompt)
prompt_layers = [PromptEncoder._repeat_interleave_batch(layer, num_images_per_prompt) for layer in hidden_states_list] # fmt: off
return prompt_embeds, prompt_layers, attention_mask
@staticmethod
def _pad_embedding(
prompt_embeds: mx.array,
max_tokens: int,
attention_mask: mx.array | None = None,
) -> tuple[mx.array, mx.array]:
batch_size, seq_len, dim = prompt_embeds.shape
if attention_mask is None:
attention_mask = mx.ones((batch_size, seq_len), dtype=prompt_embeds.dtype)
else:
attention_mask = attention_mask.astype(prompt_embeds.dtype)
if max_tokens < seq_len:
raise ValueError("`max_tokens` must be >= current sequence length.")
if max_tokens > seq_len:
pad_length = max_tokens - seq_len
padding = mx.zeros((batch_size, pad_length, dim), dtype=prompt_embeds.dtype)
prompt_embeds = mx.concatenate([prompt_embeds, padding], axis=1)
mask_padding = mx.zeros((batch_size, pad_length), dtype=attention_mask.dtype)
attention_mask = mx.concatenate([attention_mask, mask_padding], axis=1)
return prompt_embeds, attention_mask
@staticmethod
def _repeat_interleave_batch(
tensor: mx.array,
num_images_per_prompt: int,
) -> mx.array:
if num_images_per_prompt == 1:
return tensor
batch, *rest = tensor.shape
tensor = mx.expand_dims(tensor, axis=1)
tensor = mx.broadcast_to(tensor, (batch, num_images_per_prompt, *rest))
new_shape = (batch * num_images_per_prompt, *rest)
return mx.reshape(tensor, new_shape)

133
src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_attention.py Normal file
View File

@ -0,0 +1,133 @@
import math
from typing import Tuple
import mlx.core as mx
from mlx import nn
from mlx.core.fast import scaled_dot_product_attention
from .smol_lm3_3b_rope import SmolLM3_3B_RotaryEmbedding
class SmolLM3_3B_SelfAttention(nn.Module):
def __init__(
self,
hidden_size: int,
num_attention_heads: int,
num_key_value_heads: int,
max_position_embeddings: int = 65_536,
rope_theta: float = 5_000_000.0,
attention_dropout: float = 0.0,
):
super().__init__()
self.hidden_size = hidden_size
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.head_dim = hidden_size // num_attention_heads
self.num_key_value_groups = num_attention_heads // num_key_value_heads
self.scale = 1.0 / math.sqrt(self.head_dim)
self.attention_dropout = attention_dropout
# Projections – no bias, matching SmolLM3 config (attention_bias = False)
self.q_proj = nn.Linear(hidden_size, num_attention_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(hidden_size, num_key_value_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(num_attention_heads * self.head_dim, hidden_size, bias=False)
self.rotary_emb = SmolLM3_3B_RotaryEmbedding(
dim=self.head_dim,
max_position_embeddings=max_position_embeddings,
base=rope_theta,
)
def __call__(
self,
hidden_states: mx.array,
attention_mask: mx.array | None = None,
cos_sin: Tuple[mx.array, mx.array] | None = None,
) -> mx.array:
batch_size, seq_len, _ = hidden_states.shape
# Projections: Use nn.Linear directly (like Qwen/Flux models)
q = self.q_proj(hidden_states)
k = self.k_proj(hidden_states)
v = self.v_proj(hidden_states)
# Cast back to input dtype if Linear output is float32 (for consistency)
if q.dtype != hidden_states.dtype:
q = q.astype(hidden_states.dtype)
k = k.astype(hidden_states.dtype)
v = v.astype(hidden_states.dtype)
# Force evaluation to ensure computation is complete
mx.eval(q, k, v)
# Reshape to (batch, heads, seq, head_dim)
q = q.reshape(batch_size, seq_len, self.num_attention_heads, self.head_dim).transpose(0, 2, 1, 3)
k = k.reshape(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(0, 2, 1, 3)
v = v.reshape(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(0, 2, 1, 3)
# Rotary embeddings
if cos_sin is None:
cos_sin = self.rotary_emb(seq_len)
cos, sin = cos_sin
q, k = SmolLM3_3B_SelfAttention._apply_rope(q, k, cos, sin)
# Grouped-query attention: repeat kv heads if needed
if self.num_key_value_heads != self.num_attention_heads:
k = SmolLM3_3B_SelfAttention._repeat_kv(k, self.num_key_value_groups)
v = SmolLM3_3B_SelfAttention._repeat_kv(v, self.num_key_value_groups)
# Prepare mask for scaled_dot_product_attention
attn_mask = None
if attention_mask is not None:
seq_len_k = k.shape[2]
causal_mask = attention_mask[:, :, :, :seq_len_k]
if causal_mask.shape[1] == 1:
causal_mask = mx.broadcast_to(causal_mask, (batch_size, self.num_attention_heads, seq_len, seq_len_k))
attn_mask = causal_mask.astype(q.dtype)
# Use MLX fast SDPA
attn_output = scaled_dot_product_attention(
q,
k,
v,
scale=self.scale,
mask=attn_mask,
)
attn_output = attn_output.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, self.hidden_size)
attn_output = self.o_proj(attn_output)
return attn_output
@staticmethod
def _rotate_half(x: mx.array) -> mx.array:
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return mx.concatenate([-x2, x1], axis=-1)
@staticmethod
def _apply_rope(
q: mx.array,
k: mx.array,
cos: mx.array,
sin: mx.array,
) -> Tuple[mx.array, mx.array]:
q_dtype = q.dtype
k_dtype = k.dtype
q = q.astype(mx.float32)
k = k.astype(mx.float32)
cos = cos.astype(mx.float32)
sin = sin.astype(mx.float32)
q_embed = (q * cos) + (SmolLM3_3B_SelfAttention._rotate_half(q) * sin)
k_embed = (k * cos) + (SmolLM3_3B_SelfAttention._rotate_half(k) * sin)
return q_embed.astype(q_dtype), k_embed.astype(k_dtype)
@staticmethod
def _repeat_kv(hidden_states: mx.array, n_rep: int) -> mx.array:
batch, num_kv_heads, seq_len, head_dim = hidden_states.shape
hidden_states = mx.expand_dims(hidden_states, axis=2)
hidden_states = mx.broadcast_to(hidden_states, (batch, num_kv_heads, n_rep, seq_len, head_dim))
return hidden_states.reshape(batch, num_kv_heads * n_rep, seq_len, head_dim)

64
src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_encoder_layer.py Normal file
View File

@ -0,0 +1,64 @@
from typing import Tuple
import mlx.core as mx
from mlx import nn
from .smol_lm3_3b_attention import SmolLM3_3B_SelfAttention
from .smol_lm3_3b_mlp import SmolLM3_3B_MLP
from .smol_lm3_3b_rms_norm import SmolLM3_3B_RMSNorm
class SmolLM3_3B_EncoderLayer(nn.Module):
def __init__(
self,
hidden_size: int = 2048,
num_attention_heads: int = 16,
num_key_value_heads: int = 4,
intermediate_size: int = 11_008,
rms_norm_eps: float = 1e-6,
max_position_embeddings: int = 65_536,
rope_theta: float = 5_000_000.0,
attention_dropout: float = 0.0,
hidden_act: str = "silu",
):
super().__init__()
self.input_layernorm = SmolLM3_3B_RMSNorm(hidden_size, eps=rms_norm_eps)
self.self_attn = SmolLM3_3B_SelfAttention(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
num_key_value_heads=num_key_value_heads,
max_position_embeddings=max_position_embeddings,
rope_theta=rope_theta,
attention_dropout=attention_dropout,
)
self.post_attention_layernorm = SmolLM3_3B_RMSNorm(hidden_size, eps=rms_norm_eps)
self.mlp = SmolLM3_3B_MLP(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act=hidden_act,
)
def __call__(
self,
hidden_states: mx.array,
attention_mask: mx.array | None,
cos_sin: Tuple[mx.array, mx.array],
layer_idx: int | None = None,
) -> mx.array:
# Self-attention block
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
cos_sin=cos_sin,
)
hidden_states = residual + hidden_states
# Feed-forward block
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states
return hidden_states

30
src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_mlp.py Normal file
View File

@ -0,0 +1,30 @@
import mlx.core as mx
from mlx import nn
class SmolLM3_3B_MLP(nn.Module):
def __init__(
self,
hidden_size: int,
intermediate_size: int,
hidden_act: str = "silu",
):
super().__init__()
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.hidden_act = hidden_act
self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
def _activation(self, x: mx.array) -> mx.array:
if self.hidden_act == "silu":
return x * mx.sigmoid(x)
return x * mx.sigmoid(x)
def __call__(self, hidden_states: mx.array) -> mx.array:
gate = self._activation(self.gate_proj(hidden_states))
up = self.up_proj(hidden_states)
hidden = gate * up
hidden = self.down_proj(hidden)
return hidden

17
src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_rms_norm.py Normal file
View File

@ -0,0 +1,17 @@
import mlx.core as mx
from mlx import nn
class SmolLM3_3B_RMSNorm(nn.Module):
def __init__(self, hidden_size: int, eps: float = 1e-6):
super().__init__()
self.weight = mx.ones((hidden_size,))
self.eps = eps
def __call__(self, hidden_states: mx.array) -> mx.array:
input_dtype = hidden_states.dtype
hidden_states = hidden_states.astype(mx.float32)
variance = mx.mean(mx.square(hidden_states), axis=-1, keepdims=True)
hidden_states = hidden_states * mx.rsqrt(variance + self.eps)
result = self.weight.astype(mx.float32) * hidden_states
return result.astype(input_dtype)

31
src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_rope.py Normal file
View File

@ -0,0 +1,31 @@
from typing import Tuple
import mlx.core as mx
from mlx import nn
class SmolLM3_3B_RotaryEmbedding(nn.Module):
def __init__(
self,
dim: int,
max_position_embeddings: int = 65_536,
base: float = 5_000_000.0,
):
super().__init__()
self.inv_freq = 1.0 / (base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
self.max_position_embeddings = max_position_embeddings
def __call__(self, seq_len: int) -> Tuple[mx.array, mx.array]:
cos, sin = SmolLM3_3B_RotaryEmbedding._build_cos_sin(self.inv_freq, seq_len)
cos = mx.expand_dims(mx.expand_dims(cos, axis=0), axis=0)
sin = mx.expand_dims(mx.expand_dims(sin, axis=0), axis=0)
return cos, sin
@staticmethod
def _build_cos_sin(inv_freq, seq_len: int) -> Tuple[mx.array, mx.array]:
positions = mx.arange(seq_len, dtype=mx.float32)
freqs = mx.outer(positions, inv_freq)
emb = mx.concatenate([freqs, freqs], axis=-1)
cos = mx.cos(emb)
sin = mx.sin(emb)
return cos, sin

105
src/mflux/models/fibo/model/fibo_text_encoder/smol_lm3_3b_text_encoder.py Normal file
View File

@ -0,0 +1,105 @@
from typing import List
import mlx.core as mx
from mlx import nn
from .smol_lm3_3b_encoder_layer import SmolLM3_3B_EncoderLayer
from .smol_lm3_3b_rms_norm import SmolLM3_3B_RMSNorm
from .smol_lm3_3b_rope import SmolLM3_3B_RotaryEmbedding
class SmolLM3_3B_TextEncoder(nn.Module):
def __init__(
self,
vocab_size: int = 128_256,
hidden_size: int = 2048,
intermediate_size: int = 11_008,
num_hidden_layers: int = 36,
num_attention_heads: int = 16,
num_key_value_heads: int = 4,
max_position_embeddings: int = 65_536,
rope_theta: float = 5_000_000.0,
rms_norm_eps: float = 1e-6,
hidden_act: str = "silu",
):
super().__init__()
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.max_position_embeddings = max_position_embeddings
self.rope_theta = rope_theta
self.rms_norm_eps = rms_norm_eps
self.hidden_act = hidden_act
self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
self.layers: List[SmolLM3_3B_EncoderLayer] = [
SmolLM3_3B_EncoderLayer(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
num_key_value_heads=num_key_value_heads,
intermediate_size=intermediate_size,
rms_norm_eps=rms_norm_eps,
max_position_embeddings=max_position_embeddings,
rope_theta=rope_theta,
hidden_act=hidden_act,
)
for _ in range(num_hidden_layers)
]
self.norm = SmolLM3_3B_RMSNorm(hidden_size, eps=rms_norm_eps)
self.rotary_emb = SmolLM3_3B_RotaryEmbedding(
dim=hidden_size // num_attention_heads,
max_position_embeddings=max_position_embeddings,
base=rope_theta,
)
def __call__(
self,
input_ids: mx.array,
attention_mask: mx.array,
output_hidden_states: bool = True,
) -> List[mx.array] | mx.array:
batch_size, seq_len = input_ids.shape
hidden_states = self.embed_tokens(input_ids)
attention_mask_4d = SmolLM3_3B_TextEncoder._build_attention_mask(attention_mask)
cos, sin = self.rotary_emb(seq_len)
hidden_states_list: List[mx.array] = [hidden_states]
for layer_idx, layer in enumerate(self.layers):
hidden_states = layer(
hidden_states=hidden_states,
attention_mask=attention_mask_4d,
cos_sin=(cos, sin),
layer_idx=layer_idx,
)
if output_hidden_states:
hidden_states_list.append(hidden_states)
hidden_states = self.norm(hidden_states)
if output_hidden_states:
hidden_states_list[-1] = hidden_states
return hidden_states_list
return hidden_states
@staticmethod
def _build_attention_mask(attention_mask: mx.array) -> mx.array:
batch_size, seq_len = attention_mask.shape
mask_dtype = mx.float32
min_dtype_value = mx.finfo(mask_dtype).min
padding_mask = mx.where(
attention_mask == 1,
mx.zeros_like(attention_mask).astype(mask_dtype),
mx.ones_like(attention_mask).astype(mask_dtype) * min_dtype_value,
)
padding_mask = mx.expand_dims(mx.expand_dims(padding_mask, axis=1), axis=1)
idx = mx.arange(seq_len, dtype=mx.int32)
j = mx.expand_dims(idx, axis=0)
i = mx.expand_dims(idx, axis=1)
tri_bool = j > i
zeros_2d = mx.zeros((seq_len, seq_len), dtype=mask_dtype)
minval_2d = mx.ones((seq_len, seq_len), dtype=mask_dtype) * min_dtype_value
causal_tri_mask = mx.where(tri_bool, minval_2d, zeros_2d)
causal_tri_mask = mx.expand_dims(mx.expand_dims(causal_tri_mask, axis=0), axis=0)
causal_tri_mask = mx.broadcast_to(causal_tri_mask, (batch_size, 1, seq_len, seq_len))
attention_mask_4d = causal_tri_mask + padding_mask
return attention_mask_4d

3
src/mflux/models/fibo/model/fibo_transformer/__init__.py Normal file
View File

@ -0,0 +1,3 @@
from .transformer import FiboTransformer
__all__ = ["FiboTransformer"]

38
src/mflux/models/fibo/model/fibo_transformer/bria_fibo_timesteps.py Normal file
View File

@ -0,0 +1,38 @@
import math
import mlx.core as mx
from mlx import nn
class BriaFiboTimesteps(nn.Module):
def __init__(
self,
num_channels: int,
flip_sin_to_cos: bool,
downscale_freq_shift: float,
scale: float = 1.0,
time_theta: int = 10000,
):
super().__init__()
self.num_channels = num_channels
self.flip_sin_to_cos = flip_sin_to_cos
self.downscale_freq_shift = downscale_freq_shift
self.scale = scale
self.time_theta = time_theta
def __call__(self, timesteps: mx.array) -> mx.array:
half_dim = self.num_channels // 2
exponent = -math.log(self.time_theta) * mx.arange(0, half_dim, dtype=mx.float32)
exponent = exponent / (half_dim - self.downscale_freq_shift)
emb = mx.exp(exponent)
emb = mx.expand_dims(timesteps.astype(mx.float32), axis=-1) * mx.expand_dims(emb, axis=0)
emb = self.scale * emb
sin = mx.sin(emb)
cos = mx.cos(emb)
emb = mx.concatenate([sin, cos], axis=-1)
if self.flip_sin_to_cos:
emb = mx.concatenate([emb[:, half_dim:], emb[:, :half_dim]], axis=-1)
if self.num_channels % 2 == 1:
pad = mx.zeros((emb.shape[0], 1), dtype=emb.dtype)
emb = mx.concatenate([emb, pad], axis=-1)
return emb

28
src/mflux/models/fibo/model/fibo_transformer/feed_forward.py Normal file
View File

@ -0,0 +1,28 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_transformer.fibo_gelu import FiboGELU
class FiboFeedForward(nn.Module):
def __init__(
self,
dim: int,
dim_out: int | None = None,
mult: int = 4,
activation_fn: str = "gelu-approximate",
dropout: float = 0.0,
):
super().__init__()
inner_dim = int(dim * mult)
dim_out = dim if dim_out is None else dim_out
self.net: list[nn.Module] = [
FiboGELU(dim_in=dim, dim_out=inner_dim, approximate="tanh", bias=True),
nn.Dropout(dropout),
nn.Linear(inner_dim, dim_out, bias=True),
]
def __call__(self, hidden_states: mx.array) -> mx.array:
for module in self.net:
hidden_states = module(hidden_states)
return hidden_states

31
src/mflux/models/fibo/model/fibo_transformer/fibo_ada_layer_norm_zero.py Normal file
View File

@ -0,0 +1,31 @@
import mlx.core as mx
from mlx import nn
class FiboAdaLayerNormZero(nn.Module):
def __init__(self, embedding_dim: int, eps: float = 1e-6):
super().__init__()
self.embedding_dim = embedding_dim
self.eps = eps
self.linear = nn.Linear(embedding_dim, 6 * embedding_dim)
def __call__(self, hidden_states: mx.array, text_embeddings: mx.array):
emb = self.linear(nn.silu(text_embeddings))
chunk = self.embedding_dim
shift_msa = emb[:, 0 * chunk : 1 * chunk]
scale_msa = emb[:, 1 * chunk : 2 * chunk]
gate_msa = emb[:, 2 * chunk : 3 * chunk]
shift_mlp = emb[:, 3 * chunk : 4 * chunk]
scale_mlp = emb[:, 4 * chunk : 5 * chunk]
gate_mlp = emb[:, 5 * chunk : 6 * chunk]
norm_hidden_states = FiboAdaLayerNormZero._layer_norm(self.eps, hidden_states)
hidden_states = norm_hidden_states * (1 + scale_msa[:, None, :]) + shift_msa[:, None, :]
return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
@staticmethod
def _layer_norm(eps, x: mx.array) -> mx.array:
x_f32 = x.astype(mx.float32)
mean = mx.mean(x_f32, axis=-1, keepdims=True)
var = mx.mean((x_f32 - mean) ** 2, axis=-1, keepdims=True)
y = (x_f32 - mean) / mx.sqrt(var + eps)
return y.astype(x.dtype)

50
src/mflux/models/fibo/model/fibo_transformer/fibo_embed_nd.py Normal file
View File

@ -0,0 +1,50 @@
import mlx.core as mx
from mlx import nn
class FiboEmbedND(nn.Module):
def __init__(self, theta: int = 10000, axes_dim: list[int] | None = None):
super().__init__()
self.theta = theta
self.axes_dim = axes_dim or [16, 56, 56]
def __call__(self, ids: mx.array) -> tuple[mx.array, mx.array]:
if ids.ndim == 3 and ids.shape[0] == 1:
ids = ids[0]
n_axes = ids.shape[-1]
pos = ids.astype(mx.float32)
cos_out: list[mx.array] = []
sin_out: list[mx.array] = []
for i in range(n_axes):
axis_dim = self.axes_dim[i]
cos_axis, sin_axis = FiboEmbedND._get_1d_rotary_pos_embed(
dim=axis_dim,
pos=pos[:, i],
theta=self.theta,
)
cos_out.append(cos_axis)
sin_out.append(sin_axis)
freqs_cos = mx.concatenate(cos_out, axis=-1)
freqs_sin = mx.concatenate(sin_out, axis=-1)
return freqs_cos, freqs_sin
@staticmethod
def _get_1d_rotary_pos_embed(
dim: int,
pos: mx.array,
theta: float = 10000.0,
) -> tuple[mx.array, mx.array]:
if pos.ndim != 1:
pos = mx.reshape(pos, (-1,))
pos = pos.astype(mx.float32)
freqs = 1.0 / (theta ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
angles = pos[:, None] * freqs[None, :]
cos_base = mx.cos(angles)
sin_base = mx.sin(angles)
cos = mx.reshape(mx.stack([cos_base, cos_base], axis=-1), (pos.shape[0], -1))
sin = mx.reshape(mx.stack([sin_base, sin_base], axis=-1), (pos.shape[0], -1))
return cos, sin

14
src/mflux/models/fibo/model/fibo_transformer/fibo_gelu.py Normal file
View File

@ -0,0 +1,14 @@
import mlx.core as mx
from mlx import nn
class FiboGELU(nn.Module):
def __init__(self, dim_in: int, dim_out: int, approximate: str = "tanh", bias: bool = True):
super().__init__()
self.proj = nn.Linear(dim_in, dim_out, bias=bias)
self.approximate = approximate
def __call__(self, hidden_states: mx.array) -> mx.array:
hidden_states = self.proj(hidden_states)
hidden_states = nn.gelu_approx(hidden_states)
return hidden_states

115
src/mflux/models/fibo/model/fibo_transformer/fibo_joint_attention.py Normal file
View File

@ -0,0 +1,115 @@
import mlx.core as mx
from mlx import nn
from mlx.core.fast import scaled_dot_product_attention
from mflux.models.fibo.model.fibo_transformer.fibo_single_attention import FiboSingleAttention
class FiboJointAttention(nn.Module):
def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, eps: float = 1e-6):
super().__init__()
self.head_dim = attention_head_dim
self.num_heads = num_attention_heads
self.inner_dim = dim
self.to_q = nn.Linear(dim, self.inner_dim)
self.to_k = nn.Linear(dim, self.inner_dim)
self.to_v = nn.Linear(dim, self.inner_dim)
self.norm_q = nn.RMSNorm(self.head_dim, eps=eps)
self.norm_k = nn.RMSNorm(self.head_dim, eps=eps)
# Added KV projections for encoder_hidden_states
self.add_q_proj = nn.Linear(dim, self.inner_dim)
self.add_k_proj = nn.Linear(dim, self.inner_dim)
self.add_v_proj = nn.Linear(dim, self.inner_dim)
self.norm_added_q = nn.RMSNorm(self.head_dim, eps=eps)
self.norm_added_k = nn.RMSNorm(self.head_dim, eps=eps)
# Output projections
self.to_out = [nn.Linear(self.inner_dim, dim)]
self.to_add_out = nn.Linear(self.inner_dim, dim)
def __call__(
self,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
image_rotary_emb: tuple[mx.array, mx.array],
attention_mask: mx.array | None = None,
) -> tuple[mx.array, mx.array]:
batch_size, seq_img, dim = hidden_states.shape
_, seq_ctx, _ = encoder_hidden_states.shape
cos, sin = image_rotary_emb
# QKV for image stream: [B, S_img, inner_dim]
query = self.to_q(hidden_states)
key = self.to_k(hidden_states)
value = self.to_v(hidden_states)
# QKV for context stream (added_kv): [B, S_ctx, inner_dim]
enc_query = self.add_q_proj(encoder_hidden_states)
enc_key = self.add_k_proj(encoder_hidden_states)
enc_value = self.add_v_proj(encoder_hidden_states)
# Reshape to [B, S, H, D]
query = mx.reshape(query, (batch_size, seq_img, self.num_heads, self.head_dim))
key = mx.reshape(key, (batch_size, seq_img, self.num_heads, self.head_dim))
value = mx.reshape(value, (batch_size, seq_img, self.num_heads, self.head_dim))
enc_query = mx.reshape(enc_query, (batch_size, seq_ctx, self.num_heads, self.head_dim))
enc_key = mx.reshape(enc_key, (batch_size, seq_ctx, self.num_heads, self.head_dim))
enc_value = mx.reshape(enc_value, (batch_size, seq_ctx, self.num_heads, self.head_dim))
# RMSNorm over last dim: produce normalized Q/K matching BriaFiboAttention semantics.
query = self.norm_q(query.astype(mx.float32)).astype(query.dtype)
key = self.norm_k(key.astype(mx.float32)).astype(key.dtype)
enc_query = self.norm_added_q(enc_query.astype(mx.float32)).astype(enc_query.dtype)
enc_key = self.norm_added_k(enc_key.astype(mx.float32)).astype(enc_key.dtype)
# Concatenate encoder + image along sequence dim (matches BriaFiboAttnProcessor)
query = mx.concatenate([enc_query, query], axis=1)
key = mx.concatenate([enc_key, key], axis=1)
value = mx.concatenate([enc_value, value], axis=1)
seq_total = seq_ctx + seq_img
# Apply RoPE to Q,K (layout [B, S_total, H, D])
query = FiboSingleAttention.apply_rotary_emb(query, cos, sin)
key = FiboSingleAttention.apply_rotary_emb(key, cos, sin)
# Convert to [B, H, S, D] for fast SDPA
query_bhsd = mx.transpose(query, (0, 2, 1, 3))
key_bhsd = mx.transpose(key, (0, 2, 1, 3))
value_bhsd = mx.transpose(value, (0, 2, 1, 3))
# Prepare mask for scaled_dot_product_attention
attn_mask = None
if attention_mask is not None:
attn_mask = mx.broadcast_to(attention_mask, (batch_size, self.num_heads, seq_total, seq_total))
attn_mask = attn_mask.astype(query_bhsd.dtype)
scale = 1.0 / mx.sqrt(mx.array(self.head_dim, dtype=query_bhsd.dtype))
attn_output_bhsd = scaled_dot_product_attention(
query_bhsd,
key_bhsd,
value_bhsd,
scale=scale,
mask=attn_mask,
)
# Back to [B, S_total, H, D] then flatten to [B, S_total, inner_dim]
attn_output = mx.transpose(attn_output_bhsd, (0, 2, 1, 3))
attn_output = mx.reshape(attn_output, (batch_size, seq_total, self.inner_dim))
# Split back into context and image streams
context_attn_output = attn_output[:, :seq_ctx, :]
hidden_attn_output = attn_output[:, seq_ctx:, :]
# Output projections
hidden_attn_output = self.to_out[0](hidden_attn_output)
context_attn_output = self.to_add_out(context_attn_output)
return hidden_attn_output, context_attn_output

80
src/mflux/models/fibo/model/fibo_transformer/fibo_single_attention.py Normal file
View File

@ -0,0 +1,80 @@
import mlx.core as mx
from mlx import nn
from mlx.core.fast import scaled_dot_product_attention
class FiboSingleAttention(nn.Module):
def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, eps: float = 1e-6):
super().__init__()
self.head_dim = attention_head_dim
self.num_heads = num_attention_heads
self.inner_dim = dim
self.to_q = nn.Linear(dim, self.inner_dim)
self.to_k = nn.Linear(dim, self.inner_dim)
self.to_v = nn.Linear(dim, self.inner_dim)
self.norm_q = nn.RMSNorm(self.head_dim, eps=eps)
self.norm_k = nn.RMSNorm(self.head_dim, eps=eps)
def __call__(
self,
hidden_states: mx.array,
image_rotary_emb: tuple[mx.array, mx.array],
attention_mask: mx.array | None = None,
) -> mx.array:
batch_size, seq_len, _ = hidden_states.shape
cos, sin = image_rotary_emb
# [B, S, inner_dim]
query = self.to_q(hidden_states)
key = self.to_k(hidden_states)
value = self.to_v(hidden_states)
query = mx.reshape(query, (batch_size, seq_len, self.num_heads, self.head_dim))
key = mx.reshape(key, (batch_size, seq_len, self.num_heads, self.head_dim))
value = mx.reshape(value, (batch_size, seq_len, self.num_heads, self.head_dim))
# RMSNorm
query = self.norm_q(query.astype(mx.float32)).astype(query.dtype)
key = self.norm_k(key.astype(mx.float32)).astype(key.dtype)
# RoPE
query = FiboSingleAttention.apply_rotary_emb(query, cos, sin)
key = FiboSingleAttention.apply_rotary_emb(key, cos, sin)
# [B, H, S, D]
query_bhsd = mx.transpose(query, (0, 2, 1, 3))
key_bhsd = mx.transpose(key, (0, 2, 1, 3))
value_bhsd = mx.transpose(value, (0, 2, 1, 3))
scale = 1.0 / mx.sqrt(mx.array(self.head_dim, dtype=query_bhsd.dtype))
attn_output = scaled_dot_product_attention(
query_bhsd,
key_bhsd,
value_bhsd,
scale=scale,
mask=attention_mask,
)
attn_output = mx.transpose(attn_output, (0, 2, 1, 3))
attn_output = mx.reshape(attn_output, (batch_size, seq_len, self.inner_dim))
return attn_output
@staticmethod
def apply_rotary_emb(
x: mx.array,
freqs_cos: mx.array,
freqs_sin: mx.array,
) -> mx.array:
bsz, seq_len, num_heads, head_dim = x.shape
cos = mx.expand_dims(mx.expand_dims(freqs_cos, axis=0), axis=2)
sin = mx.expand_dims(mx.expand_dims(freqs_sin, axis=0), axis=2)
x2 = x.reshape(bsz, seq_len, num_heads, -1, 2)
x_real = x2[..., 0]
x_imag = x2[..., 1]
x_rotated_real = -x_imag
x_rotated_imag = x_real
x_rotated = mx.stack([x_rotated_real, x_rotated_imag], axis=-1).reshape(bsz, seq_len, num_heads, head_dim)
out = (x.astype(mx.float32) * cos + x_rotated.astype(mx.float32) * sin).astype(x.dtype)
return out

76
src/mflux/models/fibo/model/fibo_transformer/joint_transformer_block.py Normal file
View File

@ -0,0 +1,76 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_transformer.feed_forward import FiboFeedForward
from mflux.models.fibo.model.fibo_transformer.fibo_ada_layer_norm_zero import FiboAdaLayerNormZero
from mflux.models.fibo.model.fibo_transformer.fibo_joint_attention import FiboJointAttention
class FiboJointTransformerBlock(nn.Module):
def __init__(
self,
layer: int,
dim: int = 3072,
num_attention_heads: int = 24,
attention_head_dim: int = 128,
eps: float = 1e-6,
):
super().__init__()
self.layer = layer
self.norm1 = FiboAdaLayerNormZero(embedding_dim=dim)
self.norm1_context = FiboAdaLayerNormZero(embedding_dim=dim)
self.attn = FiboJointAttention(dim=dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim) # fmt: off
self.norm2 = nn.LayerNorm(dims=dim, eps=eps, affine=False)
self.ff = FiboFeedForward(dim=dim, dim_out=dim, mult=4, activation_fn="gelu-approximate")
self.norm2_context = nn.LayerNorm(dims=dim, eps=eps, affine=False)
self.ff_context = FiboFeedForward(dim=dim, dim_out=dim, mult=4, activation_fn="gelu-approximate")
def __call__(
self,
temb: mx.array,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
image_rotary_emb: mx.array,
attention_mask: mx.array | None = None,
) -> tuple[mx.array, mx.array]:
# 1. AdaLayerNormZero for both streams
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
hidden_states=hidden_states,
text_embeddings=temb,
)
norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
hidden_states=encoder_hidden_states,
text_embeddings=temb,
)
# 2. Joint attention over image + context streams.
attn_output, context_attn_output = self.attn(
hidden_states=norm_hidden_states,
encoder_hidden_states=norm_encoder_hidden_states,
image_rotary_emb=image_rotary_emb,
attention_mask=attention_mask,
)
# 3a. Process attention outputs for the image stream.
attn_output = mx.expand_dims(gate_msa, axis=1) * attn_output
hidden_states = hidden_states + attn_output
norm_hidden_states = self.norm2(hidden_states)
norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
ff_output = self.ff(norm_hidden_states)
ff_output = mx.expand_dims(gate_mlp, axis=1) * ff_output
hidden_states = hidden_states + ff_output
# 3b. Process attention outputs for the context stream.
context_attn_output = mx.expand_dims(c_gate_msa, axis=1) * context_attn_output
encoder_hidden_states = encoder_hidden_states + context_attn_output
norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
context_ff_output = self.ff_context(norm_encoder_hidden_states)
encoder_hidden_states = encoder_hidden_states + mx.expand_dims(c_gate_mlp, axis=1) * context_ff_output
return encoder_hidden_states, hidden_states

55
src/mflux/models/fibo/model/fibo_transformer/single_transformer_block.py Normal file
View File

@ -0,0 +1,55 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_transformer.fibo_single_attention import FiboSingleAttention
from mflux.models.flux.model.flux_transformer.ada_layer_norm_zero_single import AdaLayerNormZeroSingle
class FiboSingleTransformerBlock(nn.Module):
def __init__(
self,
layer: int,
dim: int = 3072,
num_attention_heads: int = 24,
attention_head_dim: int = 128,
mlp_ratio: float = 4.0,
):
super().__init__()
self.layer = layer
self.mlp_hidden_dim = int(dim * mlp_ratio)
self.norm = AdaLayerNormZeroSingle()
self.proj_mlp = nn.Linear(dim, self.mlp_hidden_dim)
self.act_mlp = nn.gelu_approx
self.proj_out = nn.Linear(dim + self.mlp_hidden_dim, dim)
self.attn = FiboSingleAttention(dim=dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim) # fmt: off
def __call__(
self,
temb: mx.array,
hidden_states: mx.array,
image_rotary_emb: tuple[mx.array, mx.array],
attention_mask: mx.array | None = None,
) -> mx.array:
# 0. Residual connection
residual = hidden_states
# 1. AdaLayerNormZeroSingle
norm_hidden_states, gate = self.norm(
hidden_states=hidden_states,
text_embeddings=temb,
)
# 2. Attention
attn_output = self.attn(
hidden_states=norm_hidden_states,
image_rotary_emb=image_rotary_emb,
attention_mask=attention_mask,
)
# 3. MLP + projection
mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
hidden_states = mx.concatenate([attn_output, mlp_hidden_states], axis=2)
gate = mx.expand_dims(gate, axis=1)
hidden_states = gate * self.proj_out(hidden_states)
hidden_states = residual + hidden_states
return hidden_states

10
src/mflux/models/fibo/model/fibo_transformer/text_projection.py Normal file
View File

@ -0,0 +1,10 @@
from mlx import nn
class BriaFiboTextProjection(nn.Module):
def __init__(self, in_features: int = 2048, hidden_size: int = 1536):
super().__init__()
self.linear = nn.Linear(in_features, hidden_size, bias=False)
def __call__(self, caption):
return self.linear(caption)

17
src/mflux/models/fibo/model/fibo_transformer/time_embed.py Normal file
View File

@ -0,0 +1,17 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_transformer.bria_fibo_timesteps import BriaFiboTimesteps
from mflux.models.flux.model.flux_transformer.timestep_embedder import TimestepEmbedder
class BriaFiboTimestepProjEmbeddings(nn.Module):
def __init__(self, embedding_dim: int = 3072, time_theta: int = 10000):
super().__init__()
self.time_proj = BriaFiboTimesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, time_theta=time_theta, scale=1.0) # fmt: off
self.timestep_embedder = TimestepEmbedder()
def __call__(self, timestep: mx.array, dtype) -> mx.array:
timesteps_proj = self.time_proj(timestep)
timesteps_emb = self.timestep_embedder(timesteps_proj.astype(dtype))
return timesteps_emb

225
src/mflux/models/fibo/model/fibo_transformer/transformer.py Normal file
View File

@ -0,0 +1,225 @@
import mlx.core as mx
from mlx import nn
from mflux.config.runtime_config import RuntimeConfig
from mflux.models.fibo.model.fibo_transformer.fibo_embed_nd import FiboEmbedND
from mflux.models.fibo.model.fibo_transformer.joint_transformer_block import FiboJointTransformerBlock
from mflux.models.fibo.model.fibo_transformer.single_transformer_block import FiboSingleTransformerBlock
from mflux.models.fibo.model.fibo_transformer.text_projection import BriaFiboTextProjection
from mflux.models.fibo.model.fibo_transformer.time_embed import BriaFiboTimestepProjEmbeddings
from mflux.models.flux.model.flux_transformer.ada_layer_norm_continuous import AdaLayerNormContinuous
class FiboTransformer(nn.Module):
def __init__(
self,
in_channels: int = 48,
num_layers: int = 8,
num_single_layers: int = 38,
):
super().__init__()
self.pos_embed = FiboEmbedND()
self.x_embedder = nn.Linear(in_channels, 3072)
self.time_embed = BriaFiboTimestepProjEmbeddings()
self.context_embedder = nn.Linear(4096, 3072)
self.transformer_blocks = [FiboJointTransformerBlock(i) for i in range(num_layers)]
self.single_transformer_blocks = [FiboSingleTransformerBlock(i) for i in range(num_single_layers)]
self.norm_out = AdaLayerNormContinuous(3072, 3072)
self.proj_out = nn.Linear(3072, in_channels)
self.caption_projection = [BriaFiboTextProjection() for _ in range(num_layers + num_single_layers)]
def __call__(
self,
t: int,
config: RuntimeConfig,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_encoder_layers: list[mx.array],
) -> mx.array:
# 1. Create embeddings
hidden_states = FiboTransformer._handle_classifier_free_guidance(hidden_states, encoder_hidden_states)
hidden_states = self.x_embedder(hidden_states)
encoder_hidden_states = self.context_embedder(encoder_hidden_states)
time_embeddings = FiboTransformer._compute_time_embeddings(t, config, hidden_states.shape[0], hidden_states.dtype, self.time_embed) # fmt: off
image_rotary_emb = FiboTransformer._compute_rotary_embeddings(encoder_hidden_states, self.pos_embed, config, hidden_states.dtype) # fmt: off
# 2. Compute attention mask
attention_mask = FiboTransformer._compute_attention_mask(
config=config,
batch_size=hidden_states.shape[0],
encoder_hidden_states=encoder_hidden_states,
max_tokens=encoder_hidden_states.shape[1],
)
# 3. Project the fibo-specific text encoder layers
text_encoder_layers = [
self.caption_projection[i](text_layer)
for i, text_layer in enumerate(text_encoder_layers)
] # fmt: off
# 4. Run the joint transformer blocks
block_id = 0
for _, block in enumerate(self.transformer_blocks):
encoder_hidden_states, hidden_states = FiboTransformer._apply_joint_transformer_block(
block=block,
time_embeddings=time_embeddings,
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_encoder_layer=text_encoder_layers[block_id],
image_rotary_emb=image_rotary_emb,
attention_mask=attention_mask,
)
block_id += 1
# 5. Run the single transformer blocks
for _, block in enumerate(self.single_transformer_blocks):
encoder_hidden_states, hidden_states = FiboTransformer._apply_single_transformer_block(
block=block,
time_embeddings=time_embeddings,
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
text_encoder_layer=text_encoder_layers[block_id],
image_rotary_emb=image_rotary_emb,
attention_mask=attention_mask,
)
block_id += 1
# 6. Project the final output
hidden_states = self.norm_out(hidden_states, time_embeddings)
hidden_states = self.proj_out(hidden_states)
return hidden_states
@staticmethod
def _apply_joint_transformer_block(
time_embeddings: mx.array,
block: FiboJointTransformerBlock,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_encoder_layer: mx.array,
image_rotary_emb: mx.array,
attention_mask: mx.array,
) -> tuple[mx.array, mx.array]:
encoder_hidden_states_half = encoder_hidden_states[:, :, :1536]
encoder_hidden_states = mx.concatenate([encoder_hidden_states_half, text_encoder_layer], axis=-1)
encoder_hidden_states, hidden_states = block(
temb=time_embeddings,
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
image_rotary_emb=image_rotary_emb,
attention_mask=attention_mask,
)
return encoder_hidden_states, hidden_states
@staticmethod
def _apply_single_transformer_block(
time_embeddings: mx.array,
block: FiboSingleTransformerBlock,
hidden_states: mx.array,
encoder_hidden_states: mx.array,
text_encoder_layer: mx.array,
image_rotary_emb: mx.array,
attention_mask: mx.array,
) -> tuple[mx.array, mx.array]:
encoder_hidden_states_half = encoder_hidden_states[:, :, :1536]
encoder_hidden_states = mx.concatenate([encoder_hidden_states_half, text_encoder_layer], axis=-1)
combined = mx.concatenate([encoder_hidden_states, hidden_states], axis=1)
combined = block(
temb=time_embeddings,
hidden_states=combined,
image_rotary_emb=image_rotary_emb,
attention_mask=attention_mask,
)
encoder_len = encoder_hidden_states.shape[1]
encoder_hidden_states = combined[:, :encoder_len, ...]
hidden_states = combined[:, encoder_len:, ...]
return encoder_hidden_states, hidden_states
@staticmethod
def _handle_classifier_free_guidance(hidden_states: mx.array, encoder_hidden_states: mx.array) -> mx.array:
batch_size = hidden_states.shape[0]
encoder_batch_size = encoder_hidden_states.shape[0]
if encoder_batch_size == batch_size * 2:
hidden_states = mx.concatenate([hidden_states, hidden_states], axis=0)
return hidden_states
@staticmethod
def _compute_time_embeddings(
t: int,
config: RuntimeConfig,
batch_size: int,
dtype: mx.Dtype,
time_embed: BriaFiboTimestepProjEmbeddings,
) -> mx.array:
timestep_value = config.scheduler.timesteps[t]
timestep = mx.full((batch_size,), timestep_value, dtype=dtype)
return time_embed(timestep, dtype=dtype)
@staticmethod
def _compute_rotary_embeddings(
encoder_hidden_states: mx.array,
pos_embed: FiboEmbedND,
config: RuntimeConfig,
dtype: mx.Dtype,
) -> mx.array:
max_tokens = encoder_hidden_states.shape[1]
txt_ids = mx.zeros((max_tokens, 3), dtype=dtype)
img_ids = FiboTransformer._prepare_latent_image_ids(height=config.height, width=config.width, dtype=dtype)
if txt_ids.ndim == 3 and txt_ids.shape[0] == 1:
txt_ids = txt_ids[0]
if img_ids.ndim == 3 and img_ids.shape[0] == 1:
img_ids = img_ids[0]
ids = mx.concatenate((txt_ids, img_ids), axis=0)
ids = mx.expand_dims(ids, axis=0)
return pos_embed(ids)
@staticmethod
def _prepare_latent_image_ids(height: int, width: int, dtype=mx.float32) -> mx.array:
vae_scale_factor = 16
latent_height = height // vae_scale_factor
latent_width = width // vae_scale_factor
row_indices = mx.arange(0, latent_height, dtype=dtype)[:, None]
row_indices = mx.broadcast_to(row_indices, (latent_height, latent_width))
col_indices = mx.arange(0, latent_width, dtype=dtype)[None, :]
col_indices = mx.broadcast_to(col_indices, (latent_height, latent_width))
zeros_channel = mx.zeros((latent_height, latent_width), dtype=dtype)
latent_image_ids = mx.stack([zeros_channel, row_indices, col_indices], axis=-1)
latent_image_ids = mx.reshape(latent_image_ids, (latent_height * latent_width, 3))
return latent_image_ids
@staticmethod
def _prepare_attention_mask(attention_mask_2d: mx.array) -> mx.array:
attention_matrix = mx.einsum("bi,bj->bij", attention_mask_2d, attention_mask_2d)
mask_dtype = attention_mask_2d.dtype
min_dtype_value = mx.finfo(mask_dtype).min
attention_matrix = mx.where(
attention_matrix == 1,
mx.zeros_like(attention_matrix).astype(mask_dtype),
(mx.ones_like(attention_matrix) * min_dtype_value).astype(mask_dtype),
)
attention_matrix = mx.expand_dims(attention_matrix, axis=1)
return attention_matrix
@staticmethod
def _compute_attention_mask(
batch_size: int,
config: RuntimeConfig,
encoder_hidden_states: mx.array,
max_tokens: int,
) -> mx.array:
vae_scale_factor = 16
latent_height = config.height // vae_scale_factor
latent_width = config.width // vae_scale_factor
latent_seq_len = latent_height * latent_width
prompt_attention_mask = mx.ones((batch_size, max_tokens), dtype=mx.float32)
latent_attention_mask = mx.ones((batch_size, latent_seq_len), dtype=mx.float32)
attention_mask_2d = mx.concatenate([prompt_attention_mask, latent_attention_mask], axis=1)
attention_mask = FiboTransformer._prepare_attention_mask(attention_mask_2d)
attention_mask = attention_mask.astype(encoder_hidden_states.dtype)
return attention_mask

1
src/mflux/models/fibo/model/fibo_vae/__init__.py Normal file
View File

@ -0,0 +1 @@
"""FIBO VAE components."""

17
src/mflux/models/fibo/model/fibo_vae/common/__init__.py Normal file
View File

@ -0,0 +1,17 @@
"""FIBO VAE common/shared components."""
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_attention_block import Wan2_2_AttentionBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_mid_block import Wan2_2_MidBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_resample import Wan2_2_Resample
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_residual_block import Wan2_2_ResidualBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_rms_norm import Wan2_2_RMSNorm
__all__ = [
"Wan2_2_AttentionBlock",
"Wan2_2_CausalConv3d",
"Wan2_2_MidBlock",
"Wan2_2_ResidualBlock",
"Wan2_2_Resample",
"Wan2_2_RMSNorm",
]

48
src/mflux/models/fibo/model/fibo_vae/common/wan_2_2_attention_block.py Normal file
View File

@ -0,0 +1,48 @@
import mlx.core as mx
from mlx import nn
from mlx.core.fast import scaled_dot_product_attention
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_rms_norm import Wan2_2_RMSNorm
class Wan2_2_AttentionBlock(nn.Module):
def __init__(self, dim: int):
super().__init__()
self.dim = dim
self.norm = Wan2_2_RMSNorm(dim, images=True)
self.to_qkv = nn.Conv2d(dim, dim * 3, kernel_size=1)
self.proj = nn.Conv2d(dim, dim, kernel_size=1)
def __call__(self, x: mx.array) -> mx.array:
identity = x
batch_size, channels, time, height, width = x.shape
x = mx.transpose(x, (0, 2, 1, 3, 4))
x = mx.reshape(x, (batch_size * time, channels, height, width))
x = self.norm(x)
x = mx.transpose(x, (0, 2, 3, 1))
qkv = self.to_qkv(x)
qkv = mx.transpose(qkv, (0, 3, 1, 2))
qkv = mx.reshape(qkv, (batch_size * time, 1, channels * 3, height * width))
qkv = mx.transpose(qkv, (0, 1, 3, 2))
q, k, v = mx.split(qkv, 3, axis=3)
scale = 1.0 / (channels**0.5)
x = scaled_dot_product_attention(
q,
k,
v,
scale=scale,
mask=None,
)
x = mx.reshape(x, (batch_size * time, height * width, channels))
x = mx.transpose(x, (0, 2, 1))
x = mx.reshape(x, (batch_size * time, channels, height, width))
x = mx.transpose(x, (0, 2, 3, 1))
x = self.proj(x)
x = mx.transpose(x, (0, 3, 1, 2))
x = mx.reshape(x, (batch_size, time, channels, height, width))
x = mx.transpose(x, (0, 2, 1, 3, 4))
return x + identity

42
src/mflux/models/fibo/model/fibo_vae/common/wan_2_2_causal_conv_3d.py Normal file
View File

@ -0,0 +1,42 @@
import mlx.core as mx
from mlx import nn
class Wan2_2_CausalConv3d(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int = 3,
stride: int = 1,
padding: int = 1,
name: str | None = None,
):
super().__init__()
self.conv3d = nn.Conv3d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=0,
)
self.padding = padding
self.stride = stride
self.kernel_size = kernel_size
self.name = name or f"conv3d_{in_channels}to{out_channels}"
def __call__(self, x: mx.array) -> mx.array:
pad_t = pad_h = pad_w = self.padding
if pad_t > 0 or pad_h > 0 or pad_w > 0:
pad_spec = [
(0, 0),
(0, 0),
(2 * pad_t, 0),
(pad_h, pad_h),
(pad_w, pad_w),
]
x = mx.pad(x, pad_spec)
x = mx.transpose(x, (0, 2, 3, 4, 1))
x = self.conv3d(x)
x = mx.transpose(x, (0, 4, 1, 2, 3))
return x

22
src/mflux/models/fibo/model/fibo_vae/common/wan_2_2_mid_block.py Normal file
View File

@ -0,0 +1,22 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_attention_block import Wan2_2_AttentionBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_residual_block import Wan2_2_ResidualBlock
class Wan2_2_MidBlock(nn.Module):
def __init__(self, dim: int, non_linearity: str = "silu", num_layers: int = 1):
super().__init__()
self.resnets = [Wan2_2_ResidualBlock(dim, dim, non_linearity)]
self.attentions = []
for _ in range(num_layers):
self.attentions.append(Wan2_2_AttentionBlock(dim))
self.resnets.append(Wan2_2_ResidualBlock(dim, dim, non_linearity))
def __call__(self, x: mx.array) -> mx.array:
x = self.resnets[0](x)
for attn, resnet in zip(self.attentions, self.resnets[1:]):
x = attn(x)
x = resnet(x)
return x

64
src/mflux/models/fibo/model/fibo_vae/common/wan_2_2_resample.py Normal file
View File

@ -0,0 +1,64 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d
class Wan2_2_Resample(nn.Module):
def __init__(self, dim: int, mode: str, upsample_out_dim: int = None):
super().__init__()
self.dim = dim
self.mode = mode
if upsample_out_dim is None:
upsample_out_dim = dim // 2
if mode == "upsample3d":
self.time_conv = Wan2_2_CausalConv3d(dim, dim * 2, kernel_size=(3, 1, 1), stride=1, padding=(1, 0, 0))
self.resample_conv = nn.Conv2d(dim, upsample_out_dim, kernel_size=3, stride=1, padding=1)
elif mode == "upsample2d":
self.resample_conv = nn.Conv2d(dim, upsample_out_dim, kernel_size=3, stride=1, padding=1)
self.time_conv = None
elif mode == "downsample2d":
self.resample_conv = nn.Conv2d(dim, dim, kernel_size=3, stride=2, padding=0)
self.time_conv = None
elif mode == "downsample3d":
self.resample_conv = nn.Conv2d(dim, dim, kernel_size=3, stride=2, padding=0)
self.time_conv = None
else:
raise ValueError(f"Unsupported resample mode: {mode}")
def __call__(self, x: mx.array, block_idx: int | None = None) -> mx.array:
b, c, t, h, w = x.shape
if self.mode in ("upsample2d", "upsample3d"):
if self.mode == "upsample3d" and self.time_conv is not None:
x = self.time_conv(x)
x = mx.reshape(x, (b, 2, c, t, h, w))
x = mx.transpose(x, (0, 2, 3, 1, 4, 5))
x = mx.reshape(x, (b, c, t * 2, h, w))
t = t * 2
x = mx.transpose(x, (0, 2, 1, 3, 4))
x = mx.reshape(x, (b * t, c, h, w))
x = mx.transpose(x, (0, 2, 3, 1))
x = mx.repeat(x, 2, axis=1)
x = mx.repeat(x, 2, axis=2)
x = self.resample_conv(x)
x = mx.transpose(x, (0, 3, 1, 2))
new_c = x.shape[1]
new_h, new_w = x.shape[2], x.shape[3]
x = mx.reshape(x, (b, t, new_c, new_h, new_w))
x = mx.transpose(x, (0, 2, 1, 3, 4))
return x
# downsample modes
x = mx.transpose(x, (0, 2, 1, 3, 4))
x = mx.reshape(x, (b * t, c, h, w))
x = mx.transpose(x, (0, 2, 3, 1))
x = mx.pad(x, [(0, 0), (0, 1), (0, 1), (0, 0)])
x = self.resample_conv(x)
x = mx.transpose(x, (0, 3, 1, 2))
new_c = x.shape[1]
new_h, new_w = x.shape[2], x.shape[3]
x = mx.reshape(x, (b, t, new_c, new_h, new_w))
x = mx.transpose(x, (0, 2, 1, 3, 4))
return x

35
src/mflux/models/fibo/model/fibo_vae/common/wan_2_2_residual_block.py Normal file
View File

@ -0,0 +1,35 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_rms_norm import Wan2_2_RMSNorm
class Wan2_2_ResidualBlock(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
non_linearity: str = "silu",
):
super().__init__()
self.norm1 = Wan2_2_RMSNorm(in_dim, images=False)
self.conv1 = Wan2_2_CausalConv3d(in_dim, out_dim, 3, padding=1)
self.norm2 = Wan2_2_RMSNorm(out_dim, images=False)
self.conv2 = Wan2_2_CausalConv3d(out_dim, out_dim, 3, padding=1)
if in_dim != out_dim:
self.conv_shortcut = Wan2_2_CausalConv3d(in_dim, out_dim, 1, padding=0)
else:
self.conv_shortcut = None
def __call__(self, x: mx.array, resnet_idx: int | None = None, block_idx: int | None = None) -> mx.array:
h = self.conv_shortcut(x) if self.conv_shortcut is not None else x
x = self.norm1(x)
x = nn.silu(x)
x = self.conv1(x)
x = self.norm2(x)
x = nn.silu(x)
x = self.conv2(x)
result = x + h
return result

32
src/mflux/models/fibo/model/fibo_vae/common/wan_2_2_rms_norm.py Normal file
View File

@ -0,0 +1,32 @@
import mlx.core as mx
from mlx import nn
class Wan2_2_RMSNorm(nn.Module):
def __init__(self, dim: int, eps: float = 1e-12, images: bool = True):
super().__init__()
self.eps = eps
self.scale = float(dim) ** 0.5
self.images = images
if images:
self.weight = mx.ones((dim, 1, 1))
else:
self.weight = mx.ones((dim, 1, 1, 1))
def __call__(self, x: mx.array) -> mx.array:
sum_sq = mx.sum(x * x, axis=1, keepdims=True)
l2_norm = mx.sqrt(sum_sq)
denom = mx.maximum(l2_norm, mx.array(self.eps, dtype=l2_norm.dtype))
x_normalized = x / denom
if x.ndim == 5 and not self.images:
weight = self.weight.reshape(1, -1, 1, 1, 1)
elif x.ndim == 4 and self.images:
weight = self.weight.reshape(1, -1, 1, 1)
else:
if x.ndim == 5:
weight = self.weight.reshape(1, -1, 1, 1, 1)
elif x.ndim == 4:
weight = self.weight.reshape(1, -1, 1, 1)
else:
weight = self.weight
return x_normalized * self.scale * weight

5
src/mflux/models/fibo/model/fibo_vae/decoder/__init__.py Normal file
View File

@ -0,0 +1,5 @@
"""FIBO VAE decoder components."""
from mflux.models.fibo.model.fibo_vae.decoder.wan_2_2_decoder_3d import Wan2_2_Decoder3d
__all__ = ["Wan2_2_Decoder3d"]

58
src/mflux/models/fibo/model/fibo_vae/decoder/wan_2_2_decoder_3d.py Normal file
View File

@ -0,0 +1,58 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_mid_block import Wan2_2_MidBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_rms_norm import Wan2_2_RMSNorm
from mflux.models.fibo.model.fibo_vae.decoder.wan_2_2_residual_up_block import Wan2_2_ResidualUpBlock
class Wan2_2_Decoder3d(nn.Module):
def __init__(
self,
dim: int = 256,
z_dim: int = 48,
dim_mult: list[int] = [1, 2, 4, 4],
num_res_blocks: int = 2,
temporal_upsample: list[bool] | None = None,
non_linearity: str = "silu",
out_channels: int = 12,
):
super().__init__()
self.dim = dim
self.z_dim = z_dim
self.dim_mult = dim_mult
self.num_res_blocks = num_res_blocks
self.temporal_upsample = temporal_upsample or []
dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
self.conv_in = Wan2_2_CausalConv3d(z_dim, dims[0], 3, padding=1, name="decoder_conv_in")
self.mid_block = Wan2_2_MidBlock(dims[0], non_linearity, num_layers=1)
self.up_blocks: list[Wan2_2_ResidualUpBlock] = []
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
up_flag = i != len(dim_mult) - 1
temporal_flag = (
bool(self.temporal_upsample and i < len(self.temporal_upsample) and self.temporal_upsample[i])
if up_flag
else False
)
up_block = Wan2_2_ResidualUpBlock(
in_dim=in_dim,
out_dim=out_dim,
num_res_blocks=num_res_blocks,
temporal_upsample=temporal_flag,
up_flag=up_flag,
non_linearity=non_linearity,
)
self.up_blocks.append(up_block)
self.norm_out = Wan2_2_RMSNorm(out_dim, images=False)
self.conv_out = Wan2_2_CausalConv3d(out_dim, out_channels, 3, padding=1, name="decoder_conv_out")
def __call__(self, x: mx.array) -> mx.array:
x = self.conv_in(x)
x = self.mid_block(x)
for i, up_block in enumerate(self.up_blocks):
x = up_block(x, block_idx=i, first_chunk=True)
x = self.norm_out(x)
x = nn.silu(x)
x = self.conv_out(x)
return x

53
src/mflux/models/fibo/model/fibo_vae/decoder/wan_2_2_dup_up_3d.py Normal file
View File

@ -0,0 +1,53 @@
import mlx.core as mx
from mlx import nn
class Wan2_2_DupUp3D(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
factor_t: int,
factor_s: int = 1,
):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.factor_t = factor_t
self.factor_s = factor_s
self.factor = self.factor_t * self.factor_s * self.factor_s
self.repeats = out_channels * self.factor // in_channels
def __call__(self, x: mx.array, first_chunk: bool = False) -> mx.array:
b, c, t, h, w = x.shape
x = mx.repeat(x, self.repeats, axis=1)
x = mx.reshape(
x,
(
b,
self.out_channels,
self.factor_t,
self.factor_s,
self.factor_s,
t,
h,
w,
),
)
x = mx.transpose(x, (0, 1, 5, 2, 6, 3, 7, 4))
x = mx.reshape(
x,
(
b,
self.out_channels,
t * self.factor_t,
h * self.factor_s,
w * self.factor_s,
),
)
if first_chunk and self.factor_t > 1:
x = x[:, :, self.factor_t - 1 :, :, :]
return x

60
src/mflux/models/fibo/model/fibo_vae/decoder/wan_2_2_residual_up_block.py Normal file
View File

@ -0,0 +1,60 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_resample import Wan2_2_Resample
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_residual_block import Wan2_2_ResidualBlock
from mflux.models.fibo.model.fibo_vae.decoder.wan_2_2_dup_up_3d import Wan2_2_DupUp3D
class Wan2_2_ResidualUpBlock(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
num_res_blocks: int,
temporal_upsample: bool = False,
up_flag: bool = False,
non_linearity: str = "silu",
):
super().__init__()
self.in_dim = in_dim
self.out_dim = out_dim
if up_flag:
self.avg_shortcut: Wan2_2_DupUp3D | None = Wan2_2_DupUp3D(
in_channels=in_dim,
out_channels=out_dim,
factor_t=2 if temporal_upsample else 1,
factor_s=2,
)
else:
self.avg_shortcut = None
resnets: list[Wan2_2_ResidualBlock] = []
current_dim = in_dim
for _ in range(num_res_blocks + 1):
resnets.append(Wan2_2_ResidualBlock(current_dim, out_dim, non_linearity))
current_dim = out_dim
self.resnets = resnets
if up_flag:
upsample_mode = "upsample3d" if temporal_upsample else "upsample2d"
self.upsampler: Wan2_2_Resample | None = Wan2_2_Resample(
out_dim, mode=upsample_mode, upsample_out_dim=out_dim
)
else:
self.upsampler = None
def __call__(self, x: mx.array, block_idx: int | None = None, first_chunk: bool = False) -> mx.array:
x_copy = x
for i, resnet in enumerate(self.resnets):
x = resnet(x, resnet_idx=i, block_idx=block_idx)
if self.upsampler is not None:
x = self.upsampler(x, block_idx=block_idx)
if self.avg_shortcut is not None:
x = x + self.avg_shortcut(x_copy, first_chunk=first_chunk)
return x

34
src/mflux/models/fibo/model/fibo_vae/decoder/wan_2_2_up_block.py Normal file
View File

@ -0,0 +1,34 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_resample import Wan2_2_Resample
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_residual_block import Wan2_2_ResidualBlock
class Wan2_2_UpBlock(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
num_res_blocks: int,
upsample_mode: str | None = None,
non_linearity: str = "silu",
):
super().__init__()
self.resnets: list[Wan2_2_ResidualBlock] = []
current_dim = in_dim
for _ in range(num_res_blocks + 1):
self.resnets.append(Wan2_2_ResidualBlock(current_dim, out_dim, non_linearity))
current_dim = out_dim
self.upsampler: Wan2_2_Resample | None = None
if upsample_mode is not None:
self.upsampler = Wan2_2_Resample(out_dim, mode=upsample_mode, upsample_out_dim=out_dim)
def __call__(self, x: mx.array, block_idx: int | None = None) -> mx.array:
for i, resnet in enumerate(self.resnets):
x = resnet(x, resnet_idx=i, block_idx=block_idx)
if self.upsampler is not None:
x = self.upsampler(x, block_idx=block_idx)
return x

5
src/mflux/models/fibo/model/fibo_vae/encoder/__init__.py Normal file
View File

@ -0,0 +1,5 @@
"""FIBO VAE encoder components."""
from mflux.models.fibo.model.fibo_vae.encoder.wan_2_2_encoder_3d import Wan2_2_Encoder3d
__all__ = ["Wan2_2_Encoder3d"]

75
src/mflux/models/fibo/model/fibo_vae/encoder/wan_2_2_avg_down_3d.py Normal file
View File

@ -0,0 +1,75 @@
import mlx.core as mx
from mlx import nn
class Wan2_2_AvgDown3D(nn.Module):
def __init__(
self,
in_channels: int,
out_channels: int,
factor_t: int,
factor_s: int = 1,
):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.factor_t = factor_t
self.factor_s = factor_s
self.factor = self.factor_t * self.factor_s * self.factor_s
assert in_channels * self.factor % out_channels == 0
self.group_size = in_channels * self.factor // out_channels
def __call__(self, x: mx.array) -> mx.array:
pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
if pad_t > 0:
x = mx.pad(
x,
[
(0, 0), # batch
(0, 0), # channels
(pad_t, 0), # time
(0, 0), # height
(0, 0), # width
],
)
b, c, t, h, w = x.shape
x = mx.reshape(
x,
(
b,
c,
t // self.factor_t,
self.factor_t,
h // self.factor_s,
self.factor_s,
w // self.factor_s,
self.factor_s,
),
)
x = mx.transpose(x, (0, 1, 3, 5, 7, 2, 4, 6))
x = mx.reshape(
x,
(
b,
c * self.factor,
t // self.factor_t,
h // self.factor_s,
w // self.factor_s,
),
)
x = mx.reshape(
x,
(
b,
self.out_channels,
self.group_size,
t // self.factor_t,
h // self.factor_s,
w // self.factor_s,
),
)
x = mx.mean(x, axis=2)
return x

57
src/mflux/models/fibo/model/fibo_vae/encoder/wan_2_2_down_block.py Normal file
View File

@ -0,0 +1,57 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_attention_block import Wan2_2_AttentionBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_resample import Wan2_2_Resample
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_residual_block import Wan2_2_ResidualBlock
from mflux.models.fibo.model.fibo_vae.encoder.wan_2_2_avg_down_3d import Wan2_2_AvgDown3D
class Wan2_2_DownBlock(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
num_res_blocks: int,
attn_scales: list[float] | None = None,
scale: float = 1.0,
temporal_downsample: bool = False,
non_linearity: str = "silu",
is_last: bool = False,
):
super().__init__()
if attn_scales is None:
attn_scales = []
resnets: list[nn.Module] = []
current_dim = in_dim
for _ in range(num_res_blocks):
resnets.append(Wan2_2_ResidualBlock(current_dim, out_dim, non_linearity))
if scale in attn_scales:
resnets.append(Wan2_2_AttentionBlock(out_dim))
current_dim = out_dim
self.resnets = resnets
# Shortcut path with downsample (mirrors AvgDown3D in diffusers)
self.avg_shortcut = Wan2_2_AvgDown3D(
in_dim,
out_dim,
factor_t=2 if temporal_downsample else 1,
factor_s=2 if not is_last else 1,
)
# Main path downsampler
if not is_last:
mode = "downsample3d" if temporal_downsample else "downsample2d"
self.downsampler = Wan2_2_Resample(out_dim, mode=mode)
else:
self.downsampler = None
def __call__(self, x: mx.array) -> mx.array:
x_copy = x
for layer in self.resnets:
x = layer(x)
if self.downsampler is not None:
x = self.downsampler(x)
return x + self.avg_shortcut(x_copy)

75
src/mflux/models/fibo/model/fibo_vae/encoder/wan_2_2_encoder_3d.py Normal file
View File

@ -0,0 +1,75 @@
import mlx.core as mx
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_mid_block import Wan2_2_MidBlock
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_rms_norm import Wan2_2_RMSNorm
from mflux.models.fibo.model.fibo_vae.encoder.wan_2_2_down_block import Wan2_2_DownBlock
class Wan2_2_Encoder3d(nn.Module):
def __init__(
self,
in_channels: int = 3,
dim: int = 128,
z_dim: int = 4,
dim_mult: list[int] | None = None,
num_res_blocks: int = 2,
attn_scales: list[float] | None = None,
temporal_downsample: list[bool] | None = None,
non_linearity: str = "silu",
is_residual: bool = False,
):
super().__init__()
if dim_mult is None:
dim_mult = [1, 2, 4, 4]
if attn_scales is None:
attn_scales = []
if temporal_downsample is None:
temporal_downsample = [False, True, True]
if is_residual:
raise NotImplementedError("Residual down blocks are not implemented for Wan2_2_Encoder3d in MLX.")
self.dim = dim
self.z_dim = z_dim
self.dim_mult = dim_mult
self.num_res_blocks = num_res_blocks
self.attn_scales = attn_scales
dims = [dim * u for u in [1] + dim_mult]
self.temporal_downsample = temporal_downsample
self.conv_in = Wan2_2_CausalConv3d(in_channels, dims[0], 3, padding=1)
scale = 1.0
self.down_blocks: list[Wan2_2_DownBlock] = []
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
block = Wan2_2_DownBlock(
in_dim=in_dim,
out_dim=out_dim,
num_res_blocks=num_res_blocks,
attn_scales=attn_scales,
scale=scale,
temporal_downsample=temporal_downsample[i] if i < len(temporal_downsample) else False,
non_linearity=non_linearity,
is_last=i == len(dim_mult) - 1,
)
self.down_blocks.append(block)
if i != len(dim_mult) - 1:
scale /= 2.0
self.mid_block = Wan2_2_MidBlock(out_dim, non_linearity, num_layers=1)
self.norm_out = Wan2_2_RMSNorm(out_dim, images=False)
self.conv_out = Wan2_2_CausalConv3d(out_dim, z_dim, 3, padding=1)
def __call__(self, x: mx.array) -> mx.array:
x = self.conv_in(x)
for block in self.down_blocks:
x = block(x)
x = self.mid_block(x)
x = self.norm_out(x)
x = nn.silu(x)
x = self.conv_out(x)
return x

111
src/mflux/models/fibo/model/fibo_vae/wan_2_2_vae.py Normal file
View File

@ -0,0 +1,111 @@
import mlx.core as mx
import numpy as np
from mlx import nn
from mflux.models.fibo.model.fibo_vae.common.wan_2_2_causal_conv_3d import Wan2_2_CausalConv3d
from mflux.models.fibo.model.fibo_vae.decoder.wan_2_2_decoder_3d import Wan2_2_Decoder3d
from mflux.models.fibo.model.fibo_vae.encoder.wan_2_2_encoder_3d import Wan2_2_Encoder3d
class Wan2_2_VAE(nn.Module):
Z_DIM = 48
ENCODER_BASE_DIM = 160
DECODER_BASE_DIM = 256
DIM_MULT = [1, 2, 4, 4]
NUM_RES_BLOCKS = 2
OUT_CHANNELS = 12
LATENTS_MEAN = np.array([-0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, -0.1382, 0.0542, 0.2813, 0.0891, 0.157, -0.0098, 0.0375, -0.1825, -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.123, -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.052, 0.3748, 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667], dtype=np.float32) # fmt: off
LATENTS_STD = np.array([0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.499, 0.4818, 0.5013, 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, 0.3971, 1.06, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744], dtype=np.float32) # fmt: off
def __init__(self):
super().__init__()
self.encoder = Wan2_2_Encoder3d(
in_channels=3,
dim=self.ENCODER_BASE_DIM,
z_dim=self.Z_DIM * 2,
dim_mult=self.DIM_MULT,
num_res_blocks=self.NUM_RES_BLOCKS,
attn_scales=[],
temporal_downsample=[False, True, True],
)
self.quant_conv = Wan2_2_CausalConv3d(self.Z_DIM * 2, self.Z_DIM * 2, 1, padding=0, name="quant_conv")
self.decoder = Wan2_2_Decoder3d(
dim=self.DECODER_BASE_DIM,
z_dim=self.Z_DIM,
dim_mult=self.DIM_MULT,
num_res_blocks=self.NUM_RES_BLOCKS,
temporal_upsample=None,
out_channels=self.OUT_CHANNELS,
)
self.post_quant_conv = Wan2_2_CausalConv3d(self.Z_DIM, self.Z_DIM, 1, padding=0, name="post_quant_conv")
def encode(self, images: mx.array) -> mx.array:
if images.ndim == 4:
x = images.reshape(images.shape[0], images.shape[1], 1, images.shape[2], images.shape[3])
elif images.ndim == 5:
x = images
else:
raise ValueError(f"Expected 4D or 5D input for VAE.encode, got shape {images.shape}")
patch_size = 2
x = self._patchify(x, patch_size=patch_size)
h = self.encoder(x)
h = self.quant_conv(h)
mean = h[:, : self.Z_DIM, :, :, :]
latents_mean = mx.array(self.LATENTS_MEAN).reshape(1, self.Z_DIM, 1, 1, 1)
latents_std = mx.array(self.LATENTS_STD).reshape(1, self.Z_DIM, 1, 1, 1)
encoded = (mean - latents_mean) / latents_std
return encoded
def decode(self, latents: mx.array) -> mx.array:
if latents.ndim == 4:
latents = latents.reshape(latents.shape[0], latents.shape[1], 1, latents.shape[2], latents.shape[3])
latents_mean = mx.array(self.LATENTS_MEAN).reshape(1, self.Z_DIM, 1, 1, 1)
latents_std = mx.array(self.LATENTS_STD).reshape(1, self.Z_DIM, 1, 1, 1)
latents = latents * latents_std + latents_mean
latents = self.post_quant_conv(latents)
decoded = self.decoder(latents)
patch_size = 2
decoded = self._unpatchify(decoded, patch_size=patch_size)
if decoded.shape[2] == 1:
decoded = decoded[:, :, 0, :, :]
return decoded
@staticmethod
def _patchify(x: mx.array, patch_size: int) -> mx.array:
if patch_size == 1:
return x
batch_size, channels, frames, height, width = x.shape
x = mx.reshape(
x,
(
batch_size,
channels,
frames,
height // patch_size,
patch_size,
width // patch_size,
patch_size,
),
)
x = mx.transpose(x, (0, 1, 6, 4, 2, 3, 5))
x = mx.reshape(
x,
(batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size),
)
return x
@staticmethod
def _unpatchify(x: mx.array, patch_size: int) -> mx.array:
if patch_size == 1:
return x
batch_size, c_patches, frames, height, width = x.shape
channels = c_patches // (patch_size * patch_size)
x = mx.reshape(x, (batch_size, channels, patch_size, patch_size, frames, height, width))
x = mx.transpose(x, (0, 1, 4, 5, 3, 6, 2))
x = mx.reshape(x, (batch_size, channels, frames, height * patch_size, width * patch_size))
return x

7
src/mflux/models/fibo/tokenizer/__init__.py Normal file
View File

@ -0,0 +1,7 @@
from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo
from mflux.models.fibo.tokenizer.smol_lm3_3b_tokenizer import FiboTokenizerHandler
__all__ = [
"FiboTokenizerHandler",
"TokenizerFibo",
]

37
src/mflux/models/fibo/tokenizer/fibo_tokenizer.py Normal file
View File

@ -0,0 +1,37 @@
import mlx.core as mx
import numpy as np
from transformers import PreTrainedTokenizer
class TokenizerFibo:
def __init__(self, tokenizer: PreTrainedTokenizer, bot_token_id: int = 128000):
self.tokenizer = tokenizer
self.bot_token_id = bot_token_id
def tokenize(
self,
prompts: list[str],
max_length: int = 2048,
padding: str = "longest",
truncation: bool = True,
add_special_tokens: bool = True,
) -> tuple[mx.array, mx.array]:
prompts = [p if p is not None else "" for p in prompts]
if all(p == "" for p in prompts):
batch_size = len(prompts)
input_ids_mx = mx.array(np.empty((batch_size, 0), dtype=np.int32))
attention_mask_mx = mx.array(np.empty((batch_size, 0), dtype=np.int32))
else:
tokenized = self.tokenizer(
prompts,
padding=padding,
max_length=max_length,
truncation=truncation,
add_special_tokens=add_special_tokens,
return_tensors="mlx",
)
input_ids_mx = tokenized["input_ids"]
attention_mask_mx = tokenized["attention_mask"]
return input_ids_mx, attention_mask_mx

14
src/mflux/models/fibo/tokenizer/qwen2vl_image_processor.py Normal file
View File

@ -0,0 +1,14 @@
from mflux.models.qwen.tokenizer.qwen_image_processor import QwenImageProcessor
class Qwen2VLImageProcessor(QwenImageProcessor):
def __init__(self):
super().__init__(
min_pixels=256 * 28 * 28,
max_pixels=1024 * 28 * 28,
patch_size=16,
temporal_patch_size=2,
merge_size=2,
image_mean=[0.5, 0.5, 0.5],
image_std=[0.5, 0.5, 0.5],
)

205
src/mflux/models/fibo/tokenizer/qwen2vl_processor.py Normal file
View File

@ -0,0 +1,205 @@
from typing import Optional, Union
import numpy as np
from PIL import Image
from mflux.models.fibo.tokenizer.qwen2vl_image_processor import Qwen2VLImageProcessor
class Qwen2VLProcessor:
def __init__(self, tokenizer):
self.tokenizer = tokenizer
self.image_processor = Qwen2VLImageProcessor()
def apply_chat_template(
self,
messages,
tokenize: bool = True,
add_generation_prompt: bool = False,
return_tensors: Optional[str] = None,
return_dict: bool = True,
**kwargs,
):
formatted = self.tokenizer.apply_chat_template(
messages,
tokenize=tokenize,
add_generation_prompt=add_generation_prompt,
return_tensors=return_tensors,
return_dict=return_dict,
**kwargs,
)
if tokenize and return_dict:
if return_tensors == "pt":
import torch
if isinstance(formatted, dict):
result = {}
for key, value in formatted.items():
if isinstance(value, torch.Tensor):
result[key] = value.numpy()
else:
result[key] = value
return result
elif return_tensors == "np" or return_tensors is None:
if isinstance(formatted, dict):
result = {}
for key, value in formatted.items():
if hasattr(value, "numpy"):
result[key] = value.numpy()
elif isinstance(value, (list, tuple)):
result[key] = np.array(value)
else:
result[key] = value
return result
return formatted
def __call__(
self,
text: Optional[Union[str, list[str]]] = None,
images: Optional[Union[Image.Image, list[Image.Image]]] = None,
padding: bool = True,
return_tensors: Optional[str] = None,
**kwargs,
):
result = {}
if images is not None:
pixel_values, image_grid_thw = self.image_processor.preprocess(images)
result["pixel_values"] = pixel_values
result["image_grid_thw"] = image_grid_thw
if text is not None:
if not isinstance(text, list):
text = [text]
text_inputs = self.tokenizer(
text,
padding=padding,
return_tensors=return_tensors or "np",
**kwargs,
)
# If images are provided, replace <|image_pad|> tokens with actual image tokens
# The number of image tokens is calculated as: prod(image_grid_thw) // merge_size^2
if images is not None and "input_ids" in text_inputs:
image_token_id = 151655 # <|image_pad|> token ID
input_ids = text_inputs["input_ids"]
# Calculate number of image tokens per image based on image_grid_thw and merge_size
# This matches transformers behavior: num_tokens = prod(image_grid_thw) // merge_size^2
merge_size = self.image_processor.merge_size
merge_length = merge_size**2
num_images = len(images) if isinstance(images, list) else 1
# Get image_grid_thw from result (already computed above)
image_grid_thw = result["image_grid_thw"]
# Calculate tokens per image
if return_tensors == "pt":
import torch
if isinstance(image_grid_thw, torch.Tensor):
image_grid_thw_np = image_grid_thw.cpu().numpy()
else:
image_grid_thw_np = image_grid_thw
else:
if isinstance(image_grid_thw, np.ndarray):
image_grid_thw_np = image_grid_thw
else:
image_grid_thw_np = np.array(image_grid_thw)
# Calculate tokens for each image
num_image_tokens_per_image_list = []
for i in range(num_images):
if i < len(image_grid_thw_np):
grid = image_grid_thw_np[i]
num_tokens = int(np.prod(grid)) // merge_length
num_image_tokens_per_image_list.append(num_tokens)
else:
num_image_tokens_per_image_list.append(256) # fallback
# Replace image pad tokens with actual image tokens
if return_tensors == "pt":
import torch
new_input_ids_list = []
for seq in input_ids:
new_seq = []
image_idx = 0
for token_id in seq:
if token_id == image_token_id and image_idx < num_images:
# Replace with calculated number of image tokens for this image
num_tokens = (
num_image_tokens_per_image_list[image_idx]
if image_idx < len(num_image_tokens_per_image_list)
else 256
)
new_seq.extend([image_token_id] * num_tokens)
image_idx += 1
else:
new_seq.append(token_id)
new_input_ids_list.append(new_seq)
# Pad sequences to same length
max_len = max(len(seq) for seq in new_input_ids_list)
padded_input_ids = []
padded_attention_mask = []
for seq in new_input_ids_list:
pad_len = max_len - len(seq)
padded_seq = (
seq + [self.tokenizer.pad_token_id] * pad_len
if self.tokenizer.pad_token_id is not None
else seq + [0] * pad_len
)
padded_input_ids.append(padded_seq)
padded_attention_mask.append([1] * len(seq) + [0] * pad_len)
text_inputs["input_ids"] = torch.tensor(padded_input_ids)
if "attention_mask" in text_inputs:
text_inputs["attention_mask"] = torch.tensor(padded_attention_mask)
else:
# NumPy version
new_input_ids_list = []
for seq in input_ids:
new_seq = []
image_idx = 0
for token_id in seq:
if token_id == image_token_id and image_idx < num_images:
# Replace with calculated number of image tokens for this image
num_tokens = (
num_image_tokens_per_image_list[image_idx]
if image_idx < len(num_image_tokens_per_image_list)
else 256
)
new_seq.extend([image_token_id] * num_tokens)
image_idx += 1
else:
new_seq.append(token_id)
new_input_ids_list.append(new_seq)
# Pad sequences to same length
max_len = max(len(seq) for seq in new_input_ids_list)
padded_input_ids = []
padded_attention_mask = []
for seq in new_input_ids_list:
pad_len = max_len - len(seq)
padded_seq = (
seq
+ [self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0] * pad_len
)
padded_input_ids.append(padded_seq)
padded_attention_mask.append([1] * len(seq) + [0] * pad_len)
text_inputs["input_ids"] = np.array(padded_input_ids)
if "attention_mask" in text_inputs:
text_inputs["attention_mask"] = np.array(padded_attention_mask)
if return_tensors == "pt":
result["input_ids"] = text_inputs["input_ids"]
result["attention_mask"] = text_inputs.get("attention_mask")
else:
result["input_ids"] = np.array(text_inputs["input_ids"])
if "attention_mask" in text_inputs:
result["attention_mask"] = np.array(text_inputs["attention_mask"])
return result

42
src/mflux/models/fibo/tokenizer/smol_lm3_3b_tokenizer.py Normal file
View File

@ -0,0 +1,42 @@
from pathlib import Path
import transformers
from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo
from mflux.utils.download import snapshot_download
class FiboTokenizerHandler:
def __init__(
self,
repo_id: str,
bot_token_id: int = 128000,
local_path: str | None = None,
):
root_path = Path(local_path) if local_path else FiboTokenizerHandler._download_or_get_cached_tokenizer(repo_id)
# Try different possible tokenizer paths
tokenizer_path = root_path / "tokenizer"
if not tokenizer_path.exists():
tokenizer_path = root_path / "text_encoder"
if not tokenizer_path.exists():
tokenizer_path = root_path
# Load the raw tokenizer
self.tokenizer_raw = transformers.AutoTokenizer.from_pretrained(
pretrained_model_name_or_path=str(tokenizer_path),
local_files_only=True,
fix_mistral_regex=True, # Fix Mistral regex pattern warning
)
# Wrap in our TokenizerFibo class
self.fibo = TokenizerFibo(self.tokenizer_raw, bot_token_id=bot_token_id)
@staticmethod
def _download_or_get_cached_tokenizer(repo_id: str) -> Path:
return Path(
snapshot_download(
repo_id=repo_id,
allow_patterns=["tokenizer/**", "text_encoder/**"],
)
)

1
src/mflux/models/fibo/variants/__init__.py Normal file
View File

@ -0,0 +1 @@
"""FIBO model variants."""

1
src/mflux/models/fibo/variants/txt2img/__init__.py Normal file
View File

@ -0,0 +1 @@
"""FIBO text-to-image variant."""

187
src/mflux/models/fibo/variants/txt2img/fibo.py Normal file
View File

@ -0,0 +1,187 @@
import mlx.core as mx
from mlx import nn
from tqdm import tqdm
from mflux.callbacks.callbacks import Callbacks
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.config.runtime_config import RuntimeConfig
from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator
from mflux.models.common.weights.model_saver import ModelSaver
from mflux.models.fibo.fibo_initializer import FIBOInitializer
from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
from mflux.models.fibo.model.fibo_text_encoder.prompt_encoder import PromptEncoder
from mflux.models.fibo.model.fibo_text_encoder.smol_lm3_3b_text_encoder import SmolLM3_3B_TextEncoder
from mflux.models.fibo.model.fibo_transformer import FiboTransformer
from mflux.models.fibo.model.fibo_vae.wan_2_2_vae import Wan2_2_VAE
from mflux.models.fibo.tokenizer.fibo_tokenizer import TokenizerFibo
from mflux.utils.exceptions import StopImageGenerationException
from mflux.utils.generated_image import GeneratedImage
from mflux.utils.image_util import ImageUtil
class FIBO(nn.Module):
vae: Wan2_2_VAE
transformer: FiboTransformer
text_encoder: SmolLM3_3B_TextEncoder
fibo_tokenizer: TokenizerFibo
def __init__(
self,
model_config: ModelConfig,
quantize: int | None = None,
local_path: str | None = None,
):
super().__init__()
self.model_config = model_config
self.bits = quantize
self.local_path = local_path
FIBOInitializer.init(
fibo_model=self,
model_config=model_config,
quantize=quantize,
local_path=local_path,
)
def generate_image(
self,
seed: int,
prompt: str,
config: Config,
negative_prompt: str | None = None,
) -> GeneratedImage:
# 0. Create a new runtime config based on the model type and input parameters
runtime_config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(runtime_config.init_time_step, runtime_config.num_inference_steps))
# 1. Create the initial latents
latents = LatentCreator.create_for_txt2img_or_img2img(
seed=seed,
height=runtime_config.height,
width=runtime_config.width,
img2img=Img2Img(
vae=self.vae,
latent_creator=FiboLatentCreator,
image_path=runtime_config.image_path,
sigmas=runtime_config.scheduler.sigmas,
init_time_step=runtime_config.init_time_step,
),
)
# 2. Encode the prompt
json_prompt, encoder_hidden_states, text_encoder_layers = PromptEncoder.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
tokenizer=self.fibo_tokenizer,
text_encoder=self.text_encoder,
)
# (Optional) Call subscribers for beginning of loop
Callbacks.before_loop(
seed=seed,
prompt=json_prompt,
latents=latents,
config=runtime_config,
)
for t in time_steps:
try:
# 3.t Predict the noise
noise_pred = self.transformer(
t=t,
config=runtime_config,
hidden_states=latents,
encoder_hidden_states=encoder_hidden_states,
text_encoder_layers=text_encoder_layers,
)
noise_pred = FIBO._apply_classifier_free_guidance(noise_pred, runtime_config.guidance)
# 4.t Take one denoise step
latents = runtime_config.scheduler.step(
model_output=noise_pred,
timestep=t,
sample=latents,
)
# (Optional) Call subscribers in-loop
Callbacks.in_loop(
t=t,
seed=seed,
prompt=json_prompt,
latents=latents,
config=runtime_config,
time_steps=time_steps,
)
# (Optional) Evaluate to enable progress tracking
mx.eval(latents)
except KeyboardInterrupt: # noqa: PERF203
Callbacks.interruption(
t=t,
seed=seed,
prompt=json_prompt,
latents=latents,
config=runtime_config,
time_steps=time_steps,
)
raise StopImageGenerationException(
f"Stopping image generation at step {t + 1}/{runtime_config.num_inference_steps}"
)
# (Optional) Call subscribers after loop
Callbacks.after_loop(
seed=seed,
prompt=json_prompt,
latents=latents,
config=runtime_config,
)
# 5. Decode the latent array and return the image
latents = FIBO._unpack_latents(latents, runtime_config.height, runtime_config.width)
decoded = self.vae.decode(latents)
return ImageUtil.to_image(
decoded_latents=decoded,
config=runtime_config,
seed=seed,
prompt=json_prompt,
quantization=self.bits,
lora_paths=None,
lora_scales=None,
image_path=runtime_config.image_path,
image_strength=runtime_config.image_strength,
generation_time=time_steps.format_dict["elapsed"],
)
@staticmethod
def _apply_classifier_free_guidance(noise_pred: mx.array, guidance: float) -> mx.array:
half = noise_pred.shape[0] // 2
noise_uncond = noise_pred[:half]
noise_text = noise_pred[half:]
return noise_uncond + guidance * (noise_text - noise_uncond)
@staticmethod
def _unpack_latents(latents: mx.array, height: int, width: int) -> mx.array:
batch_size, seq_len, channels = latents.shape
vae_scale_factor = 16
latent_height = height // vae_scale_factor
latent_width = width // vae_scale_factor
latents = mx.reshape(latents, (batch_size, latent_height, latent_width, channels))
latents = mx.transpose(latents, (0, 3, 1, 2))
return latents
def save_model(self, base_path: str) -> None:
ModelSaver.save_model(
model=self,
bits=self.bits,
base_path=base_path,
tokenizers=[
("fibo_tokenizer.tokenizer", "tokenizer"),
],
components=[
("vae", "vae"),
("transformer", "transformer"),
("text_encoder", "text_encoder"),
],
)

7
src/mflux/models/fibo/weights/__init__.py Normal file
View File

@ -0,0 +1,7 @@
"""FIBO weight handling."""
from mflux.models.fibo.weights.fibo_weight_handler import FIBOWeightHandler
from mflux.models.fibo.weights.fibo_weight_mapping import FIBOWeightMapping
from mflux.models.fibo.weights.fibo_weight_util import FIBOWeightUtil
__all__ = ["FIBOWeightHandler", "FIBOWeightMapping", "FIBOWeightUtil"]

179
src/mflux/models/fibo/weights/fibo_weight_handler.py Normal file
View File

@ -0,0 +1,179 @@
from pathlib import Path
import mlx.core as mx
import torch
from mlx.utils import tree_unflatten
from safetensors.torch import load_file as torch_load_file
from mflux.models.common.weights.mapping.weight_mapper import WeightMapper
from mflux.models.fibo.weights.fibo_weight_mapping import FIBOWeightMapping
from mflux.models.flux.weights.weight_handler import (
MetaData,
WeightHandler as FluxWeightHandler,
)
class FIBOWeightHandler:
def __init__(
self,
meta_data: MetaData,
vae: dict | None = None,
transformer: dict | None = None,
text_encoder: dict | None = None,
decoder: dict | None = None,
visual: dict | None = None,
config: dict | None = None,
):
self.vae = vae
self.transformer = transformer
self.text_encoder = text_encoder
self.decoder = decoder
self.visual = visual
self.config = config
self.meta_data = meta_data
@staticmethod
def load_regular_weights(
repo_id: str | None = None,
local_path: str | None = None,
) -> "FIBOWeightHandler":
root_path: Path | None = None
if local_path:
root_path = Path(local_path)
elif repo_id:
root_path = FluxWeightHandler.download_or_get_cached_weights(repo_id)
vae_weights = None
transformer_weights = None
text_encoder_weights = None
quantization_level: int | None = None
mflux_version: str | None = None
if root_path is not None:
vae_weights, _, _ = FIBOWeightHandler._try_load_saved_component(root_path, "vae")
text_encoder_weights, _, _ = FIBOWeightHandler._try_load_saved_component(root_path, "text_encoder") # fmt: off
transformer_weights, quantization_level, mflux_version = FIBOWeightHandler._try_load_saved_component(root_path, "transformer") # fmt: off
if vae_weights is None or transformer_weights is None or text_encoder_weights is None:
vae_weights = FIBOWeightHandler._load_vae_weights(repo_id, local_path)
transformer_weights = FIBOWeightHandler._load_transformer_weights(repo_id, local_path)
text_encoder_weights = FIBOWeightHandler._load_text_encoder_weights(repo_id, local_path)
quantization_level = None
mflux_version = None
return FIBOWeightHandler(
vae=vae_weights,
transformer=transformer_weights,
text_encoder=text_encoder_weights,
meta_data=MetaData(
quantization_level=quantization_level,
scale=None,
is_lora=False,
mflux_version=mflux_version,
),
)
@staticmethod
def _load_vae_weights(
repo_id: str | None = None,
local_path: str | None = None,
) -> dict:
root_path = FIBOWeightHandler._get_root_path(repo_id, local_path)
vae_path = root_path / "vae"
raw_weights = FIBOWeightHandler._load_safetensors_shards(vae_path)
mapping = FIBOWeightMapping.get_vae_mapping()
mapped_weights = WeightMapper.apply_mapping(raw_weights, mapping, num_blocks=4)
return mapped_weights
@staticmethod
def _load_transformer_weights(
repo_id: str | None = None,
local_path: str | None = None,
) -> dict:
root_path = FIBOWeightHandler._get_root_path(repo_id, local_path)
transformer_path = root_path / "transformer"
if transformer_path.exists() and list(transformer_path.glob("*.safetensors")):
raw_weights = FIBOWeightHandler._load_safetensors_shards(transformer_path)
else:
raw_weights = FIBOWeightHandler._load_safetensors_shards(root_path)
mapping = FIBOWeightMapping.get_transformer_mapping()
mapped_weights = WeightMapper.apply_mapping(
raw_weights,
mapping,
num_blocks=38,
num_layers=46,
)
return mapped_weights
@staticmethod
def _load_text_encoder_weights(
repo_id: str | None = None,
local_path: str | None = None,
) -> dict:
root_path = FIBOWeightHandler._get_root_path(repo_id, local_path)
text_encoder_path = root_path / "text_encoder"
raw_weights = FIBOWeightHandler._load_safetensors_shards(text_encoder_path)
mapping = FIBOWeightMapping.get_text_encoder_mapping()
mapped_weights = WeightMapper.apply_mapping(
raw_weights,
mapping,
num_blocks=36,
)
return mapped_weights
@staticmethod
def _get_root_path(
repo_id: str | None = None,
local_path: str | None = None,
) -> Path:
if local_path:
return Path(local_path)
return Path(FluxWeightHandler.download_or_get_cached_weights(repo_id or "briaai/FIBO"))
@staticmethod
def _load_safetensors_shards(path: Path) -> dict[str, mx.array]:
shard_files = sorted(f for f in path.glob("*.safetensors") if not f.name.startswith("._"))
if not shard_files:
raise FileNotFoundError(f"No safetensors files found in {path}")
all_weights: dict[str, mx.array] = {}
for shard in shard_files:
torch_weights = torch_load_file(str(shard))
for key, tensor in torch_weights.items():
if tensor.dtype == torch.bfloat16:
tensor = tensor.to(torch.float16)
all_weights[key] = mx.array(tensor.numpy())
return all_weights
@staticmethod
def _try_load_saved_component(
root_path: Path,
component_name: str,
) -> tuple[dict | None, int | None, str | None]:
component_path = root_path / component_name
if not component_path.exists():
return None, None, None
shard_files = sorted(f for f in component_path.glob("*.safetensors") if not f.name.startswith("._"))
if not shard_files:
return None, None, None
all_weights: dict[str, mx.array] = {}
quantization_level: int | None = None
mflux_version: str | None = None
for idx, shard in enumerate(shard_files):
data = mx.load(str(shard), return_metadata=True)
weights_dict = data[0]
all_weights.update(dict(weights_dict.items()))
if idx == 0 and len(data) > 1:
quantization_level = data[1].get("quantization_level")
mflux_version = data[1].get("mflux_version")
if quantization_level is None and mflux_version is None:
return None, None, None
unflattened = tree_unflatten(list(all_weights.items()))
return unflattened, quantization_level, mflux_version

635
src/mflux/models/fibo/weights/fibo_weight_mapping.py Normal file
View File

@ -0,0 +1,635 @@
from typing import List
from mflux.models.common.weights.mapping.weight_mapping import WeightMapping, WeightTarget
from mflux.models.qwen.weights.qwen_weight_mapping import (
reshape_gamma_to_1d,
transpose_conv2d_weight,
transpose_conv3d_weight,
)
class FIBOWeightMapping(WeightMapping):
@staticmethod
def get_transformer_mapping() -> List[WeightTarget]:
return [
# ========== Global projections ==========
WeightTarget(
mlx_path="time_embed.timestep_embedder.linear_1.weight",
hf_patterns=["time_embed.timestep_embedder.linear_1.weight"],
),
WeightTarget(
mlx_path="time_embed.timestep_embedder.linear_1.bias",
hf_patterns=["time_embed.timestep_embedder.linear_1.bias"],
),
WeightTarget(
mlx_path="time_embed.timestep_embedder.linear_2.weight",
hf_patterns=["time_embed.timestep_embedder.linear_2.weight"],
),
WeightTarget(
mlx_path="time_embed.timestep_embedder.linear_2.bias",
hf_patterns=["time_embed.timestep_embedder.linear_2.bias"],
),
WeightTarget(
mlx_path="context_embedder.weight",
hf_patterns=["context_embedder.weight"],
),
WeightTarget(
mlx_path="context_embedder.bias",
hf_patterns=["context_embedder.bias"],
),
WeightTarget(
mlx_path="x_embedder.weight",
hf_patterns=["x_embedder.weight"],
),
WeightTarget(
mlx_path="x_embedder.bias",
hf_patterns=["x_embedder.bias"],
),
# ========== Joint transformer blocks ==========
# AdaLayerNormZero (image + context streams)
WeightTarget(
mlx_path="transformer_blocks.{block}.norm1.linear.weight",
hf_patterns=["transformer_blocks.{block}.norm1.linear.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.norm1.linear.bias",
hf_patterns=["transformer_blocks.{block}.norm1.linear.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.norm1_context.linear.weight",
hf_patterns=["transformer_blocks.{block}.norm1_context.linear.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.norm1_context.linear.bias",
hf_patterns=["transformer_blocks.{block}.norm1_context.linear.bias"],
),
# Attention weights (BriaFiboAttention)
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.norm_q.weight",
hf_patterns=["transformer_blocks.{block}.attn.norm_q.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.norm_k.weight",
hf_patterns=["transformer_blocks.{block}.attn.norm_k.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_q.weight",
hf_patterns=["transformer_blocks.{block}.attn.to_q.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_q.bias",
hf_patterns=["transformer_blocks.{block}.attn.to_q.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_k.weight",
hf_patterns=["transformer_blocks.{block}.attn.to_k.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_k.bias",
hf_patterns=["transformer_blocks.{block}.attn.to_k.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_v.weight",
hf_patterns=["transformer_blocks.{block}.attn.to_v.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_v.bias",
hf_patterns=["transformer_blocks.{block}.attn.to_v.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_out.0.weight",
hf_patterns=["transformer_blocks.{block}.attn.to_out.0.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_out.0.bias",
hf_patterns=["transformer_blocks.{block}.attn.to_out.0.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.norm_added_q.weight",
hf_patterns=["transformer_blocks.{block}.attn.norm_added_q.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.norm_added_k.weight",
hf_patterns=["transformer_blocks.{block}.attn.norm_added_k.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.add_q_proj.weight",
hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.add_q_proj.bias",
hf_patterns=["transformer_blocks.{block}.attn.add_q_proj.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.add_k_proj.weight",
hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.add_k_proj.bias",
hf_patterns=["transformer_blocks.{block}.attn.add_k_proj.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.add_v_proj.weight",
hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.add_v_proj.bias",
hf_patterns=["transformer_blocks.{block}.attn.add_v_proj.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_add_out.weight",
hf_patterns=["transformer_blocks.{block}.attn.to_add_out.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.attn.to_add_out.bias",
hf_patterns=["transformer_blocks.{block}.attn.to_add_out.bias"],
),
# LayerNorm / FFN for image stream
WeightTarget(
mlx_path="transformer_blocks.{block}.norm2.weight",
hf_patterns=["transformer_blocks.{block}.norm2.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.norm2.bias",
hf_patterns=["transformer_blocks.{block}.norm2.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff.net.0.proj.weight",
hf_patterns=["transformer_blocks.{block}.ff.net.0.proj.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff.net.0.proj.bias",
hf_patterns=["transformer_blocks.{block}.ff.net.0.proj.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff.net.2.weight",
hf_patterns=["transformer_blocks.{block}.ff.net.2.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff.net.2.bias",
hf_patterns=["transformer_blocks.{block}.ff.net.2.bias"],
),
# LayerNorm / FFN for context stream
WeightTarget(
mlx_path="transformer_blocks.{block}.norm2_context.weight",
hf_patterns=["transformer_blocks.{block}.norm2_context.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.norm2_context.bias",
hf_patterns=["transformer_blocks.{block}.norm2_context.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff_context.net.0.proj.weight",
hf_patterns=["transformer_blocks.{block}.ff_context.net.0.proj.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff_context.net.0.proj.bias",
hf_patterns=["transformer_blocks.{block}.ff_context.net.0.proj.bias"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff_context.net.2.weight",
hf_patterns=["transformer_blocks.{block}.ff_context.net.2.weight"],
),
WeightTarget(
mlx_path="transformer_blocks.{block}.ff_context.net.2.bias",
hf_patterns=["transformer_blocks.{block}.ff_context.net.2.bias"],
),
# ========== Single transformer blocks ==========
WeightTarget(
mlx_path="single_transformer_blocks.{block}.norm.linear.weight",
hf_patterns=["single_transformer_blocks.{block}.norm.linear.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.norm.linear.bias",
hf_patterns=["single_transformer_blocks.{block}.norm.linear.bias"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.norm_q.weight",
hf_patterns=["single_transformer_blocks.{block}.attn.norm_q.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.norm_k.weight",
hf_patterns=["single_transformer_blocks.{block}.attn.norm_k.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.to_q.weight",
hf_patterns=["single_transformer_blocks.{block}.attn.to_q.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.to_q.bias",
hf_patterns=["single_transformer_blocks.{block}.attn.to_q.bias"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.to_k.weight",
hf_patterns=["single_transformer_blocks.{block}.attn.to_k.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.to_k.bias",
hf_patterns=["single_transformer_blocks.{block}.attn.to_k.bias"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.to_v.weight",
hf_patterns=["single_transformer_blocks.{block}.attn.to_v.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.attn.to_v.bias",
hf_patterns=["single_transformer_blocks.{block}.attn.to_v.bias"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.proj_mlp.weight",
hf_patterns=["single_transformer_blocks.{block}.proj_mlp.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.proj_mlp.bias",
hf_patterns=["single_transformer_blocks.{block}.proj_mlp.bias"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.proj_out.weight",
hf_patterns=["single_transformer_blocks.{block}.proj_out.weight"],
),
WeightTarget(
mlx_path="single_transformer_blocks.{block}.proj_out.bias",
hf_patterns=["single_transformer_blocks.{block}.proj_out.bias"],
),
# ========== Caption projection & output head ==========
WeightTarget(
mlx_path="norm_out.linear.weight",
hf_patterns=["norm_out.linear.weight"],
),
WeightTarget(
mlx_path="norm_out.linear.bias",
hf_patterns=["norm_out.linear.bias"],
),
WeightTarget(
mlx_path="proj_out.weight",
hf_patterns=["proj_out.weight"],
),
WeightTarget(
mlx_path="proj_out.bias",
hf_patterns=["proj_out.bias"],
),
# Caption_projection layers: we rely on explicit layer index to be
# expanded via num_layers; we pass num_layers explicitly from the loader.
WeightTarget(
mlx_path="caption_projection.{layer}.linear.weight",
hf_patterns=["caption_projection.{layer}.linear.weight"],
),
]
@staticmethod
def get_text_encoder_mapping() -> List[WeightTarget]:
return [
WeightTarget(
mlx_path="embed_tokens.weight",
hf_patterns=["model.embed_tokens.weight"],
),
WeightTarget(
mlx_path="layers.{block}.self_attn.q_proj.weight",
hf_patterns=["model.layers.{block}.self_attn.q_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.self_attn.k_proj.weight",
hf_patterns=["model.layers.{block}.self_attn.k_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.self_attn.v_proj.weight",
hf_patterns=["model.layers.{block}.self_attn.v_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.self_attn.o_proj.weight",
hf_patterns=["model.layers.{block}.self_attn.o_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.mlp.gate_proj.weight",
hf_patterns=["model.layers.{block}.mlp.gate_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.mlp.up_proj.weight",
hf_patterns=["model.layers.{block}.mlp.up_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.mlp.down_proj.weight",
hf_patterns=["model.layers.{block}.mlp.down_proj.weight"],
),
WeightTarget(
mlx_path="layers.{block}.input_layernorm.weight",
hf_patterns=["model.layers.{block}.input_layernorm.weight"],
),
WeightTarget(
mlx_path="layers.{block}.post_attention_layernorm.weight",
hf_patterns=["model.layers.{block}.post_attention_layernorm.weight"],
),
WeightTarget(
mlx_path="norm.weight",
hf_patterns=["model.norm.weight"],
),
]
@staticmethod
def get_vae_mapping() -> List[WeightTarget]:
return [
# ========== Encoder conv_in ==========
WeightTarget(
mlx_path="encoder.conv_in.conv3d.weight",
hf_patterns=["encoder.conv_in.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="encoder.conv_in.conv3d.bias",
hf_patterns=["encoder.conv_in.bias"],
),
# ========== Encoder down_blocks ==========
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.norm1.weight",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.norm1.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.weight",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv1.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv1.conv3d.bias",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv1.bias"],
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.norm2.weight",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.norm2.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.weight",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv2.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv2.conv3d.bias",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv2.bias"],
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
transform=transpose_conv3d_weight,
required=False,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias",
hf_patterns=["encoder.down_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
required=False,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.downsampler.resample_conv.weight",
hf_patterns=["encoder.down_blocks.{block}.downsampler.resample.1.weight"],
transform=transpose_conv2d_weight,
required=False,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.downsampler.resample_conv.bias",
hf_patterns=["encoder.down_blocks.{block}.downsampler.resample.1.bias"],
required=False,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.weight",
hf_patterns=["encoder.down_blocks.{block}.downsampler.time_conv.weight"],
transform=transpose_conv3d_weight,
required=False,
),
WeightTarget(
mlx_path="encoder.down_blocks.{block}.downsampler.time_conv.conv3d.bias",
hf_patterns=["encoder.down_blocks.{block}.downsampler.time_conv.bias"],
required=False,
),
# ========== Encoder mid_block ==========
WeightTarget(
mlx_path="encoder.mid_block.resnets.{i}.norm1.weight",
hf_patterns=["encoder.mid_block.resnets.{i}.norm1.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.weight",
hf_patterns=["encoder.mid_block.resnets.{i}.conv1.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="encoder.mid_block.resnets.{i}.conv1.conv3d.bias",
hf_patterns=["encoder.mid_block.resnets.{i}.conv1.bias"],
),
WeightTarget(
mlx_path="encoder.mid_block.resnets.{i}.norm2.weight",
hf_patterns=["encoder.mid_block.resnets.{i}.norm2.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.weight",
hf_patterns=["encoder.mid_block.resnets.{i}.conv2.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="encoder.mid_block.resnets.{i}.conv2.conv3d.bias",
hf_patterns=["encoder.mid_block.resnets.{i}.conv2.bias"],
),
WeightTarget(
mlx_path="encoder.mid_block.attentions.{i}.norm.weight",
hf_patterns=["encoder.mid_block.attentions.{i}.norm.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="encoder.mid_block.attentions.{i}.to_qkv.weight",
hf_patterns=["encoder.mid_block.attentions.{i}.to_qkv.weight"],
transform=transpose_conv2d_weight,
),
WeightTarget(
mlx_path="encoder.mid_block.attentions.{i}.to_qkv.bias",
hf_patterns=["encoder.mid_block.attentions.{i}.to_qkv.bias"],
),
WeightTarget(
mlx_path="encoder.mid_block.attentions.{i}.proj.weight",
hf_patterns=["encoder.mid_block.attentions.{i}.proj.weight"],
transform=transpose_conv2d_weight,
),
WeightTarget(
mlx_path="encoder.mid_block.attentions.{i}.proj.bias",
hf_patterns=["encoder.mid_block.attentions.{i}.proj.bias"],
),
# ========== Encoder output ==========
WeightTarget(
mlx_path="encoder.norm_out.weight",
hf_patterns=["encoder.norm_out.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="encoder.conv_out.conv3d.weight",
hf_patterns=["encoder.conv_out.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="encoder.conv_out.conv3d.bias",
hf_patterns=["encoder.conv_out.bias"],
),
# ========== Decoder conv_in ==========
WeightTarget(
mlx_path="decoder.conv_in.conv3d.weight",
hf_patterns=["decoder.conv_in.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="decoder.conv_in.conv3d.bias",
hf_patterns=["decoder.conv_in.bias"],
),
# ========== Decoder mid_block ==========
# Mid block resnets
WeightTarget(
mlx_path="decoder.mid_block.resnets.{i}.norm1.weight",
hf_patterns=["decoder.mid_block.resnets.{i}.norm1.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.weight",
hf_patterns=["decoder.mid_block.resnets.{i}.conv1.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="decoder.mid_block.resnets.{i}.conv1.conv3d.bias",
hf_patterns=["decoder.mid_block.resnets.{i}.conv1.bias"],
),
WeightTarget(
mlx_path="decoder.mid_block.resnets.{i}.norm2.weight",
hf_patterns=["decoder.mid_block.resnets.{i}.norm2.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.weight",
hf_patterns=["decoder.mid_block.resnets.{i}.conv2.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="decoder.mid_block.resnets.{i}.conv2.conv3d.bias",
hf_patterns=["decoder.mid_block.resnets.{i}.conv2.bias"],
),
# Mid block attention
WeightTarget(
mlx_path="decoder.mid_block.attentions.{i}.norm.weight",
hf_patterns=["decoder.mid_block.attentions.{i}.norm.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="decoder.mid_block.attentions.{i}.to_qkv.weight",
hf_patterns=["decoder.mid_block.attentions.{i}.to_qkv.weight"],
transform=transpose_conv2d_weight,
),
WeightTarget(
mlx_path="decoder.mid_block.attentions.{i}.to_qkv.bias",
hf_patterns=["decoder.mid_block.attentions.{i}.to_qkv.bias"],
),
WeightTarget(
mlx_path="decoder.mid_block.attentions.{i}.proj.weight",
hf_patterns=["decoder.mid_block.attentions.{i}.proj.weight"],
transform=transpose_conv2d_weight,
),
WeightTarget(
mlx_path="decoder.mid_block.attentions.{i}.proj.bias",
hf_patterns=["decoder.mid_block.attentions.{i}.proj.bias"],
),
# ========== Decoder up_blocks ==========
# Up blocks resnets
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.norm1.weight",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm1.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv1.conv3d.weight",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv1.conv3d.bias",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv1.bias"],
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.norm2.weight",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.norm2.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv2.conv3d.weight",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv2.conv3d.bias",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv2.bias"],
),
# Up blocks resnets - conv_shortcut (optional, when in_dim != out_dim)
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.weight",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.weight"],
transform=transpose_conv3d_weight,
required=False, # Optional - only exists when dimensions differ
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.conv3d.bias",
hf_patterns=["decoder.up_blocks.{block}.resnets.{res}.conv_shortcut.bias"],
required=False, # Optional - only exists when dimensions differ
),
# Up blocks upsamplers - time_conv (blocks 0, 1 only)
WeightTarget(
mlx_path="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.weight",
hf_patterns=["decoder.up_blocks.{block}.upsampler.time_conv.weight"],
transform=transpose_conv3d_weight,
required=False, # Only exists for blocks 0, 1
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.upsampler.time_conv.conv3d.bias",
hf_patterns=["decoder.up_blocks.{block}.upsampler.time_conv.bias"],
required=False, # Only exists for blocks 0, 1
),
# Up blocks upsamplers - resample_conv (blocks 0, 1, 2)
WeightTarget(
mlx_path="decoder.up_blocks.{block}.upsampler.resample_conv.weight",
hf_patterns=["decoder.up_blocks.{block}.upsampler.resample.1.weight"],
transform=transpose_conv2d_weight,
required=False, # Only exists for blocks 0, 1, 2
),
WeightTarget(
mlx_path="decoder.up_blocks.{block}.upsampler.resample_conv.bias",
hf_patterns=["decoder.up_blocks.{block}.upsampler.resample.1.bias"],
required=False, # Only exists for blocks 0, 1, 2
),
# ========== Decoder output ==========
WeightTarget(
mlx_path="decoder.norm_out.weight",
hf_patterns=["decoder.norm_out.gamma"],
transform=reshape_gamma_to_1d,
),
WeightTarget(
mlx_path="decoder.conv_out.conv3d.weight",
hf_patterns=["decoder.conv_out.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="decoder.conv_out.conv3d.bias",
hf_patterns=["decoder.conv_out.bias"],
),
# ========== Quant conv ==========
WeightTarget(
mlx_path="quant_conv.conv3d.weight",
hf_patterns=["quant_conv.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="quant_conv.conv3d.bias",
hf_patterns=["quant_conv.bias"],
),
# ========== Post quant conv ==========
WeightTarget(
mlx_path="post_quant_conv.conv3d.weight",
hf_patterns=["post_quant_conv.weight"],
transform=transpose_conv3d_weight,
),
WeightTarget(
mlx_path="post_quant_conv.conv3d.bias",
hf_patterns=["post_quant_conv.bias"],
),
]

Some files were not shown because too many files have changed in this diff Show More