Add ability to save the canny controlnet image early on in the generation

(useful for inspecting if edge detection is as expected)

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

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)

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

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,

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}")