diff --git a/README.md b/README.md index fdb470d..9873d9f 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,8 @@ mflux-generate --model dev --prompt "Luxury food photograph" --steps 25 --seed 2 - **`--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. + Or, with the correct python environment active, create and run a separate script like the following: ```python diff --git a/src/mflux/controlnet/controlnet_util.py b/src/mflux/controlnet/controlnet_util.py index e4eadf8..daacf33 100644 --- a/src/mflux/controlnet/controlnet_util.py +++ b/src/mflux/controlnet/controlnet_util.py @@ -1,9 +1,12 @@ import logging +import os import cv2 import numpy as np import PIL +from mflux import ImageUtil + log = logging.getLogger(__name__) @@ -22,3 +25,9 @@ class ControlnetUtil: log.warning(f"Control image has different dimensions than the model. Resizing to {width}x{height}") img = img.resize((width, height), PIL.Image.LANCZOS) return img + + @staticmethod + def save_canny_image(control_image, path: str): + base, ext = os.path.splitext(path) + new_filename = f"{base}_controlnet_canny{ext}" + ImageUtil.save(control_image, new_filename) diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py index a03e883..5b0d3e8 100644 --- a/src/mflux/controlnet/flux_controlnet.py +++ b/src/mflux/controlnet/flux_controlnet.py @@ -103,7 +103,9 @@ class Flux1Controlnet: self, seed: int, prompt: str, + output: str, control_image_path: str, + save_control_image_canny: bool = False, config: ConfigControlnet = ConfigControlnet() ) -> GeneratedImage: # fmt: off # Create a new runtime config based on the model type and input parameters @@ -114,6 +116,8 @@ class Flux1Controlnet: control_image = ImageUtil.load_image(control_image_path) control_image = ControlnetUtil.scale_image(config.height, config.width, control_image) control_image = ControlnetUtil.preprocess_canny(control_image) + if save_control_image_canny: + ControlnetUtil.save_canny_image(control_image, output) controlnet_cond = ImageUtil.to_array(control_image) controlnet_cond = self.vae.encode(controlnet_cond) controlnet_cond = (controlnet_cond / self.vae.scaling_factor) + self.vae.shift_factor diff --git a/src/mflux/generate_controlnet.py b/src/mflux/generate_controlnet.py index a8ff54e..b1ecfaf 100644 --- a/src/mflux/generate_controlnet.py +++ b/src/mflux/generate_controlnet.py @@ -10,6 +10,7 @@ def main(): parser.add_argument("--prompt", type=str, required=True, help="The textual description of the image to generate.") parser.add_argument("--controlnet-image-path", type=str, required=True, help="Local path of the image to use as input for controlnet.") parser.add_argument("--controlnet-strength", type=float, default=0.4, help="Controls how strongly the control image influences the output image. A value of 0.0 means no influence. (Default is 0.4)") + parser.add_argument("--controlnet-save-canny", action="store_true", help="If set, save the Canny edge detection reference input image.") parser.add_argument("--output", type=str, default="image.png", help="The filename for the output image. Default is \"image.png\".") parser.add_argument("--model", "-m", type=str, required=True, choices=["dev", "schnell"], help="The model to use (\"schnell\" or \"dev\").") parser.add_argument("--seed", type=int, default=None, help="Entropy Seed (Default is time-based random-seed)") @@ -45,7 +46,9 @@ def main(): image = flux.generate_image( seed=int(time.time()) if args.seed is None else args.seed, prompt=args.prompt, + output=args.output, control_image_path=args.controlnet_image_path, + save_control_image_canny=args.controlnet_save_canny, config=ConfigControlnet( num_inference_steps=args.steps, height=args.height, diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index 190b4ad..52f94de 100644 --- a/src/mflux/post_processing/image_util.py +++ b/src/mflux/post_processing/image_util.py @@ -87,7 +87,12 @@ class ImageUtil: return Image.open(path) @staticmethod - def save(image: PIL.Image.Image, path: str, metadata: dict, export_json_metadata: bool = False) -> None: + def save( + image: PIL.Image.Image, + path: str, + metadata: dict | None = None, + export_json_metadata: bool = False + ) -> None: # fmt: off file_path = Path(path) file_path.parent.mkdir(parents=True, exist_ok=True) file_name = file_path.stem @@ -111,8 +116,9 @@ class ImageUtil: json.dump(metadata, json_file, indent=4) # Embed metadata - ImageUtil._embed_metadata(metadata, file_path) - log.info(f"Metadata embedded successfully at: {file_path}") + if metadata is not None: + ImageUtil._embed_metadata(metadata, file_path) + log.info(f"Metadata embedded successfully at: {file_path}") except Exception as e: log.error(f"Error saving image: {e}")