Possibility to balance the Redux function (#190)

Co-authored-by: Alessandro Rizzo <alessandrorizzo@Alessandros-Mac-Studio.local>
Co-authored-by: filipstrand <strand.filip@gmail.com>

README.md

@@ -215,6 +215,46 @@ The `mflux-generate-in-context` command supports most of the same arguments as `
 
 See the [In-Context LoRA](#-in-context-lora) section for more details on how to use this feature effectively.
 
+#### 📜 Redux Tool Command-Line Arguments
+
+The `mflux-generate-redux` command uses most of the same arguments as `mflux-generate`, with these specific parameters:
+
+- **`--redux-image-paths`** (required, `[str]`): Paths to one or more reference images to influence the generation.
+
+- **`--redux-image-strengths`** (optional, `[float]`, default: `[1.0, 1.0, ...]`): Strength values (between 0.0 and 1.0) for each reference image. Higher values give more influence to that reference image. If not provided, all images use a strength of 1.0.
+
+See the [Redux](#-redux) section for more details on this feature.
+
+#### 📜 Fill Tool Command-Line Arguments
+
+The `mflux-generate-fill` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
+
+- **`--image-path`** (required, `str`): Path to the original image that will be modified.
+
+- **`--masked-image-path`** (required, `str`): Path to a binary mask image where white areas (255) will be regenerated and black areas (0) will be preserved.
+
+- **`--guidance`** (optional, `float`, default: `30.0`): The Fill tool works best with higher guidance values compared to regular image generation.
+
+See the [Fill](#-fill) section for more details on inpainting and outpainting.
+
+#### 📜 Depth Tool Command-Line Arguments
+
+The `mflux-generate-depth` command supports most of the same arguments as `mflux-generate`, with these specific parameters:
+
+- **`--image-path`** (optional, `str`, default: `None`): Path to the reference image from which to extract a depth map. Either this or `--depth-image-path` must be provided.
+
+- **`--depth-image-path`** (optional, `str`, default: `None`): Path to a pre-generated depth map image. Either this or `--image-path` must be provided.
+
+The `mflux-save-depth` command for extracting depth maps without generating images has these arguments:
+
+- **`--image-path`** (required, `str`): Path to the image from which to extract a depth map.
+
+- **`--output`** (optional, `str`, default: Uses the input filename with "_depth" suffix): Path where the generated depth map will be saved.
+
+- **`--quantize`** or **`-q`** (optional, `int`, default: `None`): Quantization for the Depth Pro model.
+
+See the [Depth](#-depth) section for more details on this feature.
+
 #### 📜 ControlNet Command-Line Arguments
 
 The `mflux-generate-controlnet` command supports most of the same arguments as `mflux-generate`, with these additional parameters:
@@ -941,6 +981,22 @@ mflux-generate-redux \
   --width 1154 \
   -q 8
 ```
+
+You can also control the influence of each reference image by specifying strength values (between 0.0 and 1.0) using the `--redux-image-strengths` parameter:
+
+```bash
+mflux-generate-redux \
+  --prompt "a grey statue of a cat on a white platform in front of a blue background" \
+  --redux-image-paths "image1.png" "image2.png" \
+  --redux-image-strengths 0.8 0.5 \
+  --steps 20 \
+  --height 1024 \
+  --width 768 \
+  -q 8
+```
+
+A higher strength value gives the reference image more influence over the final result. If you don't specify any strengths, a default value of 1.0 is used for all images.
+
 There is a tendency for the reference image to dominate over the input prompt.
 
 ⚠️ *Note: Using the Redux tool requires an additional 1.1GB download from [black-forest-labs/FLUX.1-Redux-dev](https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev). The download happens automatically on first use.*

src/mflux/config/config.py

@@ -19,6 +19,7 @@ class Config:
         image_strength: float | None = None,
         depth_image_path: Path | None = None,
         redux_image_paths: list[Path] | None = None,
+        redux_image_strengths: list[float] | None = None,
         masked_image_path: Path | None = None,
         controlnet_strength: float | None = None,
     ):
@@ -32,5 +33,6 @@ class Config:
         self.image_strength = image_strength
         self.depth_image_path = depth_image_path
         self.redux_image_paths = redux_image_paths
+        self.redux_image_strengths = redux_image_strengths
         self.masked_image_path = masked_image_path
         self.controlnet_strength = controlnet_strength

src/mflux/config/runtime_config.py

@@ -64,6 +64,10 @@ class RuntimeConfig:
     def redux_image_paths(self) -> list[Path] | None:
         return self.config.redux_image_paths
 
+    @property
+    def redux_image_strengths(self) -> list[float] | None:
+        return self.config.redux_image_strengths
+
     @property
     def masked_image_path(self) -> Path | None:
         return self.config.masked_image_path

src/mflux/flux_tools/redux/flux_redux.py

