git clone https://github.com/Jorgvt/paramperceptnet.git
cd paramperceptnet
pip install -e .

import json

from huggingface_hub import hf_hub_download
import flax
import orbax.checkpoint
from ml_collections import ConfigDict

from paramperceptnet.models import Baseline as PerceptNet

config_path = hf_hub_download(repo_id="Jorgvt/ppnet-baseline",
                              filename="config.json")
with open(config_path, "r") as f:
    config = ConfigDict(json.load(f))

from safetensors.flax import load_file

weights_path = hf_hub_download(repo_id="Jorgvt/ppnet-baseline",
                               filename="weights.safetensors")
variables = load_file(weights_path)
variables = flax.traverse_util.unflatten_dict(variables, sep=".")
params = variables["params"]

weights_path = hf_hub_download(repo_id="Jorgvt/ppnet-fully-trained",
                               filename="weights.msgpack")
with open(weights_path, "rb") as f:
    variables = orbax.checkpoint.msgpack_utils.msgpack_restore(f.read())
variables = jax.tree_util.tree_map(lambda x: jnp.array(x), variables)
params = variables["params"]

from jax import numpy as jnp
model = PerceptNet(config)
pred = model.apply({"params": params}, jnp.ones((1,384,512,3)))