add a system to recover from broken generation cycles

README.md

@@ -8,4 +8,5 @@ pinned: false
 app_port: 8000
 ---
 
-Media server 📡
+Media server 📡
+

src/getDatabase.mts

@@ -0,0 +1,9 @@
+import { promises as fs } from 'node:fs'
+
+import { Database } from './types.mts'
+
+export const getDatabase = async (filePath: string): Promise<Database> => {
+  const rawFile = await fs.readFile(filePath, 'utf8')
+  
+  return JSON.parse(rawFile) as Database
+}

src/getStats.mts

@@ -1,11 +1,13 @@
 import { promises as fs } from 'node:fs'
 
 export const getStats = async () => {
-  const videoFilePaths = await fs.readdir(process.env.WEBTV_VIDEO_STORAGE_PATH)
-  const audioFilePaths = await fs.readdir(process.env.WEBTV_AUDIO_STORAGE_PATH)
+  const videoFilePaths = await fs.readdir(process.env.WEBTV_VIDEO_STORAGE_PATH_NEXT)
+  const legacyVideoFilePaths = await fs.readdir(process.env.WEBTV_VIDEO_STORAGE_PATH)
+  const legacyAudioFilePaths = await fs.readdir(process.env.WEBTV_AUDIO_STORAGE_PATH)
 
   return {
     nbVideoFiles: videoFilePaths.length,
-    nbAudioFiles: audioFilePaths.length,
+    nbLegacyVideoFiles: legacyVideoFilePaths.length,
+    nbLegacyAudioFiles: legacyAudioFilePaths.length,
   }
 }

src/index.mts

@@ -2,85 +2,154 @@ import { v4 as uuid } from 'uuid'
 import { upscaleVideo } from './upscaleVideo.mts'
 import { keepVideo } from './keepVideo.mts'
 
-import { demoPrompts } from './prompts.mts'
 import { getStats } from './getStats.mts'
 import { enhanceVideo } from './enhanceVideo.mts'
 import { callZeroscope } from './callZeroscope.mts'
 import { downloadVideo } from './downloadVideo.mts'
