bmz_io

def export_to_bmz(
    model: Union[FCNModule, VAEModule],
    config: Configuration,
    path_to_archive: Union[Path, str],
    model_name: str,
    general_description: str,
    data_description: str,
    authors: list[dict],
    input_array: np.ndarray,
    output_array: np.ndarray,
    covers: Optional[list[Union[Path, str]]] = None,
    channel_names: Optional[list[str]] = None,
    model_version: str = "0.1.0",
) -> None:
    """Export the model to BioImage Model Zoo format.

    Arrays are expected to be SC(Z)YX with singleton dimensions allowed for S and C.

    `model_name` should consist of letters, numbers, dashes, underscores and parentheses
    only.

    Parameters
    ----------
    model : CAREamicsModule
        CAREamics model to export.
    config : Configuration
        Model configuration.
    path_to_archive : Union[Path, str]
        Path to the output file.
    model_name : str
        Model name.
    general_description : str
        General description of the model.
    data_description : str
        Description of the data the model was trained on.
    authors : list[dict]
        Authors of the model.
    input_array : np.ndarray
        Input array, should not have been normalized.
    output_array : np.ndarray
        Output array, should have been denormalized.
    covers : list of pathlib.Path or str, default=None
        Paths to the cover images.
    channel_names : Optional[list[str]], optional
        Channel names, by default None.
    model_version : str, default="0.1.0"
        Model version.
    """
    path_to_archive = Path(path_to_archive)

    if path_to_archive.suffix != ".zip":
        raise ValueError(
            f"Path to archive must point to a zip file, got {path_to_archive}."
        )

    if not path_to_archive.parent.exists():
        path_to_archive.parent.mkdir(parents=True, exist_ok=True)

    # versions
    careamics_version = get_careamics_version()

    # save files in temporary folder
    with tempfile.TemporaryDirectory() as tmpdirname:
        temp_path = Path(tmpdirname)

        # create environment file
        # TODO move in bioimage module
        env_path = temp_path / "environment.yml"
        env_path.write_text(create_env_text(PYTORCH_VERSION, TORCHVISION_VERSION))

        # export input and ouputs
        inputs = temp_path / "inputs.npy"
        np.save(inputs, input_array)
        outputs = temp_path / "outputs.npy"
        np.save(outputs, output_array)

        # export configuration
        config_path = save_configuration(config, temp_path / "careamics.yaml")

        # export model state dictionary
        weight_path = _export_state_dict(model, temp_path / "weights.pth")

        # export cover if necesary
        if covers is None:
            covers = [create_cover(temp_path, input_array, output_array)]

        # create model description
        model_description = create_model_description(
            config=config,
            name=model_name,
            general_description=general_description,
            data_description=data_description,
            authors=authors,
            inputs=inputs,
            outputs=outputs,
            weights_path=weight_path,
            torch_version=PYTORCH_VERSION,
            careamics_version=careamics_version,
            config_path=config_path,
            env_path=env_path,
            covers=covers,
            channel_names=channel_names,
            model_version=model_version,
        )

        # test model description
        test_kwargs = (
            model_description.config.get("bioimageio", {})
            .get("test_kwargs", {})
            .get("pytorch_state_dict", {})
        )
        summary: ValidationSummary = test_model(model_description, **test_kwargs)
        if summary.status == "failed":
            raise ValueError(f"Model description test failed: {summary}")

        # save bmz model
        save_bioimageio_package(model_description, output_path=path_to_archive)

def load_from_bmz(
    path: Union[Path, str, HttpUrl]
) -> tuple[Union[FCNModule, VAEModule], Configuration]:
    """Load a model from a BioImage Model Zoo archive.

    Parameters
    ----------
    path : Path, str or HttpUrl
        Path to the BioImage Model Zoo archive. A Http URL must point to a downloadable
        location.

    Returns
    -------
    FCNModel or VAEModel
        The loaded CAREamics model.
    Configuration
        The loaded CAREamics configuration.

    Raises
    ------
    ValueError
        If the path is not a zip file.
    """
    # load description, this creates an unzipped folder next to the archive
    model_desc = load_model_description(path)

    # extract paths
    weights_path, config_path = extract_model_path(model_desc)

    # load configuration
    config = load_configuration(config_path)

    # create careamics lightning module
    if config.algorithm_config.model.architecture == SupportedArchitecture.UNET:
        model = FCNModule(algorithm_config=config.algorithm_config)
    elif config.algorithm_config.model.architecture == SupportedArchitecture.LVAE:
        model = VAEModule(algorithm_config=config.algorithm_config)
    else:
        raise ValueError(
            f"Unsupported architecture {config.algorithm_config.model.architecture}"
        )  # TODO ugly ?

    # load model state dictionary
    _load_state_dict(model, weights_path)

    return model, config