jbilcke-hf HF staff commited on
Commit
deae345
·
1 Parent(s): 5fb05e6

work in progress on the task system

Browse files
src/config.mts ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import path from "node:path"
2
+
3
+ export const storagePath = `${process.env.VS_STORAGE_PATH || './sandbox'}`
4
+
5
+ export const tasksDirPath = path.join(storagePath, "tasks")
6
+ export const pendingTasksDirFilePath = path.join(tasksDirPath, "pending")
7
+ export const completedTasksDirFilePath = path.join(tasksDirPath, "completed")
8
+
9
+ export const videosDirPath = path.join(storagePath, "videos")
10
+ export const pendingVideosDirFilePath = path.join(videosDirPath, "pending")
11
+ export const completedVideosDirFilePath = path.join(videosDirPath, "completed")
12
+
13
+ export const shotFormatVersion = 1
14
+ export const sequenceFormatVersion = 1
src/database/constants.mts DELETED
@@ -1,6 +0,0 @@
1
-
2
- export const pendingTasksDirFilePath = './database/pending/'
3
- export const completedTasksDirFilePath = './database/completed/'
4
-
5
- export const shotFormatVersion = 1
6
- export const sequenceFormatVersion = 1
 
 
 
 
 
 
 
src/database/saveCompletedTask.mts CHANGED
@@ -2,7 +2,7 @@ import { promises as fs } from "node:fs"
2
  import path from "path"
3
 
4
  import { VideoTask } from "../types.mts"
5
- import { completedTasksDirFilePath, pendingTasksDirFilePath } from "./constants.mts"
6
 