@@ -79,6 +79,7 @@ class Flux1Redux(nn.Module):
             image_paths=config.redux_image_paths,
             image_encoder=self.image_encoder,
             image_embedder=self.image_embedder,
+            image_strengths=config.redux_image_strengths,
         )  # fmt: off
 
         # (Optional) Call subscribers for beginning of loop
@@ -148,6 +149,7 @@ class Flux1Redux(nn.Module):
             lora_paths=self.lora_paths,
             lora_scales=self.lora_scales,
             redux_image_paths=config.redux_image_paths,
+            redux_image_strengths=config.redux_image_strengths,
             image_strength=config.image_strength,
             generation_time=time_steps.format_dict["elapsed"],
         )
@@ -163,6 +165,7 @@ class Flux1Redux(nn.Module):
         image_paths: list[str] | list[Path],
         image_encoder: SiglipVisionTransformer,
         image_embedder: ReduxEncoder,
+        image_strengths: list[float] | None = None,
     ) -> tuple[mx.array, mx.array]:
         # 1. Encode the prompt
         prompt_embeds_txt, pooled_prompt_embeds = PromptEncoder.encode_prompt(
@@ -179,6 +182,7 @@ class Flux1Redux(nn.Module):
             image_paths=image_paths,
             image_encoder=image_encoder,
             image_embedder=image_embedder,
+            image_strengths=image_strengths,
         )  # fmt:off
 
         # 3. Join text embeddings with all image embeddings

src/mflux/flux_tools/redux/redux_util.py

@@ -13,13 +13,20 @@ class ReduxUtil:
         image_paths: list[str] | list[Path],
         image_encoder: SiglipVisionTransformer,
         image_embedder: ReduxEncoder,
+        image_strengths: list[float] | None = None,
     ) -> list[mx.array]:  # fmt:off
         image_embeds_list = []
-        for image_path in image_paths:
+        for idx, image_path in enumerate(image_paths):
+            # Get the strength for this image (default to 1.0 if not specified)
+            strength = 1.0
+            if image_strengths is not None and idx < len(image_strengths):
+                strength = image_strengths[idx]
+
             image_embeds = ReduxUtil._embed_single_image(
                 image_path=image_path,
                 image_encoder=image_encoder,
                 image_embedder=image_embedder,
+                strength=strength,
             )
             image_embeds_list.append(image_embeds)
         return image_embeds_list
@@ -29,9 +36,15 @@ class ReduxUtil:
         image_path: str | Path,
         image_encoder: SiglipVisionTransformer,
         image_embedder: ReduxEncoder,
+        strength: float = 1.0,
     ) -> mx.array:  # fmt:off
         image = ImageUtil.load_image(image_path).convert("RGB")
         image = ImageUtil.preprocess_for_model(image=image)
         image_latents, pooler_output = image_encoder(image)
         image_embeds = image_embedder(image_latents)
+
+        # Apply strength factor to the image embeddings
+        if strength != 1.0:
+            image_embeds = image_embeds * strength
+
         return image_embeds

src/mflux/generate_redux.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from mflux import Config, ModelConfig, StopImageGenerationException
 from mflux.callbacks.callback_registry import CallbackRegistry
 from mflux.callbacks.instances.memory_saver import MemorySaver
@@ -17,6 +19,12 @@ def main():
     parser.add_output_arguments()
     args = parser.parse_args()
 
+    # Validate and normalize redux image strengths
+    redux_image_strengths = _validate_redux_image_strengths(
+        redux_image_paths=args.redux_image_paths,
+        redux_image_strengths=args.redux_image_strengths,
+    )
+
     # 1. Load the model
     flux = Flux1Redux(
         model_config=ModelConfig.dev_redux(),
@@ -52,6 +60,7 @@ def main():
                     width=args.width,
                     guidance=args.guidance,
                     redux_image_paths=args.redux_image_paths,
+                    redux_image_strengths=redux_image_strengths,
                 ),
             )
 
@@ -64,5 +73,25 @@ def main():
             print(memory_saver.memory_stats())
 
 
+def _validate_redux_image_strengths(
+    redux_image_paths: list[Path],
+    redux_image_strengths: list[float] | None,
+) -> list[float] | None:
+    if not redux_image_strengths or len(redux_image_strengths) == 0:
+        return redux_image_strengths
+
+    # If strengths provided but not enough for all images, fill with default (1.0)
+    if len(redux_image_strengths) < len(redux_image_paths):
+        redux_image_strengths.extend([1.0] * (len(redux_image_paths) - len(redux_image_strengths)))
+
+    # If too many strengths provided, raise error
+    elif len(redux_image_strengths) > len(redux_image_paths):
+        raise ValueError(
+            f"Too many strengths provided ({len(redux_image_strengths)}), expted at most {len(redux_image_paths)}."
+        )
+
+    return redux_image_strengths
+
+
 if __name__ == "__main__":
     main()

src/mflux/post_processing/generated_image.py