+import { getDatabase } from './getDatabase.mts'
+import { callMusicgen } from './callMusicgen.mts'
+
+let hasReachedStartingPoint = false
+
+type RunMode = 'running' | 'paused' | 'dry_run'
+
+const status = `${process.env.WEBTV_STATUS || 'dry_run'}` as RunMode
+
+console.log(`Web TV server status: ${status}`)
 
 const main = async () => {
-  console.log('Generating ideas..')
-  const ideas = demoPrompts
+  if (status === 'paused') {
+    setTimeout(() => {
+      main()
+    }, 30000)
+    return
+  }
+
+  console.log('Reading persistent file structure..')
+  const stats = await getStats()
+  console.log(`New format: We have ${stats.nbVideoFiles} video files`)
+
+  console.log(`Legacy: We have ${stats.nbLegacyVideoFiles} video files and ${stats.nbLegacyAudioFiles} audio files`)
+
+  console.log('Reading prompt database..')
+  const db = await getDatabase('./database.json')
+
+  const nbTotalShots = db.sequences.reduce((a, s) => a + s.shots.length, 0)
+  console.log(`Prompt database version: ${db.version}`)
+  console.log(`We got ${db.sequences.length} sequences for ${nbTotalShots} shots in total`)
 
   console.log('Generating videos sequences..')
   const instanceId = process.env.WEBTV_WORKER_INSTANCE_ID || '0'
 
-  for (const { input, captions } of ideas) {
-    console.log(`\nVideo sequence to generate: ${input}`)
+  const startingPointExists = db.sequences.some(seq => seq.shots.some(shot => shot.shotId === db.startAtShotId))
+  
+  if (!startingPointExists) {
+    console.log(`Starting point ${db.startAtShotId} not found, we will start at the beginning`)
+    hasReachedStartingPoint = true
+  } else if (db.startAtShotId) {
+    console.log(`We are going to start at shot ${db.startAtShotId}`)
+  } else {
+    console.log('We are going to start at the beginning')
+  }
+
+  for (const sequence of db.sequences) {
+    const containsStartingPoint = sequence.shots.some(shot => shot.shotId === db.startAtShotId)
+
+    // we skip sequences that were already processed
+    if (!hasReachedStartingPoint && !containsStartingPoint) {
+      continue
+    }
+
+    console.log(`
+-----------------------------------------------------------
+Going to generate ${sequence.shots.length} for prompt:
+${sequence.videoPrompt}
+`)
 
     const sequenceId = uuid()
 
-    const silentVideos: string[] = []
+    const generatedShots: string[] = []
 
     // this is hardcoded everywhere for now, since videos longer than 3 sec require the Nvidia A100
     const videoDurationInSecs = 3
 
-    for (const caption of captions) {
-      console.log(`- generating shot: ${caption}`)
+    let shotIndex = 0
+
+    for (const shot of sequence.shots) {
+
+      if (shot.shotId === db.startAtShotId) {
+        hasReachedStartingPoint = true
+      }
+
+      if (!hasReachedStartingPoint) {
+        shotIndex++
+        continue
+      }
+
+      console.log(`- generating shot: ${shot.shotId}`)
+
+      if (status === 'dry_run') {
+        // console.log('DRY RUN')
+        shotIndex++
+        continue
+      }
+
       try {
-        const generatedVideoUrl = await callZeroscope(caption)
+        const generatedVideoUrl = await callZeroscope(shot.videoPrompt)
 
-        const videoName = `inst_${instanceId}_seq_${sequenceId}_shot_${Date.now()}.mp4`
+        const shotFileName = `inst_${instanceId}_seq_${sequenceId}_shot_${shotIndex++}_${Date.now()}.mp4`
 
-        console.log(`- downloading ${videoName} from ${generatedVideoUrl}`)
-        await downloadVideo(generatedVideoUrl, videoName)
+        console.log(`- downloading shot ${shotFileName} from ${generatedVideoUrl}`)
+        await downloadVideo(generatedVideoUrl, shotFileName)
 
-        console.log(`- downloaded ${videoName}`)
+        console.log(`- downloaded shot ${shotFileName}`)
 
-        console.log('- upscaling video..')
+        console.log('- upscaling shot..')
       
         try {
-          await upscaleVideo(videoName, caption)
+          await upscaleVideo(shotFileName, shot.videoPrompt)
         } catch (err) {
           // upscaling is finicky, if it fails we try again
-          console.log('- trying again to upscale video..')
-          await upscaleVideo(videoName, caption)
+          console.log('- trying again to upscale shot..')
+          await upscaleVideo(shotFileName, shot.videoPrompt)
         }
 
-        console.log('- enhancing video..')
-        await enhanceVideo(videoName)
+        console.log('- enhancing shot..')
+        await enhanceVideo(shotFileName)
 
-        console.log('- saving final video..')
-        await keepVideo(videoName, process.env.WEBTV_VIDEO_STORAGE_PATH_NEXT)
+        console.log('- saving final shot..')
+        await keepVideo(shotFileName, process.env.WEBTV_VIDEO_STORAGE_PATH_NEXT)
 
-        silentVideos.push(videoName)
+        generatedShots.push(shotFileName)
 
         console.log('- done!')
       } catch (err) {
         console.log(`- error: ${err}`)
       }
 
-      const totalRunTime = videoDurationInSecs * silentVideos.length
+      const totalRunTime = videoDurationInSecs * generatedShots.length
 
       if (totalRunTime <= 0) {
         continue
       }
 
-      console.log(`TODO: generate ${totalRunTime} seconds of music`)
       // TODO: generate music from MusicGen, with the correct length
       // (or we could generate a slightly larger track and let ffmpeg cut it)
+      console.log(`TODO: generate ${totalRunTime} seconds of music`)
+      await callMusicgen(sequence.audioPrompt) // this does nothing for now
 
-    }  
-
-    console.log('Finished gerating all video sequences')
+    }
 
-    console.log('Current stats:', await getStats())
+    console.log('Finished generating sequence')
   }
 
-  console.log('Finished the full cycle')
+  console.log('Finished the cycle')
+
+  hasReachedStartingPoint = true // set this to true in all cases
+
   setTimeout(() => {
     main()
-  }, 3000)
+  }, 2000)
 }
 
 setTimeout(() => {

src/types.mts

@@ -0,0 +1,23 @@
+export interface Shot {
+  shotId: string
+  index: number
+  lastGenerationAt: string
+  videoPrompt: string
+  audioPrompt: string
+}
+
+export interface Sequence {
+  sequenceId: string
+  lastGenerationAt: string
+  videoPrompt: string
+  audioPrompt: string
+  channel: string
+  tags: string[]
+  shots: Shot[]
+}
+
+export interface Database {
+  version: number
+  startAtShotId: string
+  sequences: Sequence[]
+}