model_registry.hpp

Go to the documentation of this file.

#pragma once
 
// Generic model registry primitive (promoted from tools/audio so plugins/apps —
// not just the CLI — can use it). A `ModelEntry` describes a downloadable ML model
// (HuggingFace or any URL); `resolve_checkpoint_url` turns an `hf://` ref into a
// direct HTTPS URL with path hardening. The available-models list (the catalog) is
// supplied by the consumer (e.g. tools/audio's audio catalog, or the Magenta demo's
// mrt2 catalog) — this module owns the type + resolution + lookup, not the data.
 
#include <cstdint>
#include <string>
#include <string_view>
#include <vector>
 
namespace pulp::runtime {
 
struct ModelAsset {
    std::string role;          
    std::string checkpoint_ref;
    std::string
        sha256; 
    std::uint64_t size_bytes = 0;
};
 
struct ModelEntry {
    std::string model_id;             
    std::string display_name;
    std::string description;
    std::string backend;              
    std::string checkpoint_ref;       
    std::vector<std::string> task_tags;
    std::uint64_t size_bytes = 0;
 
    std::string download_url;         
    std::string sha256;               
    bool auto_downloadable = true;    
 
    // Multi-asset bundles (e.g. Magenta: weights + vocab + encoder). When empty the
    // entry is single-file (checkpoint_ref). The downloader fetches all assets.
    std::vector<ModelAsset> assets;
 
    bool is_recommended = false;      
    std::string license;              
    std::string attribution;          
    std::string min_device;           
};

std::string resolve_checkpoint_url(const std::string& checkpoint_ref);

const ModelEntry* find_model(const std::vector<ModelEntry>& registry, std::string_view model_id);
 
}  // namespace pulp::runtime