Repository: koishijs/novelai-bot
File: src/index.ts (TypeScript)

Code before:
              // in case some client doesn't support r2 upload and follow the ye olde way.
              return forceDataPrefix(result.generations[0].img, 'image/webp')
            }
            const imgRes = await ctx.http(imgUrl, { responseType: 'arraybuffer' })
            const b64 = Buffer.from(imgRes.data).toString('base64')
            return forceDataPrefix(b64, imgRes.headers.get('content-type'))
          }
          if (config.type === 'comfyui') {
            // get filenames from history
            const promptId = res.data.prompt_id
            const check = () => ctx.http.get(trimSlash(config.endpoint) + '/history/' + promptId)
              .then((res) => res[promptId] && res[promptId].outputs)
            const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))
            let outputs
            while (!(outputs = await check())) {
              await sleep(config.pollInterval)
            }
            // get images by filename
            const imagesOutput: { data: ArrayBuffer, mime: string }[] = [];
            for (const nodeId in outputs) {
              const nodeOutput = outputs[nodeId]
              if ('images' in nodeOutput) {
                for (const image of nodeOutput['images']) {
                  const urlValues = new URLSearchParams({ filename: image['filename'], subfolder: image['subfolder'], type: image['type'] }).toString()
                  const imgRes = await ctx.http(trimSlash(config.endpoint) + '/view?' + urlValues)
                  imagesOutput.push({data: imgRes.data, mime: imgRes.headers.get('content-type')})
                  break
                }
              }
            }
            // return first image
            return forceDataPrefix(Buffer.from(imagesOutput[0].data).toString('base64'),imagesOutput[0].mime)
          }
          // event: newImage
          // id: 1
          // data:
          //                                                                        ↓ nai-v3
          if (res.headers.get('content-type') === 'application/x-zip-compressed' || res.headers.get('content-disposition')?.includes('.zip')) {
            const buffer = Buffer.from(res.data, 'binary')  // Ensure 'binary' encoding
            const zip = new AdmZip(buffer)

            // Gets all files in the ZIP file
            const zipEntries = zip.getEntries()
            const firstImageBuffer = zip.readFile(zipEntries[0])
            const b64 = Buffer.from(firstImageBuffer).toString('base64')
            return forceDataPrefix(b64, 'image/png')
          }
          return forceDataPrefix(res.data?.trimEnd().slice(27))
        }

        let dataUrl: string, count = 0

Review comment (suggestion):
```suggestion
                  imagesOutput.push({ data: imgRes.data, mime: imgRes.headers.get('content-type') })
```

Code after:
              // in case some client doesn't support r2 upload and follow the ye olde way.
              return forceDataPrefix(result.generations[0].img, 'image/webp')
            }
            const imgRes = await ctx.http(imgUrl, { responseType: 'arraybuffer' })
            const b64 = Buffer.from(imgRes.data).toString('base64')
            return forceDataPrefix(b64, imgRes.headers.get('content-type'))
          }
          if (config.type === 'comfyui') {
            // get filenames from history
            const promptId = res.data.prompt_id
            const check = () => ctx.http.get(trimSlash(config.endpoint) + '/history/' + promptId)
              .then((res) => res[promptId] && res[promptId].outputs)
            const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))
            let outputs
            while (!(outputs = await check())) {
              await sleep(config.pollInterval)
            }
            // get images by filename
            const imagesOutput: { data: ArrayBuffer, mime: string }[] = [];
            for (const nodeId in outputs) {
              const nodeOutput = outputs[nodeId]
              if ('images' in nodeOutput) {
                for (const image of nodeOutput['images']) {
                  const urlValues = new URLSearchParams({ filename: image['filename'], subfolder: image['subfolder'], type: image['type'] }).toString()
                  const imgRes = await ctx.http(trimSlash(config.endpoint) + '/view?' + urlValues)
                  imagesOutput.push({ data: imgRes.data, mime: imgRes.headers.get('content-type') })
                  break
                }
              }
            }
            // return first image
            return forceDataPrefix(Buffer.from(imagesOutput[0].data).toString('base64'), imagesOutput[0].mime)
          }
          // event: newImage
          // id: 1
          // data:
          //                                                                        ↓ nai-v3
          if (res.headers.get('content-type') === 'application/x-zip-compressed' || res.headers.get('content-disposition')?.includes('.zip')) {
            const buffer = Buffer.from(res.data, 'binary')  // Ensure 'binary' encoding
            const zip = new AdmZip(buffer)

            // Gets all files in the ZIP file
            const zipEntries = zip.getEntries()
            const firstImageBuffer = zip.readFile(zipEntries[0])
            const b64 = Buffer.from(firstImageBuffer).toString('base64')
            return forceDataPrefix(b64, 'image/png')
          }
          return forceDataPrefix(res.data?.trimEnd().slice(27))
        }

        let dataUrl: string, count = 0