From 4cb8a5b2039fccce3f4a94cb229ef441e884fe2d Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 10 May 2025 13:20:45 +0200 Subject: [PATCH 1/2] Better init pattern for depth model --- src/mflux/flux/flux_initializer.py | 6 +++--- src/mflux/models/depth_pro/depth_pro.py | 4 +++- .../models/depth_pro/depth_pro_initializer.py | 19 +++++++------------ src/mflux/save_depth.py | 4 ++-- tests/depth/test_depth_pro.py | 4 ++-- 5 files changed, 17 insertions(+), 20 deletions(-) diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index adab0d2..89d9a64 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/src/mflux/flux/flux_initializer.py @@ -2,7 +2,7 @@ from mflux import ModelConfig from mflux.controlnet.transformer_controlnet import TransformerControlnet from mflux.controlnet.weight_handler_controlnet import WeightHandlerControlnet from mflux.flux_tools.redux.weight_handler_redux import WeightHandlerRedux -from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer +from mflux.models.depth_pro.depth_pro import DepthPro from mflux.models.redux_encoder.redux_encoder import ReduxEncoder from mflux.models.siglip_vision_transformer.siglip_vision_transformer import SiglipVisionTransformer from mflux.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder @@ -115,8 +115,8 @@ class FluxInitializer: lora_repo_id=lora_repo_id, ) - # 2. Initialize the DepthPro model and assign the weights - flux_model.depth_pro = DepthProInitializer.init() + # 2. Initialize the DepthPro model + flux_model.depth_pro = DepthPro() @staticmethod def init_redux( diff --git a/src/mflux/models/depth_pro/depth_pro.py b/src/mflux/models/depth_pro/depth_pro.py index 2a4ae70..4855785 100644 --- a/src/mflux/models/depth_pro/depth_pro.py +++ b/src/mflux/models/depth_pro/depth_pro.py @@ -7,6 +7,7 @@ import numpy as np import torch from PIL import Image +from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer from mflux.models.depth_pro.depth_pro_model import DepthProModel from mflux.post_processing.image_util import ImageUtil @@ -20,9 +21,10 @@ class DepthResult: class DepthPro(nn.Module): - def __init__(self): + def __init__(self, quantize: int | None = None): super().__init__() self.depth_pro_model = DepthProModel() + DepthProInitializer.init(self.depth_pro_model, quantize=quantize) def __call__(self, image_path: str | Path, resize: bool = True) -> DepthResult: input_array, height, width = self._pre_process(image_path) diff --git a/src/mflux/models/depth_pro/depth_pro_initializer.py b/src/mflux/models/depth_pro/depth_pro_initializer.py index ba7b8c4..7d69527 100644 --- a/src/mflux/models/depth_pro/depth_pro_initializer.py +++ b/src/mflux/models/depth_pro/depth_pro_initializer.py @@ -1,16 +1,13 @@ import mlx.nn as nn -from mflux.models.depth_pro.depth_pro import DepthPro +from mflux.models.depth_pro.depth_pro_model import DepthProModel from mflux.models.depth_pro.weight_handler_depth_pro import WeightHandlerDepthPro class DepthProInitializer: @staticmethod - def init(quantize: int | None = None) -> DepthPro: - # 1. Initialize the model - depth_pro = DepthPro() - - # 2. Load the weights + def init(depth_pro_model: DepthProModel, quantize: int | None = None) -> None: + # 1. Load the weights depth_pro_weights = WeightHandlerDepthPro.load_weights() WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent0") WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample_latent1") @@ -20,11 +17,9 @@ class DepthProInitializer: WeightHandlerDepthPro.reposition_head_weights(depth_pro_weights) WeightHandlerDepthPro.reshape_transposed_convolution_weights(depth_pro_weights) - # 3. Assign the weights to the model - depth_pro.depth_pro_model.update(depth_pro_weights.weights) + # 2. Assign the weights to the model + depth_pro_model.update(depth_pro_weights.weights) - # 4. Optionally quantize the model + # 3. Optionally quantize the model if quantize: - nn.quantize(depth_pro.depth_pro_model, bits=quantize) - - return depth_pro + nn.quantize(depth_pro_model, bits=quantize) diff --git a/src/mflux/save_depth.py b/src/mflux/save_depth.py index 1e6e7fd..d38d7e2 100644 --- a/src/mflux/save_depth.py +++ b/src/mflux/save_depth.py @@ -1,5 +1,5 @@ from mflux.flux_tools.depth.depth_util import DepthUtil -from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer +from mflux.models.depth_pro.depth_pro import DepthPro from mflux.ui.cli.parsers import CommandLineParser @@ -10,7 +10,7 @@ def main(): args = parser.parse_args() # 1. Create and save the depth map - depth_pro = DepthProInitializer.init(quantize=args.quantize) + depth_pro = DepthPro(quantize=args.quantize) DepthUtil.get_or_create_depth_map(depth_pro=depth_pro, image_path=args.image_path) diff --git a/tests/depth/test_depth_pro.py b/tests/depth/test_depth_pro.py index c6d8c90..25b2b96 100644 --- a/tests/depth/test_depth_pro.py +++ b/tests/depth/test_depth_pro.py @@ -4,7 +4,7 @@ from pathlib import Path import numpy as np from PIL import Image -from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer +from mflux.models.depth_pro.depth_pro import DepthPro class TestDepthPro: @@ -17,7 +17,7 @@ class TestDepthPro: try: # Initialize DepthPro model - depth_pro = DepthProInitializer.init() + depth_pro = DepthPro() # Process the image to generate a depth map depth_result = depth_pro(str(input_image_path)) From 42b4fd6d90c8bbbc7634556e5e727e5ea710d788 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 10 May 2025 13:33:02 +0200 Subject: [PATCH 2/2] Refactor DepthPro object --- src/mflux/flux_tools/depth/depth_util.py | 12 ++++++------ src/mflux/models/depth_pro/depth_pro.py | 20 +++++++++++--------- src/mflux/save_depth.py | 13 ++++++++++--- tests/depth/test_depth_pro.py | 2 +- 4 files changed, 28 insertions(+), 19 deletions(-) diff --git a/src/mflux/flux_tools/depth/depth_util.py b/src/mflux/flux_tools/depth/depth_util.py index 521122d..89fa96c 100644 --- a/src/mflux/flux_tools/depth/depth_util.py +++ b/src/mflux/flux_tools/depth/depth_util.py @@ -23,7 +23,7 @@ class DepthUtil: image_path: str | Path | None = None, depth_image_path: str | Path | None = None, ) -> tuple[mx.array, PIL.Image.Image]: - # 1. Create the depth map or use existing one + # 1. Get an existing depth map or create a new one depth_image_path, depth_image = DepthUtil.get_or_create_depth_map( depth_pro=depth_pro, image_path=image_path, @@ -47,21 +47,21 @@ class DepthUtil: depth_pro: DepthPro, image_path: str | Path | None = None, depth_map_path: str | Path | None = None, - ) -> tuple[str | Path, PIL.Image.Image | None]: + ) -> tuple[str | Path, PIL.Image.Image]: # 1. If a depth map path is provided, use it directly if depth_map_path: if not os.path.exists(depth_map_path): raise FileNotFoundError(f"Depth map file not found: {depth_map_path}") return depth_map_path, None + if not image_path: raise ValueError("Either depth_map_path or image_path must be provided") - if not os.path.exists(image_path): - raise FileNotFoundError(f"Image file not found: {image_path}") - # 2. Generate a depth map from the image using the DepthPro model - depth_result = depth_pro(image_path) + # 2. Generate a depth map from the image + depth_result = depth_pro.create_depth_map(image_path=image_path) # 3. Save the depth map to a file with the same name + _depth suffix generated_depth_path = str(image_path).rsplit(".", 1)[0] + "_depth.png" depth_result.depth_image.save(generated_depth_path) + return generated_depth_path, depth_result.depth_image diff --git a/src/mflux/models/depth_pro/depth_pro.py b/src/mflux/models/depth_pro/depth_pro.py index 4855785..d8eae0c 100644 --- a/src/mflux/models/depth_pro/depth_pro.py +++ b/src/mflux/models/depth_pro/depth_pro.py @@ -1,8 +1,8 @@ +import os from dataclasses import dataclass from pathlib import Path import mlx.core as mx -import mlx.nn as nn import numpy as np import torch from PIL import Image @@ -20,16 +20,18 @@ class DepthResult: max_depth: float -class DepthPro(nn.Module): +class DepthPro: def __init__(self, quantize: int | None = None): - super().__init__() - self.depth_pro_model = DepthProModel() - DepthProInitializer.init(self.depth_pro_model, quantize=quantize) + self._depth_pro_model = DepthProModel() + DepthProInitializer.init(self._depth_pro_model, quantize=quantize) + + def create_depth_map(self, image_path: str | Path) -> DepthResult: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") - def __call__(self, image_path: str | Path, resize: bool = True) -> DepthResult: input_array, height, width = self._pre_process(image_path) - depth = self.depth_pro_model(input_array) - return self.post_process(depth, height=height, width=width) + depth = self._depth_pro_model(input_array) + return self._post_process(depth, height=height, width=width) @staticmethod def _pre_process(image_path): @@ -39,7 +41,7 @@ class DepthPro(nn.Module): return input_array, image.height, image.width @staticmethod - def post_process(depth: mx.array, height: int, width: int): + def _post_process(depth: mx.array, height: int, width: int): depth_min = mx.min(depth) depth_max = mx.max(depth) normalized_depth = (depth - depth_min) / (depth_max - depth_min) diff --git a/src/mflux/save_depth.py b/src/mflux/save_depth.py index d38d7e2..cd54517 100644 --- a/src/mflux/save_depth.py +++ b/src/mflux/save_depth.py @@ -1,4 +1,5 @@ -from mflux.flux_tools.depth.depth_util import DepthUtil +from pathlib import Path + from mflux.models.depth_pro.depth_pro import DepthPro from mflux.ui.cli.parsers import CommandLineParser @@ -9,9 +10,15 @@ def main(): parser.add_save_depth_arguments() args = parser.parse_args() - # 1. Create and save the depth map + # 1. Create the depth map depth_pro = DepthPro(quantize=args.quantize) - DepthUtil.get_or_create_depth_map(depth_pro=depth_pro, image_path=args.image_path) + depth_result = depth_pro.create_depth_map(image_path=args.image_path) + + # 2. Save the depth map with the same name + _depth suffix + image_path = Path(args.image_path) + output_path = image_path.with_stem(f"{image_path.stem}_depth").with_suffix(".png") + depth_result.depth_image.save(output_path) + print(f"Depth map saved to: {output_path}") if __name__ == "__main__": diff --git a/tests/depth/test_depth_pro.py b/tests/depth/test_depth_pro.py index 25b2b96..17e9be2 100644 --- a/tests/depth/test_depth_pro.py +++ b/tests/depth/test_depth_pro.py @@ -20,7 +20,7 @@ class TestDepthPro: depth_pro = DepthPro() # Process the image to generate a depth map - depth_result = depth_pro(str(input_image_path)) + depth_result = depth_pro.create_depth_map(str(input_image_path)) # Save the output image depth_result.depth_image.save(output_image_path)