diff --git a/src/mflux/callbacks/instances/battery_saver.py b/src/mflux/callbacks/instances/battery_saver.py index f773334..4e51258 100644 --- a/src/mflux/callbacks/instances/battery_saver.py +++ b/src/mflux/callbacks/instances/battery_saver.py @@ -3,7 +3,6 @@ import logging import re import subprocess - from mflux.callbacks.callback import BeforeLoopCallback from mflux.error.exceptions import StopImageGenerationException diff --git a/src/mflux/generate_redux.py b/src/mflux/generate_redux.py index b539988..8265b46 100644 --- a/src/mflux/generate_redux.py +++ b/src/mflux/generate_redux.py @@ -1,5 +1,4 @@ from argparse import Namespace - from pathlib import Path from mflux import Config, ModelConfig, StopImageGenerationException diff --git a/src/mflux/models/depth_pro/depth_pro_encoder.py b/src/mflux/models/depth_pro/depth_pro_encoder.py index feb41ed..5186d1b 100644 --- a/src/mflux/models/depth_pro/depth_pro_encoder.py +++ b/src/mflux/models/depth_pro/depth_pro_encoder.py @@ -21,7 +21,12 @@ class DepthProEncoder(nn.Module): self.upsample_lowres = nn.ConvTranspose2d(in_channels=1024, out_channels=1024, kernel_size=2, stride=2, padding=0, bias=True) # fmt: off self.fuse_lowres = nn.Conv2d(in_channels=1024 * 2, out_channels=1024, kernel_size=1, stride=1, padding=0, bias=True) # fmt: off - def __call__(self, x0: mx.array, x1: mx.array, x2: mx.array) -> list[mx.array]: + def __call__( + self, + x0: mx.array, + x1: mx.array, + x2: mx.array, + ) -> tuple[mx.array, mx.array, mx.array, mx.array, mx.array]: # 1: Run the backbone patch encoder model x_pyramid_patches = mx.concatenate((x0, x1, x2), axis=0) x_pyramid_encodings, backbone_highres_hook0, backbone_highres_hook1 = self.patch_encoder(x_pyramid_patches) @@ -55,13 +60,13 @@ class DepthProEncoder(nn.Module): x_global_features = mx.concatenate((x2_features, x_global_features), axis=1) x_global_features = ConvUtils.apply_conv(x_global_features, self.fuse_lowres) - return [ + return ( x_latent0_features, x_latent1_features, x0_features, x1_features, x_global_features, - ] + ) @staticmethod def _reshape_feature( diff --git a/src/mflux/models/depth_pro/depth_pro_model.py b/src/mflux/models/depth_pro/depth_pro_model.py index b8aa75a..a751c13 100644 --- a/src/mflux/models/depth_pro/depth_pro_model.py +++ b/src/mflux/models/depth_pro/depth_pro_model.py @@ -14,6 +14,6 @@ class DepthProModel(nn.Module): self.head = FOVHead() def __call__(self, x0: mx.array, x1: mx.array, x2: mx.array) -> tuple[mx.array, mx.array]: - encodings = self.encoder(x0, x1, x2) - features = self.decoder(encodings) - return self.head(features) + x0_lat, x1_lat, x0_feat, x1_feat, x_global = self.encoder(x0, x1, x2) + decoded = self.decoder(x0_lat, x1_lat, x0_feat, x1_feat, x_global) + return self.head(decoded) diff --git a/src/mflux/models/depth_pro/multires_conv_decoder.py b/src/mflux/models/depth_pro/multires_conv_decoder.py index 8f34fd4..7551ba3 100644 --- a/src/mflux/models/depth_pro/multires_conv_decoder.py +++ b/src/mflux/models/depth_pro/multires_conv_decoder.py @@ -23,16 +23,29 @@ class MultiresConvDecoder(nn.Module): FeatureFusionBlock2d(num_features=256, deconv=True), ] - def __call__(self, encodings: list[mx.array]) -> tuple[mx.array, mx.array]: - # Process last layer - encodings_last = encodings[4] - features = ConvUtils.apply_conv(encodings_last, self.convs[4]) + def __call__( + self, + x0_latent: mx.array, + x1_latent: mx.array, + x0_features: mx.array, + x1_features: mx.array, + x_global_features: mx.array, + ) -> mx.array: + # Process global features: + features = ConvUtils.apply_conv(x_global_features, self.convs[4]) features = self.fusions[4](features) - # Process remaining levels with skip connections - for i in [3, 2, 1, 0]: - enc = encodings[i] - features_i = ConvUtils.apply_conv(enc, self.convs[i]) - features = self.fusions[i](features, features_i) + # Process remaining levels with skip connections: + x1_skip_features = ConvUtils.apply_conv(x1_features, self.convs[3]) + features = self.fusions[3](features, x1_skip_features) + + x0_skip_features = ConvUtils.apply_conv(x0_features, self.convs[2]) + features = self.fusions[2](features, x0_skip_features) + + x1_skip_latents = ConvUtils.apply_conv(x1_latent, self.convs[1]) + features = self.fusions[1](features, x1_skip_latents) + + x0_skip_latents = ConvUtils.apply_conv(x0_latent, self.convs[0]) + features = self.fusions[0](features, x0_skip_latents) return features diff --git a/src/mflux/models/depth_pro/weight_handler_depth_pro.py b/src/mflux/models/depth_pro/weight_handler_depth_pro.py index 2b4a959..d932889 100644 --- a/src/mflux/models/depth_pro/weight_handler_depth_pro.py +++ b/src/mflux/models/depth_pro/weight_handler_depth_pro.py @@ -11,6 +11,8 @@ from mlx.utils import tree_unflatten from mflux.weights.weight_handler import MetaData from mflux.weights.weight_util import WeightUtil +logger = logging.getLogger(__name__) + class WeightHandlerDepthPro: def __init__(self, weights: dict, meta_data: MetaData): @@ -51,13 +53,13 @@ class WeightHandlerDepthPro: # 2. Download if model doesn't exist if not model_path.exists(): - logging.info("Downloading Depth Pro model from Apple...") + logger.info("Downloading Depth Pro model from Apple...") try: urllib.request.urlretrieve(APPLE_MODEL_URL, model_path) - logging.info(f"Downloaded model to {model_path}") + logger.info(f"Downloaded model to {model_path}") except (urllib.error.URLError, urllib.error.HTTPError) as e: - logging.error(f"Failed to download model: {e}") - logging.info(f"Please manually download from: {APPLE_MODEL_URL}") + logger.error(f"Failed to download model: {e}") + logger.info(f"Please manually download from: {APPLE_MODEL_URL}") if not model_path.exists(): raise FileNotFoundError(f"Model file not found at {model_path}")