refactoring

src/production/renderAnalysis.mts

@@ -0,0 +1,19 @@
+
+import { RenderedScene, RenderRequest } from "../types.mts"
+
+import { renderImageAnalysis } from "./renderImageAnalysis.mts"
+
+export async function renderAnalysis(request: RenderRequest, response: RenderedScene) {
+  
+  if (request.analyze) {
+    const isVideo = request?.nbFrames > 1
+
+    // note: this only works on images for now,
+    // but we could also analyze the first video frame to get ourselves an idea
+    const optionalAnalysisFn = !isVideo
+      ? renderImageAnalysis(request, response)
+      : Promise.resolve()
+
+    await optionalAnalysisFn
+  }
+}

src/production/renderContent.mts

@@ -0,0 +1,25 @@
+
+import { RenderedScene, RenderRequest } from "../types.mts"
+
+import { renderImage } from "./renderImage.mts"
+import { renderVideo } from "./renderVideo.mts"
+
+export async function renderContent(request: RenderRequest, response: RenderedScene) {
+  const isVideo = request?.nbFrames > 1
+
+  const renderContentFn = isVideo
+    ? renderVideo
+    : renderImage
+ 
+  try {
+    await renderContentFn(request, response)
+  } catch (err) {
+    // console.log(`renderContent() failed, trying a 2nd time..`)
+    try {
+      await renderContentFn(request, response)
+    } catch (err2) {
+      // console.log(`renderContent() failed, trying a 3th time..`)
+      await renderContentFn(request, response)
+    }
+  }
+}

src/production/renderPipeline.mts

@@ -1,53 +1,19 @@
 
 import { RenderedScene, RenderRequest } from "../types.mts"
 
-import { renderImage } from "./renderImage.mts"
-import { renderVideo } from "./renderVideo.mts"
-import { renderImageSegmentation } from "./renderImageSegmentation.mts"
-import { renderVideoSegmentation } from "./renderVideoSegmentation.mts"
-import { renderImageUpscaling } from "./renderImageUpscaling.mts"
 import { saveRenderedSceneToCache } from "../utils/filesystem/saveRenderedSceneToCache.mts"
