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/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 2a4ae70..d8eae0c 100644 --- a/src/mflux/models/depth_pro/depth_pro.py +++ b/src/mflux/models/depth_pro/depth_pro.py @@ -1,12 +1,13 @@ +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 +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 @@ -19,15 +20,18 @@ class DepthResult: max_depth: float -class DepthPro(nn.Module): - def __init__(self): - super().__init__() - self.depth_pro_model = DepthProModel() +class DepthPro: + def __init__(self, quantize: int | None = None): + 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): @@ -37,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/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..cd54517 100644 --- a/src/mflux/save_depth.py +++ b/src/mflux/save_depth.py @@ -1,5 +1,6 @@ -from mflux.flux_tools.depth.depth_util import DepthUtil -from mflux.models.depth_pro.depth_pro_initializer import DepthProInitializer +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 - depth_pro = DepthProInitializer.init(quantize=args.quantize) - DepthUtil.get_or_create_depth_map(depth_pro=depth_pro, image_path=args.image_path) + # 1. Create the depth map + depth_pro = DepthPro(quantize=args.quantize) + 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 c6d8c90..17e9be2 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,10 +17,10 @@ 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)) + depth_result = depth_pro.create_depth_map(str(input_image_path)) # Save the output image depth_result.depth_image.save(output_image_path)