fixup! discojs-core/train/models: initial

server/src/tasks.ts

@@ -4,8 +4,8 @@ import fs from 'node:fs/promises'
 import tf from '@tensorflow/tfjs'
 import '@tensorflow/tfjs-node'
 
-import type { Model, Task, Path, Digest, TaskProvider } from '@epfml/discojs-core'
-import { isTaskProvider, defaultTasks, models, serialization } from '@epfml/discojs-core'
+import type { Task, Path, Digest, TaskProvider } from '@epfml/discojs-core'
+import { Model, isTaskProvider, defaultTasks, models, serialization } from '@epfml/discojs-core'
 
 // default tasks and added ones
 // register 'taskAndModel' event to get tasks
@@ -97,19 +97,25 @@ export class TasksAndModels {
   }
 
   async addTaskAndModel (task: Task | TaskProvider, model?: Model | URL): Promise<void> {
+    let discoTask: Task
     if (isTaskProvider(task)) {
-      task = task.getTask()
+      discoTask = task.getTask()
+    } else {
+      discoTask = task
     }
 
+    let tfModel: Model
     if (model === undefined) {
-      model = await this.loadModelFromTask(task)
+      tfModel = await this.loadModelFromTask(task)
+    } else if (model instanceof Model) {
+      tfModel = model
     } else if (model instanceof URL) {
-      model = new models.TFJS(await tf.loadLayersModel(model.href))
+      tfModel = new models.TFJS(await tf.loadLayersModel(model.href))
     } else {
       throw new Error('invalid model')
     }
 
-    this.tasksAndModels = this.tasksAndModels.add([task, model])
-    this.emit('taskAndModel', task, model)
+    this.tasksAndModels = this.tasksAndModels.add([discoTask, tfModel])
+    this.emit('taskAndModel', discoTask, tfModel)
   }
 }