FIBO-Lite support (#366)

Co-authored-by: Cursor Assistant <assistant@cursor.com>

Author: filipstrand

README.md

@@ -119,9 +119,9 @@ MFLUX supports the following model families. They have different strengths and w
 
 | Model | Release date | Size | Type | Training | Description |
 | --- | --- | --- | --- | --- | --- |
-|[Z-Image](src/mflux/models/z_image/README.md) | Nov 2025 | 6B | Distilled & Base | Yes | Best all-rounder: fast, small, very good quality and realism. |
+|[Z-Image](src/mflux/models/z_image/README.md) | Nov 2025 | 6B | Distilled & Base | Yes | Fast, small, very good quality and realism. |
 |[FLUX.2](src/mflux/models/flux2/README.md) | Jan 2026 | 4B & 9B | Distilled & Base | Yes | Fastest + smallest with very good qaility and edit capabilities. |
-|[FIBO](src/mflux/models/fibo/README.md) | Oct 2025 | 8B | Base | No | Very good JSON-based prompt understanding and editability, medium speed |
+|[FIBO](src/mflux/models/fibo/README.md) | Oct 2025+ | 8B | Distilled & Base | No | Very good JSON-based prompt understanding and editability. |
 |[SeedVR2](src/mflux/models/seedvr2/README.md) | Jun 2025 | 3B & 7B | — | No | Best upscaling model. |
 |[Qwen Image](src/mflux/models/qwen/README.md) | Aug 2025+ | 20B | Base | No | Large model (slower); strong prompt understanding and world knowledge. Has edit capabilities |
 |[Depth Pro](src/mflux/models/depth_pro/README.md) | Oct 2024 | — | — | No | Very fast and accurate depth estimation model from Apple. |

src/mflux/cli/defaults/defaults.py

@@ -19,6 +19,7 @@
     "dev-krea",
     "qwen",
     "fibo",
+    "fibo-lite",
     "z-image",
     "z-image-turbo",
     "flux2-klein-4b",
@@ -34,6 +35,7 @@
     "qwen-image": 20,
     "qwen-image-edit": 20,
     "fibo": 20,
+    "fibo-lite": 8,
     "z-image": 50,
     "z-image-turbo": 9,
     "flux2-klein-4b": 4,

src/mflux/models/common/config/model_config.py

@@ -137,6 +137,11 @@ def qwen_image_edit() -> "ModelConfig":
     def fibo() -> "ModelConfig":
         return AVAILABLE_MODELS["fibo"]
 
+    @staticmethod
+    @lru_cache
+    def fibo_lite() -> "ModelConfig":
+        return AVAILABLE_MODELS["fibo-lite"]
+
     @staticmethod
     @lru_cache
     def z_image_turbo() -> "ModelConfig":
@@ -453,8 +458,20 @@ def from_name(
         supports_guidance=True,
         requires_sigma_shift=False,
     ),
-    "z-image": ModelConfig(
+    "fibo-lite": ModelConfig(
         priority=18,
+        aliases=["fibo-lite", "fibo_lite"],
+        model_name="briaai/Fibo-lite",
+        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,
+    ),
+    "z-image": ModelConfig(
+        priority=19,
         aliases=["z-image", "zimage"],
         model_name="Tongyi-MAI/Z-Image",
         base_model=None,
@@ -466,7 +483,7 @@ def from_name(
         requires_sigma_shift=True,
     ),
     "z-image-turbo": ModelConfig(
-        priority=19,
+        priority=20,
         aliases=["z-image-turbo", "zimage-turbo"],
         model_name="Tongyi-MAI/Z-Image-Turbo",
         base_model=None,
@@ -478,7 +495,7 @@ def from_name(
         requires_sigma_shift=True,
     ),
     "seedvr2-3b": ModelConfig(
-        priority=20,
+        priority=21,
         aliases=["seedvr2-3b", "seedvr2"],
         model_name="numz/SeedVR2_comfyUI",
         base_model=None,
@@ -490,7 +507,7 @@ def from_name(
         requires_sigma_shift=None,
     ),
     "seedvr2-7b": ModelConfig(
-        priority=21,
+        priority=22,
         aliases=["seedvr2-7b", "seedvr2-7B"],
         model_name="numz/SeedVR2_comfyUI",
         base_model=None,

src/mflux/models/fibo/README.md

@@ -13,6 +13,38 @@ Most text-to-image models excel at imagination—but not control. FIBO is traine
 - **Strong prompt adherence**: High alignment on PRISM-style evaluations
 - **Enterprise-grade**: 100% licensed data with governance, repeatability, and legal clarity
 
+## FIBO Lite
+
+[FIBO Lite](https://huggingface.co/briaai/Fibo-lite) is a two-stage distilled variant combining CFG distillation and SCFM for fast few-step generation. Use `--model fibo-lite` for ~10x speed: 8 steps, `guidance=1.0`, no negative prompt needed. Slight quality tradeoff vs. base FIBO.
+
+```sh
+mflux-generate-fibo \
+  --model fibo-lite \
+  --prompt "A tiny watercolor robot in a garden" \
+  --steps 8 \
+  --seed 42
+```
+
+<details>
+<summary>Python API</summary>
+
+```python
+from mflux.models.common.config import ModelConfig
+from mflux.models.fibo.variants.txt2img.fibo import FIBO
+from mflux.models.fibo_vlm.model.fibo_vlm import FiboVLM
+
+vlm = FiboVLM()
+json_prompt = vlm.generate(prompt="A tiny watercolor robot in a garden", seed=42)
+model = FIBO(model_config=ModelConfig.fibo_lite())
+image = model.generate_image(
+    seed=42,
+    prompt=json_prompt,
+    num_inference_steps=8,
+)
+image.save("robot_lite.png")
+```
+</details>
+
 ## The three modes: Generate, Refine, and Inspire
 
 ### Generate
@@ -334,5 +366,5 @@ image.save("bird_inspired.png")
 
 ## Notes
 > [!WARNING]
-> 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.
+> FIBO requires downloading the `briaai/FIBO` or `briaai/FIBO-lite` model weights (~24GB) and the `briaai/FIBO-vlm` vision-language model (~8GB), totaling ~32GB for the full model, or use quantization for smaller sizes.
 

src/mflux/models/fibo/cli/fibo_generate.py

@@ -1,6 +1,7 @@
 from mflux.callbacks.callback_manager import CallbackManager
 from mflux.cli.defaults import defaults as ui_defaults
 from mflux.cli.parser.parsers import CommandLineParser
+from mflux.models.common.config.model_config import ModelConfig
 from mflux.models.fibo.latent_creator.fibo_latent_creator import FiboLatentCreator
 from mflux.models.fibo.variants.txt2img.fibo import FIBO
 from mflux.models.fibo.variants.txt2img.util import FiboUtil
@@ -22,14 +23,23 @@ def main():
 
     # 0. Set default guidance value if not provided by user
     if args.guidance is None:
-        args.guidance = ui_defaults.GUIDANCE_SCALE
+        if args.model == "fibo-lite":
+            args.guidance = 1.0  # distilled, no CFG
+        elif args.model == "fibo":
+            args.guidance = 5.0  # base FIBO typical
+        else:
+            args.guidance = ui_defaults.GUIDANCE_SCALE
+
+    resolved_model_name = args.model if args.model in ui_defaults.MODEL_CHOICES else "fibo"
+    model_config = ModelConfig.from_name(model_name=resolved_model_name, base_model=args.base_model)
 
     json_prompt = FiboUtil.get_json_prompt(args, quantize=args.quantize)
 
     # 1. Load the FIBO model
     fibo = FIBO(
         quantize=args.quantize,
         model_path=args.model_path,
+        model_config=model_config,
     )
 
     # 2. Register callbacks

src/mflux/models/fibo/model/fibo_text_encoder/prompt_encoder.py

@@ -14,16 +14,11 @@ def encode_prompt(
         negative_prompt: str | None,
         tokenizer: Tokenizer,
         text_encoder: SmolLM3_3B_TextEncoder,
+        guidance: float = 4.0,
     ) -> 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,
@@ -32,6 +27,18 @@ def encode_prompt(
             max_sequence_length=2048,
             tokenization_prefix="positive",
         )
+
+        if guidance == 1.0:
+            encoder_hidden_states, prompt_layers = PromptEncoder._prepare_positive_only_output(
+                prompt_embeds=prompt_embeds,
+                prompt_layers=prompt_layers,
+                prompt_attention_mask=prompt_attention_mask,
+            )
+            return json_prompt, encoder_hidden_states, prompt_layers
+
+        if negative_prompt is None or negative_prompt == "":
+            negative_prompt = "ugly, blurry, low quality"
+
         neg_prompt_embeds, neg_prompt_layers, neg_prompt_attention_mask = PromptEncoder._get_prompt_embeds(
             prompt=negative_prompt,
             tokenizer=tokenizer,
@@ -54,6 +61,27 @@ def encode_prompt(
         )
         return json_prompt, encoder_hidden_states, prompt_layers
 
+    @staticmethod
+    def _prepare_positive_only_output(
+        prompt_embeds: mx.array,
+        prompt_layers: List[mx.array],
+        prompt_attention_mask: mx.array,
+    ) -> tuple[mx.array, List[mx.array]]:
+        """Prepare encoder output when guidance=1.0 (no negative prompt, positive only)."""
+        max_tokens = prompt_embeds.shape[1]
+        encoder_hidden_states, _ = PromptEncoder._pad_embedding(
+            prompt_embeds=prompt_embeds,
+            max_tokens=max_tokens,
+            attention_mask=prompt_attention_mask,
+        )
+        prompt_layers = [PromptEncoder._pad_embedding(layer, max_tokens)[0] for layer in prompt_layers]
+        total_num_layers_transformer = 46
+        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 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])

src/mflux/models/fibo/variants/txt2img/fibo.py

@@ -57,10 +57,13 @@ def generate_image(
         negative_prompt: str | None = None,
     ) -> GeneratedImage:
         # 0. Create a new config based on the model type and input parameters
+        effective_guidance = guidance
+        if "fibo-lite" in self.model_config.aliases:
+            effective_guidance = 1.0  # distilled model, cond-only
         config = Config(
             width=width,
             height=height,
-            guidance=guidance,
+            guidance=effective_guidance,
             scheduler=scheduler,
             image_path=image_path,
             image_strength=image_strength,
@@ -89,6 +92,7 @@ def generate_image(
             negative_prompt=negative_prompt,
             tokenizer=self.tokenizers["fibo"],
             text_encoder=self.text_encoder,
+            guidance=config.guidance,
         )
 
         # 3. Create callback context and call before_loop
@@ -105,7 +109,8 @@ def generate_image(
                     text_encoder_layers=text_encoder_layers,
                     encoder_hidden_states=encoder_hidden_states,
                 )
-                noise = FIBO._apply_classifier_free_guidance(noise, config.guidance)
+                if config.guidance != 1.0:
+                    noise = FIBO._apply_classifier_free_guidance(noise, config.guidance)
 
                 # 5.t Take one denoise step
                 latents = config.scheduler.step(noise=noise, timestep=t, latents=latents)

src/mflux/utils/generated_image.py

@@ -152,7 +152,8 @@ def _format_redux_strengths(self) -> list[float] | None:
         return [round(scale, 2) for scale in self.redux_image_strengths]
 
     def _is_fibo_model(self) -> bool:
-        return self.model_config.model_name == "briaai/FIBO" or str(self.model_config.base_model) == "fibo"
+        name = self.model_config.model_name
+        return name == "briaai/FIBO" or name == "briaai/Fibo-lite" or str(self.model_config.base_model) == "fibo"
 
     def _save_prompt_file(self, image_path: str | Path, overwrite: bool) -> None:
         file_path = Path(image_path)

tests/image_generation/helpers/image_generation_fibo_test_helper.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Optional
 
+from mflux.models.common.config import ModelConfig
 from mflux.models.fibo.variants.txt2img.fibo import FIBO
 from mflux.utils.image_compare import ImageCompare
 
@@ -20,8 +21,9 @@ def assert_matches_reference_image(
         negative_prompt: Optional[str] = None,
         mismatch_threshold: Optional[float] = None,
         quantize: Optional[int] = None,
+        model_config: Optional[ModelConfig] = None,
     ):
-        # resolve paths
+        model_config = model_config or ModelConfig.fibo()
         reference_image_path = ImageGeneratorFiboTestHelper.resolve_path(reference_image_path)
         output_image_path = ImageGeneratorFiboTestHelper.resolve_path(output_image_path)
 
@@ -30,6 +32,7 @@ def assert_matches_reference_image(
             model = FIBO(
                 quantize=quantize,
                 model_path=None,
+                model_config=model_config,
             )
 
             # Step 2: Generate image from prompt

tests/image_generation/test_generate_image_fibo.py

@@ -1,5 +1,6 @@
 import pytest
 
+from mflux.models.common.config import ModelConfig
 from tests.image_generation.helpers.image_generation_fibo_test_helper import ImageGeneratorFiboTestHelper
 
 OWL_PROMPT = """
@@ -123,3 +124,18 @@ def test_image_generation_fibo_refined_white_owl(self):
             guidance=4.0,
             quantize=8,
         )
+
+    @pytest.mark.slow
+    def test_image_generation_fibo_lite(self):
+        ImageGeneratorFiboTestHelper.assert_matches_reference_image(
+            reference_image_path="reference_fibo_lite.png",
+            output_image_path="output_fibo_lite.png",
+            prompt=OWL_PROMPT,
+            steps=8,
+            seed=42,
+            height=352,
+            width=640,
+            guidance=1.0,
+            quantize=8,
+            model_config=ModelConfig.fibo_lite(),
+        )