Merge pull request #127 from filipstrand/in-context-lora

In-Context LoRA
2025-03-05 08:02:33 +01:00 · 2025-03-05 08:02:33 +01:00 · fb15edc7fc
commit fb15edc7fc
parent 552eb6a3bb ab3d7abdd2
26 changed files with 777 additions and 91 deletions
--- a/README.md
+++ b/README.md
@ -22,10 +22,15 @@ Run the powerful [FLUX](https://blackforestlabs.ai/#get-flux) models from [Black
  * [💾 Saving a quantized version to disk](#-saving-a-quantized-version-to-disk)
  * [💽 Loading and running a quantized version from disk](#-loading-and-running-a-quantized-version-from-disk)
 - [💽 Running a non-quantized model directly from disk](#-running-a-non-quantized-model-directly-from-disk)
+- [🌐 Third-Party HuggingFace Model Support](#-third-party-huggingface-model-support)
 - [🎨 Image-to-Image](#-image-to-image)
 - [🔌 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)
@ -153,13 +158,15 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2

 - **`--prompt`** (required, `str`): Text description of the image to generate.

- **`--model`** or **`-m`** (required, `str`): Model to use for generation (`"schnell"` or `"dev"`).
+- **`--model`** or **`-m`** (required, `str`): Model to use for generation. Can be one of the official models (`"schnell"` or `"dev"`) or a HuggingFace repository ID for a compatible third-party model (e.g., `"Freepik/flux.1-lite-8B-alpha"`).

- **`--output`** (optional, `str`, default: `"image.png"`): Output image filename. If `--seed` `--auto-seeds` establishes N > 1 seed values, the "stem" of the output file name will automatically append `_seed_{value}`.
+- **`--base-model`** (optional, `str`, default: `None`): Specifies which base architecture a third-party model is derived from (`"schnell"` or `"dev"`). Required when using third-party models from HuggingFace.

- **`--seed`** (optional, repeatable `int` args, default: `None`): 1 or more seeds for random number generation. e.g. `--seed 42` or `--seed 123 456 789`. Default is a single time-based value.
+- **`--output`** (optional, `str`, default: `"image.png"`): Output image filename. If `--seed` or `--auto-seeds` establishes multiple seed values, the output filename will automatically be modified to include the seed value (e.g., `image_seed_42.png`).

- **`--auto-seeds`** (optional, `int`, default: `None`): Auto generate N random Seeds in a series of image generations. Superseded by `--seed` arg and `seed` values in `--config-from-metadata` files.
+- **`--seed`** (optional, repeatable `int` args, default: `None`): 1 or more seeds for random number generation. e.g. `--seed 42` or `--seed 123 456 789`. When multiple seeds are provided, MFLUX will generate one image per seed, using the same prompt and settings. Default is a single time-based value.
+
+- **`--auto-seeds`** (optional, `int`, default: `None`): Auto generate N random Seeds in a series of image generations. For example, `--auto-seeds 5` will generate 5 different images with 5 different random seeds. This is superseded by explicit `--seed` arguments and `seed` values in `--config-from-metadata` files.

 - **`--height`** (optional, `int`, default: `1024`): Height of the output image in pixels.

@ -171,7 +178,7 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2

 - **`--path`** (optional, `str`, default: `None`): Path to a local model on disk.

- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `4` or `8`).
+- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): [Quantization](#%EF%B8%8F-quantization) (choose between `3`, `4`, `6`, or `8` bits).

 - **`--lora-paths`** (optional, `[str]`, default: `None`): The paths to the [LoRA](#-LoRA) weights.

@ -179,20 +186,56 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2

 - **`--metadata`** (optional): Exports a `.json` file containing the metadata for the image with the same name. (Even without this flag, the image metadata is saved and can be viewed using `exiftool image.png`)

- **`--controlnet-image-path`** (required, `str`): Path to the local image used by ControlNet to guide output generation.
+- **`--image-path`** (optional, `str`, default: `None`): Local path to the initial image for image-to-image generation.

- **`--controlnet-strength`** (optional, `float`, default: `0.4`): Degree of influence the control image has on the output. Ranges from `0.0` (no influence) to `1.0` (full influence).
-
- **`--controlnet-save-canny`** (optional, bool, default: False): If set, saves the Canny edge detection reference image used by ControlNet.
-
- **`--init-image-path`** (optional, `str`, default: `None`): Local path to the initial image for image-to-image generation.
-
- **`--init-image-strength`** (optional, `float`, default: `0.4`): Controls how strongly the initial image influences the output image. A value of `0.0` means no influence. (Default is `0.4`)
+- **`--image-strength`** (optional, `float`, default: `0.4`): Controls how strongly the initial image influences the output image. A value of `0.0` means no influence. (Default is `0.4`)

 - **`--config-from-metadata`** or **`-C`** (optional, `str`): [EXPERIMENTAL] Path to a prior file saved via `--metadata`, or a compatible handcrafted config file adhering to the expected args schema.

 - **`--low-ram`** (optional): Reduces GPU memory usage by constraining the MLX cache size and releasing text encoders and transformer components after use. This option is only compatible with single image generation. While it may slightly decrease performance, it helps prevent system memory swapping to disk, allowing generation on systems with limited RAM.

+- **`--lora-name`** (optional, `str`, default: `None`): The name of the LoRA to download from Hugging Face.
+
+- **`--lora-repo-id`** (optional, `str`, default: `"ali-vilab/In-Context-LoRA"`): The Hugging Face repository ID for LoRAs.
+
+- **`--stepwise-image-output-dir`** (optional, `str`, default: `None`): [EXPERIMENTAL] Output directory to write step-wise images and their final composite image to. This feature may change in future versions. When specified, MFLUX will save an image for each denoising step, allowing you to visualize the generation process from noise to final image.
+
+#### 📜 In-Context LoRA Command-Line Arguments
+
+The `mflux-generate-in-context` command supports most of the same arguments as `mflux-generate`, with these additional parameters:
+
+- **`--image-path`** (required, `str`): Path to the reference image that will guide the style of the generated image.
+
+- **`--lora-style`** (optional, `str`, default: `None`): The style to use for In-Context LoRA generation. Choose from: `couple`, `storyboard`, `font`, `home`, `illustration`, `portrait`, `ppt`, `sandstorm`, `sparklers`, or `identity`.
+
+See the [In-Context LoRA](#-in-context-lora) section for more details on how to use this feature effectively.
+
+#### 📜 ControlNet Command-Line Arguments
+
+The `mflux-generate-controlnet` command supports most of the same arguments as `mflux-generate`, with these additional parameters:
+
+- **`--controlnet-image-path`** (required, `str`): Path to the local image used by ControlNet to guide output generation.
+
+- **`--controlnet-strength`** (optional, `float`, default: `0.4`): Degree of influence the control image has on the output. Ranges from `0.0` (no influence) to `1.0` (full influence).
+
+- **`--controlnet-save-canny`** (optional, bool, default: `False`): If set, saves the Canny edge detection reference image used by ControlNet.
+
+See the [ControlNet](#%EF%B8%8F-controlnet) section for more details on how to use this feature effectively.
+
+
+#### 📜 Batch Image Generation Arguments
+
+- **`--prompts-file`** (required, `str`): Local path for a file that holds a batch of prompts.
+
+- **`--global-seed`** (optional, `int`): Entropy Seed used for all prompts in the batch.
+
+#### 📜 Training Arguments
+
+- **`--train-config`** (optional, `str`): Local path of the training configuration file. This file defines all aspects of the training process including model parameters, optimizer settings, and training data. See the [Training configuration](#training-configuration) section for details on the structure of this file.
+
+- **`--train-checkpoint`** (optional, `str`): Local path of the checkpoint file which specifies how to continue the training process. Used when resuming an interrupted training run.
+
+
 <details>
 <summary>parameters supported by config files</summary>

@ -391,7 +434,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)

@ -402,6 +445,8 @@ By selecting the `--quantize` or `-q` flag to be `4`, `8`, or removing it entire
 Image generation times in this example are based on a 2021 M1 Pro (32GB) machine. Even though the images are almost identical, there is a ~2x speedup by
 running the 8-bit quantized version on this particular machine. Unlike the non-quantized version, for the 8-bit version the swap memory usage is drastically reduced and GPU utilization is close to 100% during the whole generation. Results here can vary across different machines.

+For systems with limited RAM, you can also use the `--low-ram` option which reduces GPU memory usage by constraining the MLX cache size and releasing text encoders and transformer components after use. This option is particularly helpful for preventing system memory swapping to disk on machines with less available RAM.
+
 #### 📊 Size comparisons for quantized models

 The model sizes for both `schnell` and `dev` at various quantization levels are as follows:
@ -450,7 +495,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.*
@ -511,18 +556,43 @@ processed a bit differently, which is why we require this structure above.*

 ---

+### 🌐 Third-Party HuggingFace Model Support
+
+MFLUX now supports compatible third-party models from HuggingFace that follow the FLUX architecture. This opens up the ecosystem to community-created models that may offer different capabilities, sizes, or specializations.
+
+To use a third-party model, specify the HuggingFace repository ID with the `--model` parameter and indicate which base architecture (dev or schnell) it's derived from using the `--base-model` parameter:
+
+```sh
+mflux-generate \
+    --model Freepik/flux.1-lite-8B \
+    --base-model schnell \
+    --steps 4 \
+    --seed 42 \
+    --prompt "A beautiful landscape with mountains and a lake"
+```
+
+Some examples of compatible third-party models include:
+- [Freepik/flux.1-lite-8B-alpha](https://huggingface.co/Freepik/flux.1-lite-8B-alpha) - A lighter version of FLUX
+- [shuttleai/shuttle-3-diffusion](https://huggingface.co/shuttleai/shuttle-3-diffusion) - Shuttle's implementation based on FLUX
+
+The model will be automatically downloaded from HuggingFace the first time you use it, similar to the official FLUX models.
+
+*Note: Third-party models may have different performance characteristics, capabilities, or limitations compared to the official FLUX models. Always refer to the model's documentation on HuggingFace for specific usage instructions.*
+
+---
+
 ### 🎨 Image-to-Image

 One way to condition the image generation is by starting from an existing image and let MFLUX produce new variations.
-Use the `--init-image-path` flag to specify the reference image, and the `--init-image-strength` to control how much the reference 
+Use the `--image-path` flag to specify the reference image, and the `--image-strength` to control how much the reference 
 image should guide the generation. For example, given the reference image below, the following command produced the first
 image using the  [Sketching](https://civitai.com/models/803456/sketching?modelVersionId=898364) LoRA: 

 ```sh
 mflux-generate \
 --prompt "sketching of an Eiffel architecture, masterpiece, best quality. The site is lit by lighting professionals, creating a subtle illumination effect. Ink on paper with very fine touches with colored markers, (shadings:1.1), loose lines, Schematic, Conceptual, Abstract, Gestural. Quick sketches to explore ideas and concepts." \
--init-image-path "reference.png" \
--init-image-strength 0.3 \
+--image-path "reference.png" \
+--image-strength 0.3 \
 --lora-paths Architectural_Sketching.safetensors \
 --lora-scales 1.0 \
 --model dev \
@ -601,7 +671,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.

@ -641,6 +789,8 @@ To resume training for a given checkpoint, say `0001000_checkpoint.zip`, simply
 mflux-train --train-checkpoint 0001000_checkpoint.zip
 ```

+This uses the `--train-checkpoint` command-line argument to specify the checkpoint file to resume from.
+
 There are two nice properties of the training procedure: 

 - Fully deterministic (given a specified `seed` in the training configuration)
@ -801,6 +951,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

@ -828,17 +979,20 @@ See `uv run tools/rename_images.py --help` for full CLI usage help.
 - Set up shell aliases for required args examples:
  - shortcut for dev model: `alias mflux-dev='mflux-generate --model dev'`
  - shortcut for schnell model *and* always save metadata: `alias mflux-schnell='mflux-generate --model schnell --metadata'`
+- For systems with limited memory, use the `--low-ram` flag to reduce memory usage by constraining the MLX cache size and releasing components after use
+- When generating multiple images with different seeds, use `--seed` with multiple values or `--auto-seeds` to automatically generate a series of random seeds
+- Use `--stepwise-image-output-dir` to save intermediate images at each denoising step, which can be useful for debugging or creating animations of the generation process

 ### ✅ TODO

 - [ ] [FLUX.1 Tools](https://blackforestlabs.ai/flux-1-tools/)

-### 🔬 Cool research / features to support 
+### 🔬 Cool research / features to support
+- [ ] [ConceptAttention](https://github.com/helblazer811/ConceptAttention)
 - [ ] [PuLID](https://github.com/ToTheBeginning/PuLID)
 - [ ] [depth based controlnet](https://huggingface.co/InstantX/SD3-Controlnet-Depth) via [ml-depth-pro](https://github.com/apple/ml-depth-pro) or similar?
 - [ ] [RF-Inversion](https://github.com/filipstrand/mflux/issues/91)
- [ ] [In-Context LoRA](https://github.com/ali-vilab/In-Context-LoRA)
-
+- [ ] [catvton-flux](https://github.com/nftblackmagic/catvton-flux)

 ### 🌱‍ Related projects

--- a/pyproject.toml
+++ b/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"

--- a/src/mflux/assets/in_context_example.jpg
+++ b/src/mflux/assets/in_context_example.jpg
--- a/src/mflux/assets/in_context_how_it_works.jpg
+++ b/src/mflux/assets/in_context_how_it_works.jpg
--- a/src/mflux/community/init.py
+++ b/src/mflux/community/init.py
@ -0,0 +1,3 @@
+"""
+Third-party extensions and integrations for MFlux.
+"""
--- a/src/mflux/community/in_context_lora/init.py
+++ b/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.
+"""
--- a/src/mflux/community/in_context_lora/flux_in_context_lora.py
+++ b/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)
--- a/src/mflux/community/in_context_lora/in_context_loras.py
+++ b/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]
--- a/src/mflux/config/config.py
+++ b/src/mflux/config/config.py
@ -15,8 +15,8 @@ class Config:
        width: int = 1024,
        height: int = 1024,
        guidance: float = 4.0,
-        init_image_path: Path | None = None,
-        init_image_strength: float | None = None,
+        image_path: Path | None = None,
+        image_strength: float | None = None,
        controlnet_strength: float | None = None,
    ):
        if width % 16 != 0 or height % 16 != 0:
@ -25,6 +25,6 @@ class Config:
        self.height = 16 * (height // 16)
        self.num_inference_steps = num_inference_steps
        self.guidance = guidance
-        self.init_image_path = init_image_path
-        self.init_image_strength = init_image_strength
+        self.image_path = image_path
+        self.image_strength = image_strength
        self.controlnet_strength = controlnet_strength
--- a/src/mflux/config/runtime_config.py
+++ b/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
@ -44,25 +48,29 @@ class RuntimeConfig:
        return self.model_config.num_train_steps

    @property
-    def init_image_path(self) -> str:
-        return self.config.init_image_path
+    def image_path(self) -> str:
+        return self.config.image_path

    @property
-    def init_image_strength(self) -> float:
-        return self.config.init_image_strength
+    def image_strength(self) -> float | None:
+        return self.config.image_strength

    @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:
--- a/src/mflux/flux/flux.py
+++ b/src/mflux/flux/flux.py
@ -62,7 +62,7 @@ class Flux1(nn.Module):
                vae=self.vae,
                sigmas=config.sigmas,
                init_time_step=config.init_time_step,
-                init_image_path=config.init_image_path,
+                image_path=config.image_path,
            ),
        )

@ -141,8 +141,8 @@ class Flux1(nn.Module):
            quantization=self.bits,
            lora_paths=self.lora_paths,
            lora_scales=self.lora_scales,
-            init_image_path=config.init_image_path,
-            init_image_strength=config.init_image_strength,
+            image_path=config.image_path,
+            image_strength=config.image_strength,
            generation_time=time_steps.format_dict["elapsed"],
        )

--- a/src/mflux/flux/flux_initializer.py
+++ b/src/mflux/flux/flux_initializer.py
@ -10,6 +10,7 @@ from mflux.tokenizer.t5_tokenizer import TokenizerT5
 from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.weight_handler import WeightHandler
 from mflux.weights.weight_handler_lora import WeightHandlerLoRA
+from mflux.weights.weight_handler_lora_huggingface import WeightHandlerLoRAHuggingFace
 from mflux.weights.weight_util import WeightUtil


@ -20,10 +21,13 @@ class FluxInitializer:
        model_config: ModelConfig,
        quantize: int | None,
        local_path: str | None,
-        lora_paths: list[str] | None,
-        lora_scales: list[float] | 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,
    ) -> None:
        # 0. Set paths, configs and prompt_cache for later
+        lora_paths = lora_paths or []
        flux_model.prompt_cache = {}
        flux_model.lora_paths = lora_paths
        flux_model.lora_scales = lora_scales
@ -70,9 +74,13 @@ class FluxInitializer:
        )

        # 5. Set LoRA weights
+        hf_lora_paths = WeightHandlerLoRAHuggingFace.download_loras(
+            lora_names=lora_names,
+            repo_id=lora_repo_id,
+        )
        lora_weights = WeightHandlerLoRA.load_lora_weights(
            transformer=flux_model.transformer,
-            lora_files=lora_paths,
+            lora_files=lora_paths + hf_lora_paths,
            lora_scales=lora_scales,
        )
        WeightHandlerLoRA.set_lora_weights(
@ -86,8 +94,10 @@ class FluxInitializer:
        model_config: ModelConfig,
        quantize: int | None,
        local_path: str | None,
-        lora_paths: list[str] | None,
-        lora_scales: list[float] | 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,
    ) -> None:
        # 1. Start with same init as regular Flux
        FluxInitializer.init(
@ -97,6 +107,8 @@ class FluxInitializer:
            local_path=local_path,
            lora_paths=lora_paths,
            lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
        )

        # 2. Apply ControlNet-specific initialization
--- a/src/mflux/generate.py
+++ b/src/mflux/generate.py
@ -50,8 +50,8 @@ def main():
                    height=args.height,
                    width=args.width,
                    guidance=args.guidance,
-                    init_image_path=args.init_image_path,
-                    init_image_strength=args.init_image_strength,
+                    image_path=args.image_path,
+                    image_strength=args.image_strength,
                ),
            )
            # 4. Save the image
--- a/src/mflux/generate_in_context.py
+++ b/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()
--- a/src/mflux/latent_creator/latent_creator.py
+++ b/src/mflux/latent_creator/latent_creator.py
@ -11,12 +11,12 @@ class Img2Img:
        vae: VAE,
        sigmas: mx.array,
        init_time_step: int,
-        init_image_path: int,
+        image_path: int,
    ):
        self.vae = vae
        self.sigmas = sigmas
        self.init_time_step = init_time_step
-        self.init_image_path = init_image_path
+        self.image_path = image_path


 class LatentCreator:
@ -39,7 +39,7 @@ class LatentCreator:
        img2img: Img2Img,
    ) -> mx.array:
        # 0. Determine type of image generation
-        is_text2img = img2img.init_image_path is None
+        is_text2img = img2img.image_path is None

        if is_text2img:
            # 1. Create the pure noise
@ -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
--- a/src/mflux/post_processing/generated_image.py
+++ b/src/mflux/post_processing/generated_image.py
@ -25,8 +25,8 @@ class GeneratedImage:
        lora_scales: list[float],
        controlnet_image_path: str | pathlib.Path | None = None,
        controlnet_strength: float | None = None,
-        init_image_path: str | pathlib.Path | None = None,
-        init_image_strength: float | None = None,
+        image_path: str | pathlib.Path | None = None,
+        image_strength: float | None = None,
    ):
        self.image = image
        self.model_config = model_config
@ -41,8 +41,32 @@ class GeneratedImage:
        self.lora_scales = lora_scales
        self.controlnet_image_path = controlnet_image_path
        self.controlnet_strength = controlnet_strength
-        self.init_image_path = init_image_path
-        self.init_image_strength = init_image_strength
+        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,
@ -72,8 +96,8 @@ class GeneratedImage:
            "generation_time_seconds": round(self.generation_time, 2),
            "lora_paths": [str(p) for p in self.lora_paths] if self.lora_paths else None,
            "lora_scales": [round(scale, 2) for scale in self.lora_scales] if self.lora_scales else None,
-            "init_image_path": str(self.init_image_path) if self.init_image_path else None,
-            "init_image_strength": self.init_image_strength if self.init_image_path else None,
+            "image_path": str(self.image_path) if self.image_path else None,
+            "image_strength": self.image_strength if self.image_path else None,
            "controlnet_image_path": str(self.controlnet_image_path) if self.controlnet_image_path else None,
            "controlnet_strength": round(self.controlnet_strength, 2) if self.controlnet_strength else None,
            "prompt": self.prompt,
--- a/src/mflux/post_processing/image_util.py
+++ b/src/mflux/post_processing/image_util.py
@ -26,8 +26,8 @@ class ImageUtil:
        lora_paths: list[str],
        lora_scales: list[float],
        controlnet_image_path: str | None = None,
-        init_image_path: str | None = None,
-        init_image_strength: float | None = None,
+        image_path: str | None = None,
+        image_strength: float | None = None,
    ) -> GeneratedImage:
        normalized = ImageUtil._denormalize(decoded_latents)
        normalized_numpy = ImageUtil._to_numpy(normalized)
@ -44,8 +44,8 @@ class ImageUtil:
            generation_time=generation_time,
            lora_paths=lora_paths,
            lora_scales=lora_scales,
-            init_image_path=init_image_path,
-            init_image_strength=init_image_strength,
+            image_path=image_path,
+            image_strength=image_strength,
            controlnet_image_path=controlnet_image_path,
            controlnet_strength=config.controlnet_strength,
        )
--- a/src/mflux/ui/cli/parsers.py
+++ b/src/mflux/ui/cli/parsers.py
@ -5,6 +5,7 @@ import time
 import typing as t
 from pathlib import Path

+from mflux.community.in_context_lora.in_context_loras import LORA_NAME_MAP, LORA_REPO_ID
 from mflux.ui import defaults as ui_defaults


@ -47,10 +48,15 @@ class CommandLineParser(argparse.ArgumentParser):
        self.add_argument("--base-model", type=str, required=False, choices=ui_defaults.MODEL_CHOICES, help="When using a third-party huggingface model, explicitly specify whether the base model is dev or schnell")
        self.add_argument("--quantize",  "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)")

-    def add_lora_arguments(self) -> None:
+    def add_lora_arguments(self) -> None:  # fmt: off
        self.supports_lora = True
+        lora_group = self.add_argument_group("LoRA configuration")
+        lora_group.add_argument("--lora-style", type=str, choices=sorted(LORA_NAME_MAP.keys()), help="Style of the LoRA to use (e.g., 'storyboard' for film storyboard style)")
        self.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk")
        self.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.")
+        lora_group.add_argument("--lora-name", type=str, help="Name of the LoRA to download from Hugging Face")
+        lora_group.add_argument("--lora-repo-id", type=str, default=LORA_REPO_ID, help=f"Hugging Face repository ID for LoRAs (default: {LORA_REPO_ID})")
+    # fmt: on

    def _add_image_generator_common_arguments(self) -> None:
        self.supports_image_generation = True
@ -69,8 +75,8 @@ class CommandLineParser(argparse.ArgumentParser):

    def add_image_to_image_arguments(self, required=False) -> None:
        self.supports_image_to_image = True
-        self.add_argument("--init-image-path", type=Path, required=required, default=None, help="Local path to init image")
-        self.add_argument("--init-image-strength", type=float, required=False, default=ui_defaults.INIT_IMAGE_STRENGTH, help=f"Controls how strongly the init image influences the output image. A value of 0.0 means no influence. (Default is {ui_defaults.INIT_IMAGE_STRENGTH})")
+        self.add_argument("--image-path", type=Path, required=required, default=None, help="Local path to init image")
+        self.add_argument("--image-strength", type=float, required=False, default=ui_defaults.IMAGE_STRENGTH, help=f"Controls how strongly the init image influences the output image. A value of 0.0 means no influence. (Default is {ui_defaults.IMAGE_STRENGTH})")

    def add_batch_image_generator_arguments(self) -> None:
        self.add_argument("--prompts-file", type=Path, required=True, default=argparse.SUPPRESS, help="Local path for a file that holds a batch of prompts.")
@ -152,10 +158,10 @@ class CommandLineParser(argparse.ArgumentParser):
                    namespace.lora_scales = prior_gen_metadata.get("lora_scales", []) + namespace.lora_scales

            if self.supports_image_to_image:
-                if namespace.init_image_path is None:
-                    namespace.init_image_path = prior_gen_metadata.get("init_image_path", None)
-                if namespace.init_image_strength == self.get_default("init_image_strength") and (init_img_strength_from_metadata := prior_gen_metadata.get("init_image_strength", None)):
-                    namespace.init_image_strength = init_img_strength_from_metadata
+                if namespace.image_path is None:
+                    namespace.image_path = prior_gen_metadata.get("image_path", None)
+                if namespace.image_strength == self.get_default("image_strength") and (img_strength_from_metadata := prior_gen_metadata.get("image_strength", None)):
+                    namespace.image_strength = img_strength_from_metadata

            if self.supports_controlnet:
                if namespace.controlnet_image_path is None:
--- a/src/mflux/ui/defaults.py
+++ b/src/mflux/ui/defaults.py
@ -1,7 +1,7 @@
 CONTROLNET_STRENGTH = 0.4
 GUIDANCE_SCALE = 3.5
 HEIGHT, WIDTH = 1024, 1024
-INIT_IMAGE_STRENGTH = 0.4  # for image-to-image init_image
+IMAGE_STRENGTH = 0.4
 MODEL_CHOICES = ["dev", "schnell"]
 MODEL_INFERENCE_STEPS = {
    "dev": 14,
--- a/src/mflux/weights/weight_handler_lora_huggingface.py
+++ b/src/mflux/weights/weight_handler_lora_huggingface.py
@ -0,0 +1,74 @@
+import os
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+
+class WeightHandlerLoRAHuggingFace:
+    @staticmethod
+    def download_loras(
+        lora_names: list[str] = None,
+        repo_id: str = None,
+        cache_dir: str = None,
+    ) -> list[str]:
+        if repo_id is None:
+            return []
+
+        lora_paths = []
+        if lora_names:
+            for lora_name in lora_names:
+                lora_path = WeightHandlerLoRAHuggingFace._download_lora(
+                    repo_id=repo_id,
+                    lora_name=lora_name,
+                    cache_dir=cache_dir,
+                )
+                lora_paths.append(lora_path)
+
+        return lora_paths
+
+    @staticmethod
+    def _download_lora(
+        repo_id: str,
+        lora_name: str,
+        cache_dir: str = None,
+    ) -> str:
+        # Create cache directory if it doesn't exist
+        if cache_dir is None:
+            cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "mflux", "loras")
+
+        os.makedirs(cache_dir, exist_ok=True)
+
+        # Check if the file already exists in the cache
+        cached_file_path = os.path.join(cache_dir, lora_name)
+        if os.path.exists(cached_file_path):
+            print(f"Using cached LoRA: {cached_file_path}")
+            return cached_file_path
+
+        # Download the LoRA from Hugging Face
+        print(f"Downloading LoRA '{lora_name}' from {repo_id}...")
+        download_path = Path(
+            snapshot_download(
+                repo_id=repo_id,
+                allow_patterns=[f"*{lora_name}*"],
+                cache_dir=cache_dir,
+            )
+        )
+
+        # Find the downloaded file
+        for file in download_path.glob(f"**/*{lora_name}*"):
+            if file.is_file() and file.suffix in [".safetensors", ".bin"]:
+                # Copy or link the file to the cache directory with the expected name
+                target_path = os.path.join(cache_dir, lora_name)
+                if not os.path.exists(target_path):
+                    # Create a symlink or copy the file
+                    try:
+                        os.symlink(file, target_path)
+                    except (OSError, AttributeError):
+                        import shutil
+
+                        shutil.copy2(file, target_path)
+
+                print(f"LoRA downloaded to: {target_path}")
+                return target_path
+
+        raise FileNotFoundError(f"Could not find LoRA file '{lora_name}' in the downloaded repository.")
--- a/tests/arg_parser/test_cli_argparser.py
+++ b/tests/arg_parser/test_cli_argparser.py
@ -75,8 +75,8 @@ def base_metadata_dict() -> dict:
        "generation_time_seconds": 42.0,
        "lora_paths": None,
        "lora_scales": None,
-        "init_image": None,
-        "init_image_strength": None,
+        "image": None,
+        "image_strength": None,
        "controlnet_image": None,
        "controlnet_strength": None,
        "controlnet_save_canny": False,
@ -300,32 +300,32 @@ def test_image_to_image_args(mflux_generate_parser, mflux_generate_minimal_argv,
    metadata_file = temp_dir / "image_to_image.json"
    test_path = "/some/awesome/image.png"
    with metadata_file.open("wt") as m:
-        base_metadata_dict["init_image_path"] = test_path
+        base_metadata_dict["image_path"] = test_path
        json.dump(base_metadata_dict, m, indent=4)

    # test user default value
    with patch("sys.argv", mflux_generate_minimal_argv + ["-m", "dev"]):
        args = mflux_generate_parser.parse_args()
-        assert args.init_image_path is None
-        assert args.init_image_strength == 0.4  # default
+        assert args.image_path is None
+        assert args.image_strength == 0.4  # default

    # test metadata config accepted
    with patch('sys.argv', mflux_generate_minimal_argv + ['--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
        args = mflux_generate_parser.parse_args()
-        assert args.init_image_path == test_path
-        assert args.init_image_strength == 0.4  # default
+        assert args.image_path == test_path
+        assert args.image_strength == 0.4  # default

    # test strength override
-    with patch('sys.argv', mflux_generate_minimal_argv + ['--init-image-strength', '0.7', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--image-strength', '0.7', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
        args = mflux_generate_parser.parse_args()
-        assert args.init_image_path == test_path
-        assert args.init_image_strength == 0.7
+        assert args.image_path == test_path
+        assert args.image_strength == 0.7

    # test image path override
-    with patch('sys.argv', mflux_generate_minimal_argv + ['--init-image-path', '/some/better/image.png', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
+    with patch('sys.argv', mflux_generate_minimal_argv + ['--image-path', '/some/better/image.png', '--config-from-metadata', metadata_file.as_posix()]):  # fmt: off
        args = mflux_generate_parser.parse_args()
-        assert args.init_image_path == Path("/some/better/image.png")
-        assert args.init_image_strength == 0.4  # default
+        assert args.image_path == Path("/some/better/image.png")
+        assert args.image_strength == 0.4  # default


 def test_controlnet_args(mflux_generate_controlnet_parser, mflux_generate_controlnet_minimal_argv, base_metadata_dict, temp_dir):  # fmt: off
--- a/tests/image_generation/helpers/init.py
+++ b/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",
+]
--- a/tests/image_generation/helpers/image_generation_in_context_test_helper.py
+++ b/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
--- a/tests/image_generation/helpers/image_generation_test_helper.py
+++ b/tests/image_generation/helpers/image_generation_test_helper.py
@ -18,8 +18,8 @@ class ImageGeneratorTestHelper:
        seed: int,
        height: int = None,
        width: int = None,
-        init_image_path: str | None = None,
-        init_image_strength: float | None = None,
+        image_path: str | None = None,
+        image_strength: float | None = None,
        lora_paths: list[str] | None = None,
        lora_scales: list[float] | None = None,
    ):
@ -42,8 +42,8 @@ class ImageGeneratorTestHelper:
                prompt=prompt,
                config=Config(
                    num_inference_steps=steps,
-                    init_image_path=ImageGeneratorTestHelper.resolve_path(init_image_path),
-                    init_image_strength=init_image_strength,
+                    image_path=ImageGeneratorTestHelper.resolve_path(image_path),
+                    image_strength=image_strength,
                    height=height,
                    width=width,
                ),
--- a/tests/image_generation/test_generate_image.py
+++ b/tests/image_generation/test_generate_image.py
@ -60,8 +60,8 @@ class TestImageGenerator:
    def test_image_generation_dev_image_to_image(self):
        ImageGeneratorTestHelper.assert_matches_reference_image(
            reference_image_path="reference_dev_image_to_image_result.png",
-            init_image_path="reference_dev_lora.png",
-            init_image_strength=0.4,
+            image_path="reference_dev_lora.png",
+            image_strength=0.4,
            output_image_path=TestImageGenerator.OUTPUT_IMAGE_FILENAME,
            model_config=ModelConfig.dev(),
            steps=8,
--- a/tests/image_generation/test_generate_in_context.py
+++ b/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",
+        )