7
  export const saveCompletedTask = async (task: VideoTask) => {
8
  const fileName = `${task.id}.json`
 
2
  import path from "path"
3
 
4
  import { VideoTask } from "../types.mts"
5
+ import { completedTasksDirFilePath, pendingTasksDirFilePath } from "../config.mts"
6
 
7
  export const saveCompletedTask = async (task: VideoTask) => {
8
  const fileName = `${task.id}.json`
src/database/savePendingTask.mts CHANGED
@@ -2,7 +2,7 @@ import { promises as fs } from "node:fs"
2
  import path from "path"
3
 
4
  import { VideoTask } from "../types.mts"
5
- import { pendingTasksDirFilePath } from "./constants.mts"
6
 
7
  export const savePendingTask = async (task: VideoTask) => {
8
  const fileName = `${task.id}.json`
 
2
  import path from "path"
3
 
4
  import { VideoTask } from "../types.mts"
5
+ import { pendingTasksDirFilePath } from "../config.mts"
6
 
7
  export const savePendingTask = async (task: VideoTask) => {
8
  const fileName = `${task.id}.json`
src/database/updatePendingTask.mts ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import { promises as fs } from "node:fs"
2
+ import path from "path"
3
+
4
+ import { VideoTask } from "../types.mts"
5
+ import { pendingTasksDirFilePath } from "../config.mts"
6
+
7
+ export const updatePendingTask = async (task: VideoTask) => {
8
+ try {
9
+ const fileName = `${task.id}.json`
10
+ const filePath = path.join(pendingTasksDirFilePath, fileName)
11
+ await fs.writeFile(filePath, JSON.stringify(task, null, 2), "utf8")
12
+ } catch (err) {
13
+ console.error(`Failed to update the task. Probably an issue with the serialized object or the file system: ${err}`)
14
+ // we do not forward the exception, there is no need
15
+ // we will just try again the job later (even if it means losing a bit of data)
16
+ }
17
+ }
src/main.mts CHANGED
@@ -1,4 +1,5 @@
1
  import { getPendingTasks } from "./database/getPendingTasks.mts"
 
2
 
3
  export const main = async () => {
4
  const tasks = await getPendingTasks()
@@ -10,7 +11,11 @@ export const main = async () => {
10
  }
11
 
12
  console.log(`there are ${tasks.length} pending tasks`)
13
-
 
 
 
 
14
  setTimeout(() => {
15
  main()
16
  }, 1000)
 
1
  import { getPendingTasks } from "./database/getPendingTasks.mts"
2
+ import { processTask } from "./services/processTask.mts"
3
 
4
  export const main = async () => {
5
  const tasks = await getPendingTasks()
 
11
  }
12
 
13
  console.log(`there are ${tasks.length} pending tasks`)
14
+ for (const task of tasks) {
15
+ await processTask(task)
16
+ }
17
+ console.log(`processed ${tasks.length} tasks`)
18
+
19
  setTimeout(() => {
20
  main()
21
  }, 1000)
src/services/downloadVideo.mts CHANGED
@@ -1,15 +1,16 @@
1
  import path from 'node:path'
2
  import fs from 'node:fs'
3
-
4
- import tmpDir from 'temp-dir'
5
 
6
  export const downloadVideo = async (remoteUrl: string, fileName: string): Promise<string> => {
7
 
8
- const filePath = path.resolve(tmpDir, fileName)
9
 
10
  const controller = new AbortController()
11
  const timeoutId = setTimeout(() => controller.abort(), 15 * 60 * 60 * 1000) // 15 minutes
12
 
 
 
13
  // download the video
14
  const response = await fetch(remoteUrl, {
15
  signal: controller.signal
 
1
  import path from 'node:path'
2
  import fs from 'node:fs'
3
+ import { pendingVideosDirFilePath } from '../config.mts'
 
4
 
5
  export const downloadVideo = async (remoteUrl: string, fileName: string): Promise<string> => {
6
 
7
+ const filePath = path.resolve(pendingVideosDirFilePath, fileName)
8
 
9
  const controller = new AbortController()
10
  const timeoutId = setTimeout(() => controller.abort(), 15 * 60 * 60 * 1000) // 15 minutes
11
 
12
+ // TODO finish the timeout?
13
+
14
  // download the video
15
  const response = await fetch(remoteUrl, {
16
  signal: controller.signal
src/services/processTask.mts CHANGED
@@ -1,5 +1,68 @@
 
 
 
1
  import { VideoTask } from "../types.mts";
 
 
2
 
3
  export const processTask = async (task: VideoTask) => {
4
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  }
 
1
+ import { saveCompletedTask } from "../database/saveCompletedTask.mts";
2
+ import { savePendingTask } from "../database/savePendingTask.mts";
3
+ import { updatePendingTask } from "../database/updatePendingTask.mts";
4
  import { VideoTask } from "../types.mts";
5
+ import { downloadVideo } from "./downloadVideo.mts";
6
+ import { generateVideo } from "./generateVideo.mts";
7
 
8
  export const processTask = async (task: VideoTask) => {
9
+ console.log(`processing video task ${task.id}`)
10
+
11
+ // something isn't right, the task is already completed
12
+ if (task.completed) {
13
+ console.log(`video task ${task.id} is already completed`)
14
+ await saveCompletedTask(task)
15
+ return
16
+ }
17
+
18
+ let nbCompletedShots = 0
19
+ for (const shot of task.shots) {
20
+ // skip completed shots
21
+ if (shot.completed) {
22
+ nbCompletedShots++
23
+ continue
24
+ }
25
+
26
+ console.log(`need to complete shot ${shot.id}`)
27
+
28
+ const shotFileName = `${shot.id}.mp4`
29
+
30
+ if (!shot.hasGeneratedVideo) {
31
+ console.log("generating primordial pixel soup (raw video)..")
32
+ let generatedVideoUrl = ""
33
+
34
+ // currenty we cannot generate too many frames at once,
35
+ // otherwise the upscaler will have trouble
36
+
37
+ // so for now, we fix it to 24 frames
38
+ // const nbFramesForBaseModel = Math.min(3, Math.max(1, Math.round(duration))) * 8
39
+ const nbFramesForBaseModel = 24
40
+
41
+ try {
42
+ generatedVideoUrl = await generateVideo(shot.shotPrompt, {
43
+ seed: shot.seed,
44
+ nbFrames: nbFramesForBaseModel,
45
+ nbSteps: shot.steps,
46
+ })
47
+
48
+ console.log("downloading video..")
49
+
50
+ await downloadVideo(generatedVideoUrl, shotFileName)
51
+
52
+ } catch (err) {
53
+ // something is wrong, let's put the whole thing back into the queue
54
+ task.error = `failed to generate shot ${shot.id} (will try again later)`
55
+ await updatePendingTask(task)
56
+ break
57
+ }
58
+
59
+
60
+ }
61
+
62
+ if (!shot.hasUpscaledVideo) {
63
+
64
+ }
65
+
66
+ }
67
+
68
  }
src/types.mts CHANGED
@@ -167,8 +167,7 @@ export interface VideoShotData {
167
  completedAt: string
168
  completed: boolean
169
  error: string
170
- tmpFilePath: string
171
- finalFilePath: string
172
  }
173
 
174
  export type VideoShot = VideoShotMeta & VideoShotData
@@ -221,8 +220,7 @@ export interface VideoSequenceData {
221
  completedAt: string
222
  completed: boolean
223
  error: string
224
- tmpFilePath: string
225
- finalFilePath: string
226
  }
227
 
228
  export type VideoSequence = VideoSequenceMeta & VideoSequenceData
 
167
  completedAt: string
168
  completed: boolean
169
  error: string
170
+ filePath: string
 
171
  }
172
 
173
  export type VideoShot = VideoShotMeta & VideoShotData
 
220
  completedAt: string
221
  completed: boolean
222
  error: string
223
+ filePath: string
 
224
  }
225
 
226
  export type VideoSequence = VideoSequenceMeta & VideoSequenceData
src/utils/parseShotRequest.mts CHANGED
@@ -77,8 +77,7 @@ export const parseShotRequest = async (sequence: VideoSequence, maybeShotMeta: V
77
  completedAt: '',
78
  completed: false,
79
  error: '',
80
- tmpFilePath: '',
81
- finalFilePath: '',
82
  }
83
 
84
  return shot
 
77
  completedAt: '',
78
  completed: false,
79
  error: '',
80
+ filePath: '',
 
81
  }
82
 
83
  return shot
src/utils/parseVideoRequest.mts CHANGED
@@ -55,8 +55,7 @@ export const parseVideoRequest = async (request: VideoSequenceRequest): Promise<
55
  completedAt: null,
56
  completed: false,
57
  error: '',
58
- tmpFilePath: '',
59
- finalFilePath: '',
60
 
61
  // ------- the VideoShot -----
62
 
 
55
  completedAt: null,
56
  completed: false,
57
  error: '',
58
+ filePath: '',
 
59
 
60
  // ------- the VideoShot -----
61