Merge pull request #187 from filipstrand/depth-model-refactor

Depth model refactor

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
+        # 2. Initialize the DepthPro model
-        flux_model.depth_pro = DepthProInitializer.init()
+        flux_model.depth_pro = DepthPro()
 
     @staticmethod
     def init_redux(

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
+        # 2. Generate a depth map from the image
-        depth_result = depth_pro(image_path)
+        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

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):
+class DepthPro:
-    def __init__(self):
+    def __init__(self, quantize: int | None = None):
-        super().__init__()
+        self._depth_pro_model = DepthProModel()
-        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)
+        depth = self._depth_pro_model(input_array)
-        return self.post_process(depth, height=height, width=width)
+        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)

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:
+    def init(depth_pro_model: DepthProModel, quantize: int | None = None) -> None:
-        # 1. Initialize the model
+        # 1. Load the weights
-        depth_pro = DepthPro()
-
-        # 2. 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
+        # 2. Assign the weights to the model
-        depth_pro.depth_pro_model.update(depth_pro_weights.weights)
+        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)
+            nn.quantize(depth_pro_model, bits=quantize)
-
-        return depth_pro

src/mflux/save_depth.py

@@ -1,5 +1,6 @@
-from mflux.flux_tools.depth.depth_util import DepthUtil
+from pathlib import Path
-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
 
 
@@ -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 = 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)
+    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__":

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)