Better init pattern for depth model

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(

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)

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)

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)
 
 

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