license import '../flags' ; import { env } from '../environment' ; import { assert } from '../util' ; import { arrayBufferToBase64String , base64StringToArrayBuffer , getModelArtifactsInfoForJSON } from './io_utils' ; import { IORouter , IORouterRegistry } from './router_registry' ; import { IOHandler , ModelArtifacts , ModelArtifactsInfo , ModelJSON , ModelStoreManager , SaveResult } from './types' ; const PATH_SEPARATOR = '/' ; const PATH_PREFIX = 'tensorflowjs_models' ; const INFO_SUFFIX = 'info' ; const MODEL_TOPOLOGY_SUFFIX = 'model_topology' ; const WEIGHT_SPECS_SUFFIX = 'weight_specs' ; const WEIGHT_DATA_SUFFIX = 'weight_data' ; const MODEL_METADATA_SUFFIX = 'model_metadata' ; returns export function purgeLocalStorageArtifacts ( ) { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) || typeof window === 'undefined' || typeof window . localStorage === 'undefined' ) { throw new Error ( 'purgeLocalStorageModels() cannot proceed because local storage is ' + 'unavailable in the current environment.' ) ; } const LS = window . localStorage ; const purgedModelPaths = [ ] ; for ( let i = 0 ; i < LS . length ; ++ i ) { const key = LS . key ( i ) ; const prefix = PATH_PREFIX + PATH_SEPARATOR ; if ( key . startsWith ( prefix ) && key . length > prefix . length ) { LS . removeItem ( key ) ; const modelName = getModelPathFromKey ( key ) ; if ( purgedModelPaths . indexOf ( modelName ) === - 1 ) { purgedModelPaths . push ( modelName ) ; } } } return purgedModelPaths ; } type LocalStorageKeys = { /** Key of the localStorage entry storing `ModelArtifactsInfo`. */ info : string ; /**
     * Key of the localStorage entry storing the 'modelTopology' key of
     * `model.json`
     */ topology : string ; /**
     * Key of the localStorage entry storing the `weightsManifest.weights` entries
     * of `model.json`
     */ weightSpecs : string ; /** Key of the localStorage entry storing the weight data in Base64 */ weightData : string ; see ModelMetadata modelMetadata : string ; } ; type ModelMetadata = Omit < ModelJSON , 'modelTopology' | 'weightsManifest' > ; function getModelKeys ( path ) { return { info : [ PATH_PREFIX , path , INFO_SUFFIX ] . join ( PATH_SEPARATOR ) , topology : [ PATH_PREFIX , path , MODEL_TOPOLOGY_SUFFIX ] . join ( PATH_SEPARATOR ) , weightSpecs : [ PATH_PREFIX , path , WEIGHT_SPECS_SUFFIX ] . join ( PATH_SEPARATOR ) , weightData : [ PATH_PREFIX , path , WEIGHT_DATA_SUFFIX ] . join ( PATH_SEPARATOR ) , modelMetadata : [ PATH_PREFIX , path , MODEL_METADATA_SUFFIX ] . join ( PATH_SEPARATOR ) } ; } function removeItems ( keys ) { for ( const key of Object . values ( keys ) ) { window . localStorage . removeItem ( key ) ; } } param key function getModelPathFromKey ( key ) { const items = key . split ( PATH_SEPARATOR ) ; if ( items . length < 3 ) { throw new Error ( `Invalid key format: ${ key }` ) ; } return items . slice ( 1 , items . length - 1 ) . join ( PATH_SEPARATOR ) ; } function maybeStripScheme ( key ) { return key . startsWith ( BrowserLocalStorage . URL_SCHEME ) ? key . slice ( BrowserLocalStorage . URL_SCHEME . length ) : key ; } /**
 * IOHandler subclass: Browser Local Storage.
 *
 * See the doc string to `browserLocalStorage` for more details.
 */ export class BrowserLocalStorage implements IOHandler { protected readonly LS : Storage ; protected readonly modelPath : string ; protected readonly keys : LocalStorageKeys ; static readonly URL_SCHEME = 'localstorage://' ; constructor ( modelPath ) { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) || typeof window === 'undefined' || typeof window . localStorage === 'undefined' ) { throw new Error ( 'The current environment does not support local storage.' ) ; } this . LS = window . localStorage ; if ( modelPath == null || ! modelPath ) { throw new Error ( 'For local storage, modelPath must not be null, undefined or empty.' ) ; } this . modelPath = modelPath ; this . keys = getModelKeys ( this . modelPath ) ; } param modelArtifacts returns async save ( modelArtifacts ) { if ( modelArtifacts . modelTopology instanceof ArrayBuffer ) { throw new Error ( 'BrowserLocalStorage.save() does not support saving model topology ' + 'in binary formats yet.' ) ; } else { const topology = JSON . stringify ( modelArtifacts . modelTopology ) ; const weightSpecs = JSON . stringify ( modelArtifacts . weightSpecs ) ; const modelArtifactsInfo = getModelArtifactsInfoForJSON ( modelArtifacts ) ; try { this . LS . setItem ( this . keys . info , JSON . stringify ( modelArtifactsInfo ) ) ; this . LS . setItem ( this . keys . topology , topology ) ; this . LS . setItem ( this . keys . weightSpecs , weightSpecs ) ; this . LS . setItem ( this . keys . weightData , arrayBufferToBase64String ( modelArtifacts . weightData ) ) ; const metadata = { format : modelArtifacts . format , generatedBy : modelArtifacts . generatedBy , convertedBy : modelArtifacts . convertedBy , signature : modelArtifacts . signature != null ? modelArtifacts . signature : undefined , userDefinedMetadata : modelArtifacts . userDefinedMetadata != null ? modelArtifacts . userDefinedMetadata : undefined , modelInitializer : modelArtifacts . modelInitializer != null ? modelArtifacts . modelInitializer : undefined , trainingConfig : modelArtifacts . trainingConfig != null ? modelArtifacts . trainingConfig : undefined } ; this . LS . setItem ( this . keys . modelMetadata , JSON . stringify ( metadata ) ) ; return { modelArtifactsInfo } ; } catch ( err ) { removeItems ( this . keys ) ; throw new Error ( `Failed to save model '${ this . modelPath }' to local storage: ` + `size quota being exceeded is a possible cause of this failure: ` + `modelTopologyBytes=${ modelArtifactsInfo . modelTopologyBytes }, ` + `weightSpecsBytes=${ modelArtifactsInfo . weightSpecsBytes }, ` + `weightDataBytes=${ modelArtifactsInfo . weightDataBytes }.` ) ; } } } returns async load ( ) { const info = JSON . parse ( this . LS . getItem ( this . keys . info ) ) as ModelArtifactsInfo ; if ( info == null ) { throw new Error ( `In local storage, there is no model with name '${ this . modelPath }'` ) ; } if ( info . modelTopologyType !== 'JSON' ) { throw new Error ( 'BrowserLocalStorage does not support loading non-JSON model ' + 'topology yet.' ) ; } const out = { } ; const topology = JSON . parse ( this . LS . getItem ( this . keys . topology ) ) ; if ( topology == null ) { throw new Error ( `In local storage, the topology of model '${ this . modelPath }' ` + `is missing.` ) ; } out . modelTopology = topology ; const weightSpecs = JSON . parse ( this . LS . getItem ( this . keys . weightSpecs ) ) ; if ( weightSpecs == null ) { throw new Error ( `In local storage, the weight specs of model '${ this . modelPath }' ` + `are missing.` ) ; } out . weightSpecs = weightSpecs ; const metadataString = this . LS . getItem ( this . keys . modelMetadata ) ; if ( metadataString != null ) { const metadata = JSON . parse ( metadataString ) as ModelMetadata ; out . format = metadata . format ; out . generatedBy = metadata . generatedBy ; out . convertedBy = metadata . convertedBy ; if ( metadata . signature != null ) { out . signature = metadata . signature ; } if ( metadata . userDefinedMetadata != null ) { out . userDefinedMetadata = metadata . userDefinedMetadata ; } if ( metadata . modelInitializer != null ) { out . modelInitializer = metadata . modelInitializer ; } if ( metadata . trainingConfig != null ) { out . trainingConfig = metadata . trainingConfig ; } } const weightDataBase64 = this . LS . getItem ( this . keys . weightData ) ; if ( weightDataBase64 == null ) { throw new Error ( `In local storage, the binary weight values of model ` + `'${ this . modelPath }' are missing.` ) ; } out . weightData = base64StringToArrayBuffer ( weightDataBase64 ) ; return out ; } } export const localStorageRouter = ( url ) => { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) ) { return null ; } else { if ( ! Array . isArray ( url ) && url . startsWith ( BrowserLocalStorage . URL_SCHEME ) ) { return browserLocalStorage ( url . slice ( BrowserLocalStorage . URL_SCHEME . length ) ) ; } else { return null ; } } } ; IORouterRegistry . registerSaveRouter ( localStorageRouter ) ; IORouterRegistry . registerLoadRouter ( localStorageRouter ) ; param modelPath returns export function browserLocalStorage ( modelPath ) { return new BrowserLocalStorage ( modelPath ) ; } export class BrowserLocalStorageManager implements ModelStoreManager { private readonly LS : Storage ; constructor ( ) { assert ( env ( ) . getBool ( 'IS_BROWSER' ) , ( ) => 'Current environment is not a web browser' ) ; assert ( typeof window === 'undefined' || typeof window . localStorage !== 'undefined' , ( ) => 'Current browser does not appear to support localStorage' ) ; this . LS = window . localStorage ; } async listModels ( ) { const out = { } ; const prefix = PATH_PREFIX + PATH_SEPARATOR ; const suffix = PATH_SEPARATOR + INFO_SUFFIX ; for ( let i = 0 ; i < this . LS . length ; ++ i ) { const key = this . LS . key ( i ) ; if ( key . startsWith ( prefix ) && key . endsWith ( suffix ) ) { const modelPath = getModelPathFromKey ( key ) ; out [ modelPath ] = JSON . parse ( this . LS . getItem ( key ) ) as ModelArtifactsInfo ; } } return out ; } async removeModel ( path ) { path = maybeStripScheme ( path ) ; const keys = getModelKeys ( path ) ; if ( this . LS . getItem ( keys . info ) == null ) { throw new Error ( `Cannot find model at path '${ path }'` ) ; } const info = JSON . parse ( this . LS . getItem ( keys . info ) ) as ModelArtifactsInfo ; removeItems ( keys ) ; return info ; } }