license /**
 * Imports
 */ import * as tf from '@tensorflow/tfjs-core' ; import * as support from '../core/sketch_support' ; property property property property property export interface SketchRNNInfo { max_seq_len : number ; mode : number ; name : string ; scale_factor : number ; version : number ; } property property property property property property property export interface StrokePDF { pi : Float32Array ; muX : Float32Array ; muY : Float32Array ; sigmaX : Float32Array ; sigmaY : Float32Array ; corr : Float32Array ; pen : Float32Array ; } property property export interface LSTMState { c : Float32Array ; h : Float32Array ; } /**
 * Main SketchRNN model class.
 *
 * Implementation of decoder model in https://arxiv.org/abs/1704.03477
 *
 * TODO(hardmaru): make a "batch" continueSequence-like method
 * that runs fully on GPU.
 */ export class SketchRNN { private checkpointURL : string ; private forgetBias : tf . Scalar ; private initialized : boolean ; public info : SketchRNNInfo ; public numUnits : number ; public pixelFactor : number ; public scaleFactor : number ; private weights : Float32Array [ ] ; private weightDims : number [ ] [ ] ; private outputKernel : tf . Tensor2D ; private outputBias : tf . Tensor1D ; private lstmKernel : tf . Tensor2D ; private lstmBias : tf . Tensor1D ; private rawVars : tf . Tensor [ ] ; private NMIXTURE = 20 ; param checkpointURL constructor ( checkpointURL ) { this . checkpointURL = checkpointURL ; this . initialized = false ; } /**
     * Returns true if model is intialized.
     */ isInitialized ( ) { return this . initialized ; } /**
     * Instantiates class information inputs from the JSON model.
     * TODO(hardmaru): to add support for new tfjs checkpoints.
     */ private instantiateFromJSON ( info , weightDims , weightStrings ) { this . forgetBias = tf . scalar ( 1.0 ) ; this . info = info ; this . setPixelFactor ( 2.0 ) ; this . weightDims = weightDims ; this . numUnits = this . weightDims [ 0 ] [ 0 ] ; let rawWeights ; const maxWeight = 10.0 ; this . weights = [ ] ; for ( let i = 0 ; i < weightStrings . length ; i ++ ) { rawWeights = new Float32Array ( support . stringToArray ( weightStrings [ i ] ) ) ; const N = rawWeights . length ; for ( let j = 0 ; j < N ; j ++ ) { rawWeights [ j ] = maxWeight * rawWeights [ j ] / 32767 ; } this . weights . push ( rawWeights ) ; } this . outputKernel = tf . tensor2d ( this . weights [ 0 ] , [ this . weightDims [ 0 ] [ 0 ] , this . weightDims [ 0 ] [ 1 ] ] ) ; this . outputBias = tf . tensor1d ( this . weights [ 1 ] ) ; const lstmKernelXH = tf . tensor2d ( this . weights [ 2 ] , [ this . weightDims [ 2 ] [ 0 ] , this . weightDims [ 2 ] [ 1 ] ] ) ; const lstmKernelHH = tf . tensor2d ( this . weights [ 3 ] , [ this . weightDims [ 3 ] [ 0 ] , this . weightDims [ 3 ] [ 1 ] ] ) ; const axis = 0 ; this . lstmKernel = tf . concat2d ( [ lstmKernelXH , lstmKernelHH ] , axis ) ; this . lstmBias = tf . tensor1d ( this . weights [ 4 ] ) ; this . rawVars = [ this . outputKernel , this . outputBias , this . lstmKernel , this . lstmBias ] ; } /**
     * Loads variables from the JSON model
     */ async initialize ( ) { this . dispose ( ) ; const vars = await fetch ( this . checkpointURL ) . then ( ( response ) => response . json ( ) ) ; this . instantiateFromJSON ( vars [ 0 ] , vars [ 1 ] , vars [ 2 ] ) ; this . initialized = true ; console . log ( 'Initialized SketchRNN.' ) ; } dispose ( ) { if ( this . rawVars ) { for ( let i = 0 ; i < this . rawVars . length ; i ++ ) { this . rawVars [ i ] . dispose ( ) ; } this . rawVars = undefined ; } if ( this . forgetBias ) { this . forgetBias . dispose ( ) ; this . forgetBias = undefined ; } this . initialized = false ; } param scale returns setPixelFactor ( scale ) { this . pixelFactor = scale ; this . scaleFactor = this . info . scale_factor / this . pixelFactor ; } param stroke param state returns update ( stroke , state ) { const out = tf . tidy ( ( ) => { const numUnits = this . numUnits ; const s = this . scaleFactor ; const normStroke = [ stroke [ 0 ] / s , stroke [ 1 ] / s , stroke [ 2 ] , stroke [ 3 ] , stroke [ 4 ] ] ; const x = tf . tensor2d ( normStroke , [ 1 , 5 ] ) ; const c = tf . tensor2d ( state . c , [ 1 , numUnits ] ) ; const h = tf . tensor2d ( state . h , [ 1 , numUnits ] ) ; const newState = tf . basicLSTMCell ( this . forgetBias , this . lstmKernel , this . lstmBias , x , c , h ) ; return tf . concat ( newState , 1 ) ; } ) ; const newCH = out . dataSync ( ) ; out . dispose ( ) ; const newC = newCH . slice ( 0 , this . numUnits ) ; const newH = newCH . slice ( this . numUnits , this . numUnits * 2 ) ; const finalState = { c : new Float32Array ( newC ) , h : new Float32Array ( newH ) } ; return finalState ; } param strokes param state param steps returns updateStrokes ( strokes , state , steps ? ) { const out = tf . tidy ( ( ) => { const numUnits = this . numUnits ; const s = this . scaleFactor ; let normStroke ; let x ; let c ; let h ; let newState ; let numSteps = strokes . length ; if ( steps ) { numSteps = steps ; } c = tf . tensor2d ( state . c , [ 1 , numUnits ] ) ; h = tf . tensor2d ( state . h , [ 1 , numUnits ] ) ; for ( let i = 0 ; i < numSteps ; i ++ ) { normStroke = [ strokes [ i ] [ 0 ] / s , strokes [ i ] [ 1 ] / s , strokes [ i ] [ 2 ] , strokes [ i ] [ 3 ] , strokes [ i ] [ 4 ] ] ; x = tf . tensor2d ( normStroke , [ 1 , 5 ] ) ; newState = tf . basicLSTMCell ( this . forgetBias , this . lstmKernel , this . lstmBias , x , c , h ) ; c = newState [ 0 ] ; h = newState [ 1 ] ; } return tf . concat ( newState , 1 ) ; } ) ; const newCH = out . dataSync ( ) ; out . dispose ( ) ; const newC = newCH . slice ( 0 , this . numUnits ) ; const newH = newCH . slice ( this . numUnits , this . numUnits * 2 ) ; const finalState = { c : new Float32Array ( newC ) , h : new Float32Array ( newH ) } ; return finalState ; } param state param temperature param softmaxTemperature returns getPDF ( state , temperature = 0.65 , softmaxTemperature ? ) { const temp = temperature ; let discreteTemp = 0.5 + temp * 0.5 ; if ( softmaxTemperature ) { discreteTemp = softmaxTemperature ; } const NOUT = this . NMIXTURE ; const out = tf . tidy ( ( ) => { const numUnits = this . numUnits ; const h = tf . tensor2d ( state . h , [ 1 , numUnits ] ) ; const sqrttemp = tf . scalar ( Math . sqrt ( temp ) ) ; const softtemp = tf . scalar ( discreteTemp ) ; const z = tf . add ( tf . matMul ( h , this . outputKernel ) , this . outputBias ) . squeeze ( ) ; const [ rawPen , rst ] = tf . split ( z , [ 3 , NOUT * 6 ] ) ; const [ rawPi , mu1 , mu2 , rawSigma1 , rawSigma2 , rawCorr ] = tf . split ( rst , 6 ) ; const pen = tf . softmax ( rawPen . div ( softtemp ) ) ; const pi = tf . softmax ( rawPi . div ( softtemp ) ) ; const sigma1 = tf . exp ( rawSigma1 ) . mul ( sqrttemp ) ; const sigma2 = tf . exp ( rawSigma2 ) . mul ( sqrttemp ) ; const corr = tf . tanh ( rawCorr ) ; const result = [ pi , mu1 , mu2 , sigma1 , sigma2 , corr , pen ] ; return tf . concat ( result ) ; } ) ; const result = out . dataSync ( ) ; out . dispose ( ) ; const pdf = { pi : new Float32Array ( result . slice ( 0 , NOUT ) ) , muX : new Float32Array ( result . slice ( 1 * NOUT , 2 * NOUT ) ) , muY : new Float32Array ( result . slice ( 2 * NOUT , 3 * NOUT ) ) , sigmaX : new Float32Array ( result . slice ( 3 * NOUT , 4 * NOUT ) ) , sigmaY : new Float32Array ( result . slice ( 4 * NOUT , 15 * NOUT ) ) , corr : new Float32Array ( result . slice ( 5 * NOUT , 6 * NOUT ) ) , pen : new Float32Array ( result . slice ( 6 * NOUT , 6 * NOUT + 3 ) ) } ; return pdf ; } returns zeroState ( ) { const result = { c : new Float32Array ( this . numUnits ) , h : new Float32Array ( this . numUnits ) } ; return result ; } param rnnState returns copyState ( rnnState ) { const result = { c : new Float32Array ( rnnState . c ) , h : new Float32Array ( rnnState . h ) } ; return result ; } returns zeroInput ( ) { return [ 0 , 0 , 1 , 0 , 0 ] ; } param pdf returns sample ( pdf ) { const idx = support . sampleSoftmax ( pdf . pi ) ; const mu1 = pdf . muX [ idx ] ; const mu2 = pdf . muY [ idx ] ; const sigma1 = pdf . sigmaX [ idx ] ; const sigma2 = pdf . sigmaY [ idx ] ; const corr = pdf . corr [ idx ] ; const penIdx = support . sampleSoftmax ( pdf . pen ) ; const penstate = [ 0 , 0 , 0 ] ; penstate [ penIdx ] = 1 ; const delta = support . birandn ( mu1 , mu2 , sigma1 , sigma2 , corr ) ; const stroke = [ delta [ 0 ] * this . scaleFactor , delta [ 1 ] * this . scaleFactor , penstate [ 0 ] , penstate [ 1 ] , penstate [ 2 ] ] ; return stroke ; } param line param tolerance returns simplifyLine ( line , tolerance ? ) { return support . simplifyLine ( line , tolerance ) ; } param line param tolerance returns simplifyLines ( lines , tolerance ? ) { return support . simplifyLines ( lines , tolerance ) ; } param lines returns linesToStroke ( lines ) { return support . linesToStrokes ( lines ) ; } param line param lastPoint returns lineToStroke ( line , lastPoint ) { return support . lineToStroke ( line , lastPoint ) ; } }