-import { renderImageAnalysis } from "./renderImageAnalysis.mts"
+import { renderSegmentation } from "./renderSegmentation.mts"
+import { renderUpscaling } from "./renderUpscaling.mts"
+import { renderContent } from "./renderContent.mts"
+import { renderAnalysis } from "./renderAnalysis.mts"
 
 export async function renderPipeline(request: RenderRequest, response: RenderedScene) {
-  const isVideo = request?.nbFrames > 1
-
-  const renderContent = isVideo ? renderVideo : renderImage
-  const renderSegmentation  = isVideo ? renderVideoSegmentation : renderImageSegmentation 
-
-  if (isVideo) {
-    // console.log(`rendering a video..`)
-  } else {
-    // console.log(`rendering an image..`)
-  }
-
-  try {
-    await renderContent(request, response)
-  } catch (err) {
-    // console.log(`renderContent() failed, trying a 2nd time..`)
-    try {
-      await renderContent(request, response)
-    } catch (err2) {
-      // console.log(`renderContent() failed, trying a 3th time..`)
-      await renderContent(request, response)
-    }
-  }
-
-  // we upscale images with esrgan
-  // and for videos, well.. let's just skip this part,
-  // but later we could use Zeroscope V2 XL maybe?
-  const optionalUpscalingStep = isVideo
-    ? Promise.resolve()
-    : renderImageUpscaling(request, response)
-
-  const optionalAnalysisStep = request.analyze
-    ? renderImageAnalysis(request, response)
-    : Promise.resolve()
+  await renderContent(request, response)
 
   await Promise.all([
     renderSegmentation(request, response),
-    optionalAnalysisStep,
-    optionalUpscalingStep
+    renderAnalysis(request, response),
+    renderUpscaling(request, response)
   ])
 
   /*

src/production/renderSegmentation.mts

@@ -0,0 +1,18 @@
+
+import { RenderedScene, RenderRequest } from "../types.mts"
+
+import { renderImageSegmentation } from "./renderImageSegmentation.mts"
+import { renderVideoSegmentation } from "./renderVideoSegmentation.mts"
+
+export async function renderSegmentation(request: RenderRequest, response: RenderedScene) {
+  
+  if (request.segmentation === "firstframe" || request.segmentation === "allframes") {
+    const isVideo = request?.nbFrames > 1
+
+    const renderSegmentationFn = isVideo
+      ? renderVideoSegmentation
+      : renderImageSegmentation 
+
+    await renderSegmentationFn(request, response)
+  }
+}

src/production/renderUpscaling.mts

@@ -0,0 +1,20 @@
+
+import { RenderedScene, RenderRequest } from "../types.mts"
+
+import { renderImageUpscaling } from "./renderImageUpscaling.mts"
+import { renderVideoUpscaling } from "./renderVideoUpscaling.mts"
+
+export async function renderUpscaling(request: RenderRequest, response: RenderedScene) {
+
+  if (request.upscalingFactor > 1) {
+
+    const isVideo = request?.nbFrames > 1
+
+    // we upscale images with esrgan, and video with Zeroscope XL
+    const renderFn = isVideo
+      ? renderVideoUpscaling
+      : renderImageUpscaling
+
+    await renderFn(request, response)
+  }
+}

src/production/renderVideoUpscaling.mts

@@ -0,0 +1,26 @@
+import { upscaleVideoToBase64URL } from "../providers/video-upscaling/upscaleVideoToBase64URL.mts"
+import { RenderedScene, RenderRequest } from "../types.mts"
+
+export async function renderVideoUpscaling(
+  request: RenderRequest,
+  response: RenderedScene,
+): Promise<RenderedScene> {
+  
+  try {
+    // note: this converts a base64 PNG to a base64 JPG (which is good, actually!)
+    response.assetUrl = await upscaleVideoToBase64URL(response.assetUrl, request.prompt)
+    // console.log(`upscaling worked on the first try!`)
+  } catch (err) {
+    // console.error(`upscaling failed the first time.. let's try again..`)
+    try {
+      response.assetUrl = await upscaleVideoToBase64URL(response.assetUrl, request.prompt)
+      // console.log(`upscaling worked on the second try!`)
+    } catch (err) {
+      console.error(`upscaling failed on the second attempt.. let's keep the low-res image then :|`)
+      // no need to log a catastrophic failure here, since we still have the original (low-res image)
+      // to work with
+    }
+  }
+
+  return response
+}

src/providers/image-segmentation/segmentImage.mts

@@ -2,7 +2,7 @@ import puppeteer from "puppeteer"
 
 import { sleep } from "../../utils/misc/sleep.mts"
 import { ImageSegment } from "../../types.mts"
-import { downloadImageAsBase64 } from "../../utils/download/downloadFileAsBase64.mts"
+import { downloadFileAsBase64 } from "../../utils/download/downloadFileAsBase64.mts"
 import { resizeBase64Image } from "../../utils/image/resizeBase64Image.mts"
 
 // we don't use replicas yet, because it ain't easy to get their hostname
@@ -78,7 +78,7 @@ export async function segmentImage(
     // const tmpMaskFileName = `${uuidv4()}.png`
     // await downloadFileToTmp(maskUrl, tmpMaskFileName)
 
-    const rawPngInBase64 = await downloadImageAsBase64(tmpMaskDownloadUrl)
+    const rawPngInBase64 = await downloadFileAsBase64(tmpMaskDownloadUrl)
 
     const maskUrl = await resizeBase64Image(rawPngInBase64, width, height)
 

src/providers/video-upscaling/upscaleVideo.mts

@@ -32,7 +32,7 @@ export async function upscaleVideo(fileName: string, prompt: string) {
     })
 
     const secretField = await page.$('input[type=text]')
-    await secretField.type(prompt)
+    await secretField.type(secretToken)
 
     const promptField = await page.$('textarea')
     await promptField.type(prompt)

src/providers/video-upscaling/upscaleVideoToBase64URL.mts

@@ -0,0 +1,79 @@
+import path from "node:path"
+
+
+import puppeteer from "puppeteer"
+
+import { pendingFilesDirFilePath } from '../../config.mts'
+import { downloadFileAsBase64URL } from "../../utils/download/downloadFileAsBase64URL.mts"
+
+const instances: string[] = [
+  `${process.env.VC_VIDEO_UPSCALE_SPACE_API_URL_1 || ""}`
+].filter(instance => instance?.length > 0)
+
+const secretToken = `${process.env.VC_MICROSERVICE_SECRET_TOKEN || ""}`
+
+// TODO we should use an inference endpoint instead (or a space which bakes generation + upscale at the same time)
+export async function upscaleVideoToBase64URL(fileName: string, prompt: string) {
+  const instance = instances.shift()
+  instances.push(instance)
+
+  const browser = await puppeteer.launch({
+    // headless: true,
+    protocolTimeout: 800000,
+  })
+
+  try {
+    const page = await browser.newPage()
+
+    await page.goto(instance, {
+      waitUntil: 'networkidle2',
+    })
+
+    const secretField = await page.$('input[type=text]')
+    await secretField.type(secretToken)
+
+    const promptField = await page.$('textarea')
+    await promptField.type(prompt)
+
+    const inputFilePath = path.join(pendingFilesDirFilePath, fileName)
+    // console.log(`local file to upscale: ${inputFilePath}`)
+    
+    // await new Promise(r => setTimeout(r, 1000))
+
+    const fileField = await page.$('input[type=file]')
+
+    // console.log(`uploading file..`)
+    await fileField.uploadFile(inputFilePath)
+
+    // console.log('looking for the button to submit')
+    const submitButton = await page.$('button.lg')
+
+    // console.log('clicking on the button')
+    await submitButton.click()
+
+    /*
+    const client = await page.target().createCDPSession()
+
+    await client.send('Page.setDownloadBehavior', {
+      behavior: 'allow',
+      downloadPath: tmpDir,
+    })
+    */
+
+    await page.waitForSelector('a[download="xl_result.mp4"]', {
+      timeout: 800000, // need to be large enough in case someone else attemps to use our space
+    })
+
+    const upscaledFileUrl = await page.$$eval('a[download="xl_result.mp4"]', el => el.map(x => x.getAttribute("href"))[0])
+
+    // we download the whole file
+    // it's only a few seconds of video, so it should be < 2MB
+    const assetUrl = await downloadFileAsBase64URL(upscaledFileUrl)
+
+    return assetUrl
+  } catch (err) {
+    throw err
+  } finally {
+    await browser.close()
+  }
+}

src/utils/download/downloadFileAsBase64.mts

@@ -1,4 +1,4 @@
-export const downloadImageAsBase64 = async (remoteUrl: string): Promise<string> => {
+export const downloadFileAsBase64 = async (remoteUrl: string): Promise<string> => {
   const controller = new AbortController()
 
   // download the image

src/utils/download/downloadFileAsBase64URL.mts

@@ -0,0 +1,20 @@
+export const downloadFileAsBase64URL = async (remoteUrl: string): Promise<string> => {
+  const controller = new AbortController()
+
+  // download the file
+  const response = await fetch(remoteUrl, {
+    signal: controller.signal
+  })
+
+  // get as Buffer
+  const arrayBuffer = await response.arrayBuffer()
+  const buffer = Buffer.from(arrayBuffer)
+
+  // convert it to base64
+  const base64 = buffer.toString('base64')
+
+  const contentType = response.headers.get('content-type')
+
+  const assetUrl = `data:${contentType};base64,${base64}`
+  return assetUrl
+};