New: Add In-Context-LoRA feature

README.md

@@ -26,6 +26,10 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
 - [🔌 LoRA](#-lora)
   * [Multi-LoRA](#multi-lora)
   * [Supported LoRA formats (updated)](#supported-lora-formats-updated)
+- [🎭 In-Context LoRA](#-in-context-lora)
+  * [Available Styles](#available-styles)
+  * [How It Works](#how-it-works)
+  * [Tips for Best Results](#tips-for-best-results)
 - [🎛️ Dreambooth fine-tuning](#-dreambooth-fine-tuning)
   * [Training configuration](#training-configuration)
   * [Training example](#training-example)
@@ -391,7 +395,7 @@ mflux-generate \
     --quantize 8 \
     --height 1920 \
     --width 1024 \
-    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
+    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond's calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
 ```
 ![image](src/mflux/assets/comparison6.jpg)
 
@@ -450,7 +454,7 @@ mflux-generate \
     --seed 2 \
     --height 1920 \
     --width 1024 \
-    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond’s calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
+    --prompt "Tranquil pond in a bamboo forest at dawn, the sun is barely starting to peak over the horizon, panda practices Tai Chi near the edge of the pond, atmospheric perspective through the mist of morning dew, sunbeams, its movements are graceful and fluid — creating a sense of harmony and balance, the pond's calm waters reflecting the scene, inviting a sense of meditation and connection with nature, style of Howard Terpning and Jessica Rossier"
 ```
 
 *Note: When loading a quantized model from disk, there is no need to pass in `-q` flag, since we can infer this from the weight metadata.*
@@ -601,7 +605,85 @@ To report additional formats, examples or other any suggestions related to LoRA
 
 ---
 
-### 🎛 Dreambooth fine-tuning
+### 🎭 In-Context LoRA
+
+In-Context LoRA is a powerful technique that allows you to generate images in a specific style based on a reference image, without requiring model fine-tuning. This approach uses specialized LoRA weights that enable the model to understand and apply the visual context from your reference image to a new generation.
+
+This feature is based on the [In-Context LoRA for Diffusion Transformers](https://github.com/ali-vilab/In-Context-LoRA) project by Ali-ViLab.
+
+![image](src/mflux/assets/in_context_example.jpg)
+
+To use In-Context LoRA, you need:
+1. A reference image
+2. A style LoRA (optional - the in-context ability works without LoRAs, but they can significantly enhance the results)
+
+
+#### Available Styles
+
+MFLUX provides several pre-defined styles from the [Hugging Face ali-vilab/In-Context-LoRA repository](https://huggingface.co/ali-vilab/In-Context-LoRA) that you can use with the `--lora-style` argument:
+
+| Style Name     | Description                               |
+|----------------|-------------------------------------------|
+| `couple`       | Couple profile photography style          |
+| `storyboard`   | Film storyboard sketching style           |
+| `font`         | Font design and typography style          |
+| `home`         | Home decoration and interior design style |
+| `illustration` | Portrait illustration style               |
+| `portrait`     | Portrait photography style                |
+| `ppt`          | Presentation template style               |
+| `sandstorm`    | Sandstorm visual effect                   |
+| `sparklers`    | Sparklers visual effect                   |
+| `identity`     | Visual identity and branding design style |
+
+#### How It Works
+
+The In-Context LoRA generation creates a side-by-side image where:
+- The left side shows your reference image with noise applied
+- The right side shows the new generation that follows your prompt while maintaining the visual context
+
+The final output is automatically cropped to show only the right half (the generated image).
+
+![image](src/mflux/assets/in_context_how_it_works.jpg)
+
+
+#### Prompting for In-Context LoRA
+
+For best results with In-Context LoRA, your prompt should describe both the reference image and the target image you want to generate. Use markers like `[IMAGE1]`, `[LEFT]`, or `[RIGHT]` to distinguish between the two parts.
+
+Here's an example:
+
+```sh
+mflux-generate-in-context \
+  --model dev \
+  --steps 20 \
+  --quantize 8 \
+  --seed 42 \
+  --height 1024 \
+  --width 1024 \
+  --image-path "reference.png" \
+  --lora-style identity \
+  --prompt "In this set of two images, a bold modern typeface with the brand name 'DEMA' is introduced and is shown on a company merchandise product photo; [IMAGE1] a simplistic black logo featuring a modern typeface with the brand name 'DEMA' on a bright light green/yellowish background; [IMAGE2] the design is printed on a green/yellowish hoodie as a company merchandise product photo with a plain white background."
+```
+
+This prompt clearly describes both the reference image (after `[IMAGE1]`) and the desired output (after `[IMAGE2]`). Other marker pairs you can use include:
+- `[LEFT]` and `[RIGHT]`
+- `[TOP]` and `[BOTTOM]`
+- `[REFERENCE]` and `[OUTPUT]`
+
+**Important**: In the current implementation, the reference image is ALWAYS placed on the left side of the composition, and the generated image on the right side. When using marker pairs in your prompt, the first marker (e.g., `[IMAGE1]`, `[LEFT]`, `[REFERENCE]`) always refers to your reference image, while the second marker (e.g., `[IMAGE2]`, `[RIGHT]`, `[OUTPUT]`) refers to what you want to generate.
+
+#### Tips for Best Results
+
+1. **Choose the right reference image**: Select a reference image with a clear composition and structure that matches your intended output.
+2. **Adjust guidance**: Higher guidance values (7.0-9.0) tend to produce results that more closely follow your prompt.
+3. **Try different styles**: Each style produces distinctly different results - experiment to find the one that best matches your vision.
+4. **Increase steps**: For higher quality results, use 25-30 steps.
+5. **Detailed prompting**: Be specific about both the reference image and your desired output in your prompt.
+6. **Try without LoRA**: While LoRAs enhance the results, you can experiment without them to see the base in-context capabilities.
+
+---
+
+### 🎛️ Dreambooth fine-tuning
 
 As of release [v.0.5.0](https://github.com/filipstrand/mflux/releases/tag/v.0.5.0), MFLUX has support for fine-tuning your own LoRA adapters using the [Dreambooth](https://dreambooth.github.io) technique.
 
@@ -801,6 +883,7 @@ with different prompts and LoRA adapters active.
 - Some LoRA adapters does not work.
 - Currently, the supported controlnet is the [canny-only version](https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny).
 - Dreambooth training currently does not support sending in training parameters as flags.
+- In-Context LoRA currently only supports a left-right image setup (reference image on left, generated image on right).
 
 ### Optional Tool: Batch Image Renamer
 

pyproject.toml

@@ -56,6 +56,7 @@ homepage = "https://github.com/filipstrand/mflux"
 [project.scripts]
 mflux-generate = "mflux.generate:main"
 mflux-generate-controlnet = "mflux.generate_controlnet:main"
+mflux-generate-in-context = "mflux.generate_in_context:main"
 mflux-save = "mflux.save:main"
 mflux-train = "mflux.train:main"
 

src/mflux/community/__init__.py

@@ -0,0 +1,3 @@
+"""
+Third-party extensions and integrations for MFlux.
+"""

src/mflux/community/in_context_lora/__init__.py

@@ -0,0 +1,6 @@
+"""
+In-context LoRA implementation for MFlux.
+
+This module provides functionality for conditioning image generation using a reference image
+directly in the latent space, without requiring model fine-tuning.
+"""

src/mflux/community/in_context_lora/flux_in_context_lora.py

@@ -0,0 +1,193 @@
+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.flux.flux_initializer import FluxInitializer
+from mflux.latent_creator.latent_creator import Img2Img, LatentCreator
+from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
+from mflux.models.text_encoder.prompt_encoder import PromptEncoder
+from mflux.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
+from mflux.models.transformer.transformer import Transformer
+from mflux.models.vae.vae import VAE
+from mflux.post_processing.array_util import ArrayUtil
+from mflux.post_processing.generated_image import GeneratedImage
+from mflux.post_processing.image_util import ImageUtil
+
+
+class Flux1InContextLoRA(nn.Module):
+    vae: VAE
+    transformer: Transformer
+    t5_text_encoder: T5Encoder
+    clip_text_encoder: CLIPEncoder
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        quantize: int | None = None,
+        local_path: str | None = None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
+    ):
+        super().__init__()
+        FluxInitializer.init(
+            flux_model=self,
+            model_config=model_config,
+            quantize=quantize,
+            local_path=local_path,
+            lora_paths=lora_paths,
+            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
+        )
+
+    def generate_image(
+        self,
+        seed: int,
+        prompt: str,
+        config: Config,
+    ) -> GeneratedImage:
+        # 0. Create a new runtime config based on the model type and input parameters
+        config = RuntimeConfig(config, self.model_config)
+        time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
+
+        # 1. Encode the reference image
+        encoded_image = LatentCreator.encode_image(
+            height=config.height,
+            width=config.width,
+            img2img=Img2Img(
+                vae=self.vae,
+                sigmas=config.sigmas,
+                init_time_step=config.init_time_step,
+                image_path=config.image_path,
+            ),
+        )
+
+        # 2. Create the initial latents and keep the initial static noise for later blending
+        static_noise = Flux1InContextLoRA._create_in_context_latents(seed=seed, config=config)
+        latents = mx.array(static_noise)
+
+        # 3. Encode the prompt
+        prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt(
+            prompt=prompt,
+            prompt_cache=self.prompt_cache,
+            t5_tokenizer=self.t5_tokenizer,
+            clip_tokenizer=self.clip_tokenizer,
+            t5_text_encoder=self.t5_text_encoder,
+            clip_text_encoder=self.clip_text_encoder,
+        )
+
+        # (Optional) Call subscribers for beginning of loop
+        Callbacks.before_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )
+
+        for t in time_steps:
+            try:
+                # 4.t Predict the noise
+                noise = self.transformer(
+                    t=t,
+                    config=config,
+                    hidden_states=latents,
+                    prompt_embeds=prompt_embeds,
+                    pooled_prompt_embeds=pooled_prompt_embeds,
+                )
+
+                # 5.t Take one denoise step and update latents
+                dt = config.sigmas[t + 1] - config.sigmas[t]
+                latents += noise * dt
+
+                # 6.t Override left hand side of latents by linearly interpolating between latents and static noise
+                latents = Flux1InContextLoRA._update_latents(
+                    t=t,
+                    config=config,
+                    latents=latents,
+                    encoded_image=encoded_image,
+                    static_noise=static_noise,
+                )
+
+                # (Optional) Call subscribers at end of loop
+                Callbacks.in_loop(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
+
+                # (Optional) Evaluate to enable progress tracking
+                mx.eval(latents)
+
+            except KeyboardInterrupt:  # noqa: PERF203
+                Callbacks.interruption(
+                    t=t,
+                    seed=seed,
+                    prompt=prompt,
+                    latents=latents,
+                    config=config,
+                    time_steps=time_steps,
+                )
+
+        # 6. Decode the latent array and return the image
+        latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
+        decoded = self.vae.decode(latents)
+        return ImageUtil.to_image(
+            decoded_latents=decoded,
+            config=config,
+            seed=seed,
+            prompt=prompt,
+            quantization=self.bits,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            image_path=config.image_path,
+            image_strength=config.image_strength,
+            generation_time=time_steps.format_dict["elapsed"],
+        )
+
+    @staticmethod
+    def _create_in_context_latents(seed: int, config: RuntimeConfig):
+        # 1. Double the width for side-by-side generation
+        config.width = 2 * config.width
+
+        # 2. Create the initial latents with doubled width
+        latent_height = config.height // 8
+        latent_width = config.width // 8
+
+        # 3. Create noise with appropriate dimensions
+        static_noise = mx.random.normal(shape=[1, 16, latent_height, latent_width], key=mx.random.key(seed))
+        latents = ArrayUtil.pack_latents(latents=static_noise, height=config.height, width=config.width)
+        return latents
+
+    @staticmethod
+    def _update_latents(
+        t: int,
+        config: RuntimeConfig,
+        latents: mx.array,
+        encoded_image: mx.array,
+        static_noise: mx.array,
+    ) -> mx.array:  # fmt:off
+        # 1. Unpack the latents
+        unpacked = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
+        unpacked_static_noise = ArrayUtil.unpack_latents(latents=static_noise, height=config.height, width=config.width)
+
+        # 2. Calculate latent_width from the config (original width is half of current width)
+        latent_width = (config.width // 2) // 8
+
+        # 3. Override the left side with the reference image blended with appropriate noise for current timestep
+        unpacked[:, :, :, 0:latent_width] = LatentCreator.add_noise_by_interpolation(
+            clean=encoded_image[:, :, :, 0:latent_width],
+            noise=unpacked_static_noise[:, :, :, 0:latent_width],
+            sigma=config.sigmas[t+1]
+        )  # fmt:off
+
+        # 4. Repack the latents
+        return ArrayUtil.pack_latents(latents=unpacked, height=config.height, width=config.width)

src/mflux/community/in_context_lora/in_context_loras.py

@@ -0,0 +1,25 @@
+"""Configuration for In-Context LoRAs from Hugging Face."""
+
+# Default Hugging Face repository for In-Context LoRAs
+LORA_REPO_ID = "ali-vilab/In-Context-LoRA"
+
+# Mapping from simple names to actual LoRA filenames
+LORA_NAME_MAP = {
+    "couple": "couple-profile.safetensors",
+    "storyboard": "film-storyboard.safetensors",
+    "font": "font-design.safetensors",
+    "home": "home-decoration.safetensors",
+    "illustration": "portrait-illustration.safetensors",
+    "portrait": "portrait-photography.safetensors",
+    "ppt": "ppt-templates.safetensors",
+    "sandstorm": "sandstorm-visual-effect.safetensors",
+    "sparklers": "sparklers-visual-effect.safetensors",
+    "identity": "visual-identity-design.safetensors",
+}
+
+
+def get_lora_filename(simple_name: str) -> str:
+    if simple_name not in LORA_NAME_MAP:
+        valid_names = ", ".join(sorted(LORA_NAME_MAP.keys()))
+        raise ValueError(f"Unknown LoRA name: {simple_name}. Valid names are: {valid_names}")
+    return LORA_NAME_MAP[simple_name]

src/mflux/config/runtime_config.py

@@ -27,6 +27,10 @@ class RuntimeConfig:
     def width(self) -> int:
         return self.config.width
 
+    @width.setter
+    def width(self, value):
+        self.config.width = value
+
     @property
     def guidance(self) -> float:
         return self.config.guidance
@@ -53,16 +57,20 @@ class RuntimeConfig:
 
     @property
     def init_time_step(self) -> int:
-        is_txt2img = self.config.init_image_path is None or self.config.init_image_strength == 0.0
+        is_img2img = (
+            self.config.image_path is not None and
+            self.image_strength is not None and
+            self.image_strength > 0.0
+        )  # fmt: off
 
-        if is_txt2img:
-            return 0
-        else:
+        if is_img2img:
             # 1. Clamp strength to [0, 1]
-            strength = max(0.0, min(1.0, self.config.init_image_strength))
+            strength = max(0.0, min(1.0, self.config.image_strength))
 
             # 2. Return start time in [1, floor(num_steps * strength)]
             return max(1, int(self.num_inference_steps * strength))
+        else:
+            return 0
 
     @property
     def controlnet_strength(self) -> float | None:

src/mflux/generate_in_context.py

@@ -0,0 +1,70 @@
+from mflux import Config, StopImageGenerationException
+from mflux.callbacks.callback_registry import CallbackRegistry
+from mflux.callbacks.instances.memory_saver import MemorySaver
+from mflux.callbacks.instances.stepwise_handler import StepwiseHandler
+from mflux.community.in_context_lora.flux_in_context_lora import Flux1InContextLoRA
+from mflux.community.in_context_lora.in_context_loras import LORA_REPO_ID, get_lora_filename
+from mflux.config.model_config import ModelConfig
+from mflux.ui.cli.parsers import CommandLineParser
+
+
+def main():
+    # 0. Parse command line arguments
+    parser = CommandLineParser(description="Generate an image using in-context LoRA with a reference image.")
+    parser.add_general_arguments()
+    parser.add_model_arguments(require_model_arg=False)
+    parser.add_lora_arguments()
+    parser.add_image_generator_arguments(supports_metadata_config=True)
+    parser.add_image_to_image_arguments(required=True)
+    parser.add_output_arguments()
+    args = parser.parse_args()
+
+    # 1. Load the model
+    flux = Flux1InContextLoRA(
+        model_config=ModelConfig.dev(),
+        quantize=args.quantize,
+        lora_names=[get_lora_filename(args.lora_style)] if args.lora_style else None,
+        lora_repo_id=LORA_REPO_ID if args.lora_style else None,
+        lora_paths=args.lora_paths if not args.lora_style else None,
+        lora_scales=args.lora_scales,
+    )
+
+    # 2. Register the optional callbacks
+    if args.stepwise_image_output_dir:
+        handler = StepwiseHandler(flux=flux, output_dir=args.stepwise_image_output_dir)
+        CallbackRegistry.register_before_loop(handler)
+        CallbackRegistry.register_in_loop(handler)
+        CallbackRegistry.register_interrupt(handler)
+
+    memory_saver = None
+    if args.low_ram:
+        memory_saver = MemorySaver(flux)
+        CallbackRegistry.register_before_loop(memory_saver)
+        CallbackRegistry.register_in_loop(memory_saver)
+        CallbackRegistry.register_after_loop(memory_saver)
+
+    try:
+        for seed in args.seed:
+            # 3. Generate an image for each seed value
+            image = flux.generate_image(
+                seed=seed,
+                prompt=args.prompt,
+                config=Config(
+                    num_inference_steps=args.steps,
+                    height=args.height,
+                    width=args.width,
+                    guidance=args.guidance,
+                    image_path=args.image_path,
+                ),
+            )
+            # 4. Save the image
+            image.get_right_half().save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
+    except StopImageGenerationException as stop_exc:
+        print(stop_exc)
+    finally:
+        if memory_saver:
+            print(memory_saver.memory_stats())
+
+
+if __name__ == "__main__":
+    main()

src/mflux/latent_creator/latent_creator.py

@@ -57,12 +57,7 @@ class LatentCreator:
             )
 
             # 2. Encode the image
-            scaled_user_image = ImageUtil.scale_to_dimensions(
-                image=ImageUtil.load_image(img2img.init_image_path).convert("RGB"),
-                target_width=width,
-                target_height=height,
-            )
-            encoded = img2img.vae.encode(ImageUtil.to_array(scaled_user_image))
+            encoded = LatentCreator.encode_image(height=height, width=width, img2img=img2img)
             latents = ArrayUtil.pack_latents(latents=encoded, height=height, width=width)
 
             # 3. Find the appropriate sigma value
@@ -75,6 +70,16 @@ class LatentCreator:
                 sigma=sigma
             )  # fmt: off
 
+    @staticmethod
+    def encode_image(height: int, width: int, img2img: Img2Img):
+        scaled_user_image = ImageUtil.scale_to_dimensions(
+            image=ImageUtil.load_image(img2img.image_path).convert("RGB"),
+            target_width=width,
+            target_height=height,
+        )
+        encoded = img2img.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

src/mflux/post_processing/generated_image.py

@@ -44,6 +44,30 @@ class GeneratedImage:
         self.image_path = image_path
         self.image_strength = image_strength
 
+    def get_right_half(self) -> "GeneratedImage":
+        # Calculate the coordinates for the right half
+        width, height = self.image.size
+        right_half = self.image.crop((width // 2, 0, width, height))
+
+        # Create a new GeneratedImage with the right half and the same metadata
+        return GeneratedImage(
+            image=right_half,
+            model_config=self.model_config,
+            seed=self.seed,
+            prompt=self.prompt,
+            steps=self.steps,
+            guidance=self.guidance,
+            precision=self.precision,
+            quantization=self.quantization,
+            generation_time=self.generation_time,
+            lora_paths=self.lora_paths,
+            lora_scales=self.lora_scales,
+            controlnet_image_path=self.controlnet_image_path,
+            controlnet_strength=self.controlnet_strength,
+            image_path=self.image_path,
+            image_strength=self.image_strength,
+        )
+
     def save(
         self,
         path: t.Union[str, pathlib.Path],

tests/image_generation/helpers/__init__.py

@@ -0,0 +1,9 @@
+from tests.image_generation.helpers.image_generation_controlnet_test_helper import ImageGeneratorControlnetTestHelper
+from tests.image_generation.helpers.image_generation_in_context_test_helper import ImageGeneratorInContextTestHelper
+from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
+
+__all__ = [
+    "ImageGeneratorTestHelper",
+    "ImageGeneratorControlnetTestHelper",
+    "ImageGeneratorInContextTestHelper",
+]

tests/image_generation/helpers/image_generation_in_context_test_helper.py

@@ -0,0 +1,76 @@
+import os
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+from mflux import Config, ModelConfig
+from mflux.community.in_context_lora.flux_in_context_lora import Flux1InContextLoRA
+from mflux.community.in_context_lora.in_context_loras import LORA_REPO_ID, get_lora_filename
+
+
+class ImageGeneratorInContextTestHelper:
+    @staticmethod
+    def assert_matches_reference_image(
+        reference_image_path: str,
+        output_image_path: str,
+        model_config: ModelConfig,
+        prompt: str,
+        steps: int,
+        seed: int,
+        height: int = None,
+        width: int = None,
+        image_path: str | None = None,
+        lora_style: str | None = None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+    ):
+        # resolve paths
+        reference_image_path = ImageGeneratorInContextTestHelper.resolve_path(reference_image_path)
+        output_image_path = ImageGeneratorInContextTestHelper.resolve_path(output_image_path)
+        image_path = ImageGeneratorInContextTestHelper.resolve_path(image_path)
+        lora_paths = (
+            [str(ImageGeneratorInContextTestHelper.resolve_path(p)) for p in lora_paths] if lora_paths else None
+        )
+
+        try:
+            # given
+            flux = Flux1InContextLoRA(
+                model_config=model_config,
+                quantize=8,
+                lora_names=[get_lora_filename(lora_style)] if lora_style else None,
+                lora_repo_id=LORA_REPO_ID if lora_style else None,
+                lora_paths=lora_paths,
+                lora_scales=lora_scales
+            )  # fmt: off
+            # when
+            image = flux.generate_image(
+                seed=seed,
+                prompt=prompt,
+                config=Config(
+                    num_inference_steps=steps,
+                    image_path=image_path,
+                    height=height,
+                    width=width,
+                ),
+            )
+            # Save only the right half of the image (the generated part)
+            image.get_right_half().save(path=output_image_path, overwrite=True)
+
+            # then
+            np.testing.assert_array_equal(
+                np.array(Image.open(output_image_path)),
+                np.array(Image.open(reference_image_path)),
+                err_msg=f"Generated image doesn't match reference image. Check {output_image_path} vs {reference_image_path}",
+            )
+
+        finally:
+            # cleanup
+            if os.path.exists(output_image_path):
+                os.remove(output_image_path)
+
+    @staticmethod
+    def resolve_path(path) -> Path | None:
+        if path is None:
+            return None
+        return Path(__file__).parent.parent.parent / "resources" / path

tests/image_generation/test_generate_in_context.py

@@ -0,0 +1,20 @@
+from mflux import ModelConfig
+from tests.image_generation.helpers.image_generation_in_context_test_helper import ImageGeneratorInContextTestHelper
+
+
+class TestImageGeneratorInContext:
+    OUTPUT_IMAGE_FILENAME = "output.png"
+
+    def test_in_context_lora_identity(self):
+        ImageGeneratorInContextTestHelper.assert_matches_reference_image(
+            reference_image_path="ic_lora_reference_in_context_identity.png",
+            output_image_path=TestImageGeneratorInContext.OUTPUT_IMAGE_FILENAME,
+            model_config=ModelConfig.dev(),
+            steps=25,
+            seed=42,
+            height=320,
+            width=320,
+            prompt="In this set of two images, a bold modern typeface with the brand name 'DEMA' is introduced and is shown on a company merchandise product photo; [IMAGE1] a simplistic black logo featuring a modern typeface with the brand name 'DEMA' on a bright light green/yellowish background; [IMAGE2] the design is printed on a green/yellowish hoodie as a company merchandise product photo with a plain white background.",
+            image_path="ic_lora_reference_logo.png",
+            lora_style="identity",
+        )