working on SDXL + segmentation

src/index.mts

@@ -4,7 +4,7 @@ import path from "node:path"
 import { validate as uuidValidate } from "uuid"
 import express from "express"
 
-import { Video, VideoStatus, VideoAPIRequest, RenderRequest, RenderAPIResponse } from "./types.mts"
+import { Video, VideoStatus, VideoAPIRequest, RenderRequest, RenderedScene } from "./types.mts"
 import { parseVideoRequest } from "./utils/parseVideoRequest.mts"
 import { savePendingVideo } from "./scheduler/savePendingVideo.mts"
 import { getVideo } from "./scheduler/getVideo.mts"
@@ -48,8 +48,8 @@ app.post("/render", async (req, res) => {
     return
   }
 
-  let result: RenderAPIResponse = {
-    videoUrl: "",
+  let result: RenderedScene = {
+    assetUrl: "",
     maskBase64: "",
     error: "",
     segments: []
@@ -83,15 +83,51 @@ app.post("/render", async (req, res) => {
   }
 })
 
+// a "fast track" pipeline
 /*
 app.post("/segment", async (req, res) => {
-  const payload = req.body as 
+
+  const request = req.body as RenderRequest
+  console.log(req.body)
+
+  let result: RenderedScene = {
+    assetUrl: "",
+    maskBase64: "",
+    error: "",
+    segments: []
+  }
+  
   try {
-    await segmentImage()
+    result = await renderScene(request)
+  } catch (err) {
+    // console.log("failed to render scene!")
+    result.error = `failed to render scene: ${err}`
+  }
+
+  if (result.error === "already rendering") {
+    console.log("server busy")
+    res.status(200)
+    res.write(JSON.stringify({ url: "", error: result.error }))
+    res.end()
+    return
+  } else if (result.error.length > 0) {
+    // console.log("server error")
+    res.status(500)
+    res.write(JSON.stringify({ url: "", error: result.error }))
+    res.end()
+    return
+  } else {
+    // console.log("all good")
+    res.status(200)
+    res.write(JSON.stringify(result))
+    res.end()
+    return
   }
 })
 */
 
+
+
 app.post("/:ownerId", async (req, res) => {
   const request = req.body as VideoAPIRequest
 

src/production/renderScene.mts

@@ -1,104 +1,13 @@
-import { v4 as uuidv4 } from "uuid"
-
-import { ImageSegment, RenderAPIResponse, RenderRequest } from "../types.mts"
-import { downloadFileToTmp } from "../utils/downloadFileToTmp.mts"
-import { generateSeed } from "../utils/generateSeed.mts"
-import { getValidNumber } from "../utils/getValidNumber.mts"
-import { generateVideo } from "./generateVideo.mts"
-import { getFirstVideoFrame } from "../utils/getFirstVideoFrame.mts"
-import { segmentImage } from "../utils/segmentImage.mts"
-
-const state = {
-  isRendering: false
-}
-
-const seed = generateSeed()
-
-export async function renderScene(scene: RenderRequest): Promise<RenderAPIResponse> {
-  // console.log("renderScene")
-
-  // let's disable this for now
-  // this is only reliable if nothing crashes anyway..
-  /*
-  if (state.isRendering) {
-    // console.log("renderScene: isRendering")
-    return {
-      videoUrl: "",
-      error: "already rendering",
-      maskBase64: "",
-      segments: [],
-    }
-  }
-  */
-
-  // onsole.log("marking as isRendering")
-  state.isRendering = true
-
-  let url = ""
-  let error = ""
-
-  try {
-    url = await generateVideo(scene.prompt, {
-      seed: getValidNumber(scene.seed, 0, 4294967295, generateSeed()),
-      nbFrames: getValidNumber(scene.nbFrames, 8, 24, 16), // 2 seconds by default
-      nbSteps: getValidNumber(scene.nbSteps, 1, 50, 10), // use 10 by default to go fast, but not too sloppy
-    })
-    // console.log("successfull generation")
-    error = ""
-  } catch (err) {
-    error = `failed to render scene: ${err}`
+import { RenderedScene, RenderRequest } from "../types.mts"
+import { renderStaticScene } from "./renderStaticScene.mts"
+import { renderVideoScene } from "./renderVideoScene.mts"
+
+export async function renderScene(scene: RenderRequest): Promise<RenderedScene> {
+  if (scene?.nbFrames === 1) {
+    console.log(`calling renderStaticScene`)
+    return renderStaticScene(scene)
+  } else {
+    console.log(`calling renderVideoScene`)
+    return renderVideoScene(scene)
   }
-
-
-
-  // TODO add segmentation here
-  const actionnables = Array.isArray(scene.actionnables) ? scene.actionnables : []
-
-  let mask = ""
-  let segments: ImageSegment[] = []
-
-  if (actionnables.length > 0) {
-    console.log("we have some actionnables:", actionnables)
-    if (scene.segmentation === "firstframe") {
-      console.log("going to grab the first frame")
-      const tmpVideoFilePath = await downloadFileToTmp(url, `${uuidv4()}`)
-      console.log("downloaded the first frame to ", tmpVideoFilePath)
-      const firstFrameFilePath = await getFirstVideoFrame(tmpVideoFilePath)
-      console.log("downloaded the first frame to ", firstFrameFilePath)
-      
-      if (!firstFrameFilePath) {
-        console.error("failed to get the image")
-        error = "failed to segment the image"
-      } else {
-        console.log("got the first frame! segmenting..")
-        const result = await segmentImage(firstFrameFilePath, actionnables)
-        mask = result.pngInBase64
-        segments = result.segments
-        // console.log("success!", {  segments })
-      }
-      /*
-      const jpgBase64 = await getFirstVideoFrame(tmpVideoFileName)
-      if (!jpgBase64) {
-        console.error("failed to get the image")
-        error = "failed to segment the image"
-      } else {
-        console.log(`got the first frame (${jpgBase64.length})`)
-
-        console.log("TODO: call segmentImage with the base64 image")
-        await segmentImage()
-      }
-      */
-    }
-  }
-
-  // console.log("marking as not rendering anymore")
-  state.isRendering = false
-  error = ""
-
-  return {
-    videoUrl: url,
-    error,
-    maskBase64: mask,
-    segments
-  } as RenderAPIResponse
 }

src/production/renderStaticScene.mts

@@ -0,0 +1,84 @@
+import path from "node:path"
+
+import { v4 as uuidv4 } from "uuid"
+import tmpDir from "temp-dir"
+
+import { ImageSegment, RenderedScene, RenderRequest } from "../types.mts"
+import { downloadFileToTmp } from "../utils/downloadFileToTmp.mts"
+import { segmentImage } from "../utils/segmentImage.mts"
+import { generateImageSDXLAsBase64 } from "../utils/generateImageSDXL.mts"
+import { writeBase64ToFile } from "../utils/writeBase64ToFile.mts"
+
+export async function renderStaticScene(scene: RenderRequest): Promise<RenderedScene> {
+
+  let imageBase64 = ""
+  let error = ""
+
+  try {
+    console.log(`calling generateImageSDXLAsBase64 with: `, JSON.stringify({
+      positivePrompt: scene.prompt,
+      seed: scene.seed || undefined,
+      nbSteps: scene.nbSteps || undefined,
+      width: 1024,
+      height: 512
+    }, null, 2))
+    imageBase64 = await generateImageSDXLAsBase64({
+      positivePrompt: scene.prompt,
+      seed: scene.seed || undefined,
+      nbSteps: scene.nbSteps || undefined,
+      width: 1024,
+      height: 512
+    })
+    console.log("successful generation!", imageBase64.slice(0, 30))
+    error = ""
+    if (!imageBase64?.length) {
+      throw new Error(`the generated image is empty`)
+    }
+  } catch (err) {
+    error = `failed to render scene: ${err}`
+    return {
+      assetUrl: imageBase64,
+      error,
+      maskBase64: "",
+      segments: []
+    } as RenderedScene
+  }
+
+  const actionnables = Array.isArray(scene.actionnables) ? scene.actionnables : []
+
+  let mask = ""
+  let segments: ImageSegment[] = []
+
+  if (actionnables.length > 0) {
+    console.log("we have some actionnables:", actionnables)
+    console.log("going to grab the first frame")
+
+    const tmpImageFilePath = path.join(tmpDir, `${uuidv4()}.png`)
+
+    console.log("beginning:", imageBase64.slice(0, 100))
+    await writeBase64ToFile(imageBase64, tmpImageFilePath)
+    console.log("wrote the image to ", tmpImageFilePath)
+  
+    if (!tmpImageFilePath) {
+      console.error("failed to get the image")
+      error = "failed to segment the image"
+    } else {
+      console.log("got the first frame! segmenting..")
+      const result = await segmentImage(tmpImageFilePath, actionnables)
+      mask = result.pngInBase64
+      segments = result.segments
+      console.log("success!", {  segments })
+    }
+  } else {
+    console.log("no actionnables: just returning the image, then")
+  }
+
+  error = ""
+
+  return {
+    assetUrl: imageBase64,
+    error,
+    maskBase64: mask,
+    segments
+  } as RenderedScene
+}

src/production/renderVideoScene.mts

@@ -0,0 +1,81 @@
+import { v4 as uuidv4 } from "uuid"
+
+import { ImageSegment, RenderedScene, RenderRequest } from "../types.mts"
+import { downloadFileToTmp } from "../utils/downloadFileToTmp.mts"
+import { generateSeed } from "../utils/generateSeed.mts"
+import { getValidNumber } from "../utils/getValidNumber.mts"
+import { generateVideo } from "./generateVideo.mts"
+import { getFirstVideoFrame } from "../utils/getFirstVideoFrame.mts"
+import { segmentImage } from "../utils/segmentImage.mts"
+
+export async function renderVideoScene(scene: RenderRequest): Promise<RenderedScene> {
+
+  let url = ""
+  let error = ""
+
+  try {
+    url = await generateVideo(scene.prompt, {
+      seed: getValidNumber(scene.seed, 0, 2147483647, generateSeed()),
+      nbFrames: getValidNumber(scene.nbFrames, 8, 24, 16), // 2 seconds by default
+      nbSteps: getValidNumber(scene.nbSteps, 1, 50, 10), // use 10 by default to go fast, but not too sloppy
+    })
+    // console.log("successfull generation")
+    error = ""
+    if (!url?.length) {
+      throw new Error(`url for the generated image is empty`)
+    }
+  } catch (err) {
+    error = `failed to render scene: ${err}`
+  }
+
+
+
+  // TODO add segmentation here
+  const actionnables = Array.isArray(scene.actionnables) ? scene.actionnables : []
+
+  let mask = ""
+  let segments: ImageSegment[] = []
+
+  if (actionnables.length > 0) {
+    console.log("we have some actionnables:", actionnables)
+    if (scene.segmentation === "firstframe") {
+      console.log("going to grab the first frame")
+      const tmpVideoFilePath = await downloadFileToTmp(url, `${uuidv4()}`)
+      console.log("downloaded the first frame to ", tmpVideoFilePath)
+      const firstFrameFilePath = await getFirstVideoFrame(tmpVideoFilePath)
+      console.log("downloaded the first frame to ", firstFrameFilePath)
+      
+      if (!firstFrameFilePath) {
+        console.error("failed to get the image")
+        error = "failed to segment the image"
+      } else {
+        console.log("got the first frame! segmenting..")
+        const result = await segmentImage(firstFrameFilePath, actionnables)
+        mask = result.pngInBase64
+        segments = result.segments
+        // console.log("success!", {  segments })
+      }
+      /*
+      const jpgBase64 = await getFirstVideoFrame(tmpVideoFileName)
+      if (!jpgBase64) {
+        console.error("failed to get the image")
+        error = "failed to segment the image"
+      } else {
+        console.log(`got the first frame (${jpgBase64.length})`)
+
+        console.log("TODO: call segmentImage with the base64 image")
+        await segmentImage()
+      }
+      */
+    }
+  }
+
+  error = ""
+
+  return {
+    assetUrl: url,
+    error,
+    maskBase64: mask,
+    segments
+  } as RenderedScene
+}

src/types.mts

@@ -302,12 +302,13 @@ export interface ImageSegmentationRequest {
 export interface ImageSegment {
   id: number
   box: number[]
+  color: number[]
   label: string
   score: number 
 }
 
-export interface RenderAPIResponse {
-  videoUrl: string 
+export interface RenderedScene {
+  assetUrl: string 
   error: string
   maskBase64: string
   segments: ImageSegment[]

src/utils/generateImage.mts

@@ -0,0 +1,51 @@
+import { HfInference } from "@huggingface/inference"
+import { getValidNumber } from "./getValidNumber.mts";
+import { generateSeed } from "./generateSeed.mts";
+
+const hf = new HfInference(process.env.VC_HF_API_TOKEN)
+
+export async function generateImage(options: {
+  positivePrompt: string;
+  negativePrompt: string;
+  seed?: number;
+  width?: number;
+  height?: number;
+  nbSteps?: number;
+}) {
+  
+  const positivePrompt = options?.positivePrompt || ""
+  if (!positivePrompt) {
+    throw new Error("missing prompt")
+  }
+  const negativePrompt = options?.negativePrompt || ""
+  const seed = getValidNumber(options?.seed, 0, 2147483647, generateSeed())
+  const width = getValidNumber(options?.width, 256, 1024, 512)
+  const height = getValidNumber(options?.height, 256, 1024, 512)
+  const nbSteps = getValidNumber(options?.nbSteps, 5, 50, 25)
+
+  const blob = await hf.textToImage({
+    inputs: [
+      positivePrompt,
+      "bautiful",
+      "award winning",
+      "intricate details",
+      "high resolution"
+    ].filter(word => word)
+    .join(", "),
+    model: "stabilityai/stable-diffusion-2-1",
+    parameters: {
+      negative_prompt: [
+        negativePrompt,
+        "blurry",
+        // "artificial",
+       //  "cropped",
+        "low quality",
+        "ugly"
+      ].filter(word => word)
+      .join(", ")
+    }
+  })
+  const buffer = Buffer.from(await blob.arrayBuffer())
+
+  return buffer
+}

src/utils/generateImageSDXL.mts

@@ -0,0 +1,76 @@
+import { client } from "@gradio/client"
+
+import { generateSeed } from "../utils/generateSeed.mts"
+import { getValidNumber } from "./getValidNumber.mts"
+
+// we don't use replicas yet, because it ain't easy to get their hostname
+const instances: string[] = [
+  `${process.env.VC_SDXL_SPACE_API_URL_1 || ""}`,
+].filter(instance => instance?.length > 0)
+
+export async function generateImageSDXLAsBase64(options: {
+  positivePrompt: string;
+  negativePrompt?: string;
+  seed?: number;
+  width?: number;
+  height?: number;
+  nbSteps?: number;
+}) {
+
+  const positivePrompt = options?.positivePrompt || ""
+  if (!positivePrompt) {
+    throw new Error("missing prompt")
+  }
+  const negativePrompt = options?.negativePrompt || ""
+  const seed = getValidNumber(options?.seed, 0, 2147483647, generateSeed())
+  const width = getValidNumber(options?.width, 256, 1024, 512)
+  const height = getValidNumber(options?.height, 256, 1024, 512)
+  const nbSteps = getValidNumber(options?.nbSteps, 5, 100, 20)
+
+  const instance = instances.shift()
+  instances.push(instance)
+
+  const positive = [
+    positivePrompt,
+    "beautiful",
+    "award winning",
+    "intricate details",
+    "high resolution"
+  ].filter(word => word)
+  .join(", ")
+
+  const negative =  [
+    negativePrompt,
+    "blurry",
+    // "artificial",
+    // "cropped",
+    "low quality",
+    "ugly"
+  ].filter(word => word)
+  .join(", ")
+
+  const api = await client(instance, {
+    hf_token: `${process.env.VC_HF_API_TOKEN}` as any
+  })
+
+  
+  const rawResponse = (await api.predict("/run", [		
+    positive, // string  in 'Prompt' Textbox component		
+    negative, // string  in 'Negative prompt' Textbox component		
+    positive, // string  in 'Prompt 2' Textbox component		
+    negative, // string  in 'Negative prompt 2' Textbox component		
+    true, // boolean  in 'Use negative prompt' Checkbox component		
+    false, // boolean  in 'Use prompt 2' Checkbox component		
+    false, // boolean  in 'Use negative prompt 2' Checkbox component		
+    seed, // number (numeric value between 0 and 2147483647) in 'Seed' Slider component		
+    width, // number (numeric value between 256 and 1024) in 'Width' Slider component		
+    height, // number (numeric value between 256 and 1024) in 'Height' Slider component		
+    7, // number (numeric value between 1 and 20) in 'Guidance scale for base' Slider component		
+    7, // number (numeric value between 1 and 20) in 'Guidance scale for refiner' Slider component		
+    nbSteps, // number (numeric value between 10 and 100) in 'Number of inference steps for base' Slider component		
+    nbSteps, // number (numeric value between 10 and 100) in 'Number of inference steps for refiner' Slider component		
+    true, // boolean  in 'Apply refiner' Checkbox component
+  ])) as any
+    
+  return rawResponse?.data?.[0] as string
+}

src/utils/parseShotRequest.mts

@@ -44,7 +44,7 @@ export const parseShotRequest = async (sequence: VideoSequence, maybeShotMeta: P
     actorDialoguePrompt: `${maybeShotMeta.actorDialoguePrompt || ""}`,
 
     // a video sequence SHOULD NOT HAVE a consistent seed, to avoid weird geometry similarities
-    seed: getValidNumber(maybeShotMeta.seed, 0, 4294967295, generateSeed()),
+    seed: getValidNumber(maybeShotMeta.seed, 0, 2147483647, generateSeed()),
 
     // a video sequence SHOULD HAVE a consistent grain
     noise: sequence.noise,

src/utils/parseVideoRequest.mts

@@ -57,7 +57,7 @@ export const parseVideoRequest = async (ownerId: string, request: VideoAPIReques
     // describe the main actor dialogue line
     actorDialoguePrompt: `${request.sequence.actorDialoguePrompt || ''}`,
 
-    seed: getValidNumber(request.sequence.seed, 0, 4294967295, generateSeed()),
+    seed: getValidNumber(request.sequence.seed, 0, 2147483647, generateSeed()),
 
     noise: request.sequence.noise === true,
     noiseAmount: request.sequence.noise === true ? 2 : 0,

src/utils/segmentImage.mts

@@ -29,7 +29,7 @@ export async function segmentImage(
 
   const browser = await puppeteer.launch({
     headless: true,
-    protocolTimeout: 70000,
+    protocolTimeout: 120000,
   })
 
   const page = await browser.newPage()
@@ -42,8 +42,6 @@ export async function segmentImage(
   // console.log(`uploading file..`)
   await fileField.uploadFile(inputImageFilePath)
 
-  await sleep(500)
-
   const firstTextarea = await page.$('textarea[data-testid="textbox"]')
 
   const conceptsToDetect = actionnables.join(" . ")
@@ -52,13 +50,13 @@ export async function segmentImage(
   // console.log('looking for the button to submit')
   const submitButton = await page.$('button.lg')
 
-  await sleep(500)
+  await sleep(200)
 
   // console.log('clicking on the button')
   await submitButton.click()
 
   await page.waitForSelector('img[data-testid="detailed-image"]', {
-    timeout: 70000, // need to be large enough in case someone else attemps to use our space
+    timeout: 120000, // need to be large enough in case someone else attemps to use our space
   })
 
   const maskUrl = await page.$$eval('img[data-testid="detailed-image"]', el => el.map(x => x.getAttribute("src"))[0])

src/utils/segmentImageFromURL.mts

@@ -0,0 +1,25 @@
+import { v4 as uuidv4 } from "uuid"
+
+import { downloadFileToTmp } from "./downloadFileToTmp.mts"
+import { segmentImage } from "./segmentImage.mts"
+
+// TODO we should use an inference endpoint instead
+
+// note: on a large T4 (8 vCPU)
+// it takes about 30 seconds to compute
+export async function segmentImageFromURL(
+  inputUrl: string,
+  actionnables: string[]
+) {
+  if (!actionnables?.length) {
+    throw new Error("cannot segment image without actionnables!")
+  }
+  console.log(`segmenting image from URL: "${inputUrl}"`)
+  const tmpFileName = `${uuidv4()}`
+  const tmpFilePath = await downloadFileToTmp(inputUrl, tmpFileName)
+
+  const results = await segmentImage(tmpFilePath, actionnables)
+
+  console.log("image has been segmented!", results)
+  return results
+}

src/utils/writeBase64ToFile.mts

@@ -0,0 +1,18 @@
+import { promises as fs } from "node:fs"
+
+export async function writeBase64ToFile(content: string, filePath: string): Promise<void> {
+  
+  // Remove "data:image/png;base64," from the start of the data url
+  const base64Data = content.split(",")[1]
+
+  // Convert base64 to binary
+  const data = Buffer.from(base64Data, "base64")
+
+  // Write binary data to file
+  try {
+    await fs.writeFile(filePath, data)
+    console.log("File written successfully")
+  } catch (error) {
+    console.error("An error occurred:", error)
+  }
+}