26 lines
1.2 KiB
Python
26 lines
1.2 KiB
Python
import mlx.nn as nn
|
|
|
|
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(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")
|
|
WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample0")
|
|
WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample1")
|
|
WeightHandlerDepthPro.reposition_encoder_weights(depth_pro_weights, "upsample2")
|
|
WeightHandlerDepthPro.reposition_head_weights(depth_pro_weights)
|
|
WeightHandlerDepthPro.reshape_transposed_convolution_weights(depth_pro_weights)
|
|
|
|
# 2. Assign the weights to the model
|
|
depth_pro_model.update(depth_pro_weights.weights)
|
|
|
|
# 3. Optionally quantize the model
|
|
if quantize:
|
|
nn.quantize(depth_pro_model, bits=quantize)
|