@@ -29,6 +29,7 @@ class GeneratedImage:
         masked_image_path: str | Path | None = None,
         depth_image_path: str | Path | None = None,
         redux_image_paths: list[str] | list[Path] | None = None,
+        redux_image_strengths: list[float] | None = None,
     ):
         self.image = image
         self.model_config = model_config
@@ -48,6 +49,7 @@ class GeneratedImage:
         self.masked_image_path = masked_image_path
         self.depth_image_path = depth_image_path
         self.redux_image_paths = redux_image_paths
+        self.redux_image_strengths = redux_image_strengths
 
     def get_right_half(self) -> "GeneratedImage":
         # Calculate the coordinates for the right half
@@ -85,13 +87,16 @@ class GeneratedImage:
 
         ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata, overwrite)
 
+    def _format_redux_strengths(self) -> list[float] | None:
+        if not self.redux_image_strengths:
+            return None
+        return [round(scale, 2) for scale in self.redux_image_strengths]
+
     def _get_metadata(self) -> dict:
         """Generate metadata for reference as well as input data for
         command line --config-from-metadata arg in future generations.
         """
         return {
-            # mflux_version is used by future metadata readers
-            # to determine supportability of metadata-derived workflows
             "mflux_version": GeneratedImage.get_version(),
             "model": self.model_config.model_name,
             "base_model": str(self.model_config.base_model),
@@ -110,6 +115,7 @@ class GeneratedImage:
             "masked_image_path": str(self.masked_image_path) if self.masked_image_path else None,
             "depth_image_path": str(self.depth_image_path) if self.depth_image_path else None,
             "redux_image_paths": str(self.redux_image_paths) if self.redux_image_paths else None,
+            "redux_image_strengths": self._format_redux_strengths(),
             "prompt": self.prompt,
         }
 

src/mflux/post_processing/image_util.py

@@ -29,6 +29,7 @@ class ImageUtil:
         controlnet_image_path: str | Path | None = None,
         image_path: str | Path | None = None,
         redux_image_paths: list[str] | list[Path] | None = None,
+        redux_image_strengths: list[float] | None = None,
         image_strength: float | None = None,
         masked_image_path: str | Path | None = None,
         depth_image_path: str | Path | None = None,
@@ -55,6 +56,7 @@ class ImageUtil:
             masked_image_path=masked_image_path,
             depth_image_path=depth_image_path,
             redux_image_paths=redux_image_paths,
+            redux_image_strengths=redux_image_strengths,
         )
 
     @staticmethod

src/mflux/ui/cli/parsers.py

@@ -103,6 +103,7 @@ class CommandLineParser(argparse.ArgumentParser):
 
     def add_redux_arguments(self) -> None:
         self.add_argument("--redux-image-paths", type=Path, nargs="*", required=True, help="Local path to the source image")
+        self.add_argument("--redux-image-strengths", type=float, nargs="*", default=None, help="Strength values (between 0.0 and 1.0) for each reference image. Default is 1.0 for all images.")
 
     def add_output_arguments(self) -> None:
         self.add_argument("--metadata", action="store_true", help="Export image metadata as a JSON file.")

tests/arg_parser/test_cli_argparser.py

@@ -595,6 +595,7 @@ def test_redux_args(mflux_redux_parser, mflux_redux_minimal_argv):
         assert len(args.redux_image_paths) == 2
         assert args.redux_image_paths[0] == Path("image1.png")
         assert args.redux_image_paths[1] == Path("image2.png")
+        assert args.redux_image_strengths is None  # Default should be None
 
     # Test with more image paths
     with patch("sys.argv", ["mflux-redux", "--redux-image-paths", "image1.png", "image2.png", "image3.png"]):
@@ -604,6 +605,21 @@ def test_redux_args(mflux_redux_parser, mflux_redux_minimal_argv):
         assert args.redux_image_paths[1] == Path("image2.png")
         assert args.redux_image_paths[2] == Path("image3.png")
 
+    # Test with redux_image_strengths parameter
+    with patch("sys.argv", mflux_redux_minimal_argv + ["--redux-image-strengths", "0.8", "0.5"]):
+        args = mflux_redux_parser.parse_args()
+        assert len(args.redux_image_paths) == 2
+        assert len(args.redux_image_strengths) == 2
+        assert args.redux_image_strengths[0] == pytest.approx(0.8)
+        assert args.redux_image_strengths[1] == pytest.approx(0.5)
+
+    # Test with single redux_image_strength
+    with patch("sys.argv", mflux_redux_minimal_argv + ["--redux-image-strengths", "0.3"]):
+        args = mflux_redux_parser.parse_args()
+        assert len(args.redux_image_paths) == 2
+        assert len(args.redux_image_strengths) == 1
+        assert args.redux_image_strengths[0] == pytest.approx(0.3)
+
     # Test with model argument
     with patch("sys.argv", mflux_redux_minimal_argv + ["--model", "dev"]):
         args = mflux_redux_parser.parse_args()