monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Depth model refactor
This commit is contained in:
Filip Strand 2025-05-10 13:45:27 +02:00 committed by GitHub
commit 8fb1eb4116
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 44 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/mflux/flux/flux_initializer.py
View File

@ -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(

12
src/mflux/flux_tools/depth/depth_util.py
View File

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

22
src/mflux/models/depth_pro/depth_pro.py
View File

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

19
src/mflux/models/depth_pro/depth_pro_initializer.py
View File

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

17
src/mflux/save_depth.py
View File

@ -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__":

6
tests/depth/test_depth_pro.py
View File

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