diff --git a/README.md b/README.md index 024a04f..1b7b170 100644 --- a/README.md +++ b/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.* diff --git a/src/mflux/config/config.py b/src/mflux/config/config.py index 8fd4817..2f357a1 100644 --- a/src/mflux/config/config.py +++ b/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 diff --git a/src/mflux/config/runtime_config.py b/src/mflux/config/runtime_config.py index 6cba76c..6ee11bd 100644 --- a/src/mflux/config/runtime_config.py +++ b/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 diff --git a/src/mflux/flux_tools/redux/flux_redux.py b/src/mflux/flux_tools/redux/flux_redux.py index d507715..3be8623 100644 --- a/src/mflux/flux_tools/redux/flux_redux.py +++ b/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 diff --git a/src/mflux/flux_tools/redux/redux_util.py b/src/mflux/flux_tools/redux/redux_util.py index dfc5692..01cb57f 100644 --- a/src/mflux/flux_tools/redux/redux_util.py +++ b/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 diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py index db5bff3..b52ed3f 100644 --- a/src/mflux/generate_redux.py +++ b/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() diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index 8232d0f..87b8bc5 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/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, } diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index b7b013a..0501c7e 100644 --- a/src/mflux/post_processing/image_util.py +++ b/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 diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index f1ae894..343c00c 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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.") diff --git a/tests/arg_parser/test_cli_argparser.py b/tests/arg_parser/test_cli_argparser.py index f16f328..8ed12d9 100644 --- a/tests/arg_parser/test_cli_argparser.py +++ b/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()