From 39da2a57daad56b82f67e1fbce1459946f6efa8f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=B4ng=20Ti=E1=BA=BFn=20Tr=E1=BA=A7n?=
<tctien342@gmail.com>
Date: Wed, 15 Jan 2025 12:28:04 +0700
Subject: [PATCH] feat(output): initial support for video file
---
app/[locale]/main/WorkflowSidePicker.tsx | 2 +-
app/[locale]/main/workflow/TaskItem.tsx | 2 +-
components/AttachmentDetail.tsx | 2 +-
components/AttachmentReview.tsx | 74 ++++++--
components/ImageGallery.tsx | 4 +-
.../steps/CreateInputNode.tsx | 12 ++
.../steps/CreateOutputNode.tsx | 11 +-
.../AddWorkflowDialog/steps/Finalize.tsx | 5 +-
components/ui-ext/download-button.tsx | 56 ++++--
entities/attachment.ts | 2 +-
entities/enum.ts | 1 +
server/handlers/workflow.ts | 6 +-
server/routers/attachment.ts | 1 +
server/routers/snippet.ts | 4 +-
server/routers/workflow.ts | 5 +-
server/routers/workflow_task.ts | 16 +-
server/schemas/attachment.ts | 2 +-
server/utils/file.ts | 16 ++
services/comfyui.service.ts | 167 ++++++++++++------
utils/workflow.ts | 30 +++-
20 files changed, 316 insertions(+), 102 deletions(-)
create mode 100644 server/utils/file.ts
diff --git a/app/[locale]/main/WorkflowSidePicker.tsx b/app/[locale]/main/WorkflowSidePicker.tsx
index 5c283d7..300cf49 100644
--- a/app/[locale]/main/WorkflowSidePicker.tsx
+++ b/app/[locale]/main/WorkflowSidePicker.tsx
@@ -83,7 +83,7 @@ export const WorkflowSidePicker: IComponent = () => {
if (crrInput.type === EValueType.Number) {
inputRecord[key] = Number(inputData[key] || crrInput.default)
}
- if ([EValueType.File, EValueType.Image].includes(crrInput.type as EValueType)) {
+ if ([EValueType.File, EValueType.Video, EValueType.Image].includes(crrInput.type as EValueType)) {
const files = inputData[key] as File[]
if (!files || files.length === 0) {
toast({
diff --git a/app/[locale]/main/workflow/TaskItem.tsx b/app/[locale]/main/workflow/TaskItem.tsx
index e9a85f1..e49e9b6 100644
--- a/app/[locale]/main/workflow/TaskItem.tsx
+++ b/app/[locale]/main/workflow/TaskItem.tsx
@@ -241,7 +241,7 @@ export const TaskItem: IComponent<{
/>
{runningTime >= 0 && <MiniBadge Icon={Hourglass} title='Take' count={`${runningTime}s`} />}
{task.repeatCount > 1 && <MiniBadge Icon={Repeat} title='Repeat' count={task.repeatCount} />}
- {!!attachments?.length && <MiniBadge Icon={Image} title='Images' count={attachments.length} />}
+ {!!attachments?.length && <MiniBadge Icon={Image} title='Files' count={attachments.length} />}
{task.computedCost > 0 && (
<MiniBadge Icon={DollarSign} title='Credits' count={task.computedCost.toFixed(2)} />
)}
diff --git a/components/AttachmentDetail.tsx b/components/AttachmentDetail.tsx
index fc76ac2..7576b2a 100644
--- a/components/AttachmentDetail.tsx
+++ b/components/AttachmentDetail.tsx
@@ -31,7 +31,7 @@ export const AttachmentDetail: IComponent<{
const renderMapperInput = useCallback((config: IMapperInput, inputVal: any) => {
const value = inputVal || config.default
- if (config.type === EValueType.Image) {
+ if ([EValueType.Image, EValueType.Video].includes(config.type as EValueType)) {
let src = value
if (isArray(src) && src.length === 1) {
src = src[0]
diff --git a/components/AttachmentReview.tsx b/components/AttachmentReview.tsx
index f74de1a..30a849e 100644
--- a/components/AttachmentReview.tsx
+++ b/components/AttachmentReview.tsx
@@ -5,7 +5,7 @@ import { Attachment } from '@/entities/attachment'
import { cn } from '@/lib/utils'
import { PhotoView } from 'react-photo-view'
import { Button } from './ui/button'
-import { Download, MoreHorizontal, Star } from 'lucide-react'
+import { Download, ImageIcon, MoreHorizontal, Star } from 'lucide-react'
import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip'
import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from './ui/dropdown-menu'
import LoadableImage from './LoadableImage'
@@ -24,6 +24,8 @@ import useCurrentMousePosRef from '@/hooks/useCurrentMousePos'
import { useScrollingStatusRef } from '@/hooks/useScrollingStatus'
import useCopyAction from '@/hooks/useCopyAction'
import { useActionDebounce } from '@/hooks/useAction'
+import { EValueType } from '@/entities/enum'
+import { PlayCircleIcon } from '@heroicons/react/24/outline'
const AttachmentTooltipPopup: IComponent<{
taskId?: string
@@ -108,6 +110,8 @@ export const AttachmentReview: IComponent<{
}
)
+ const isVideo = image?.type === EValueType.Video
+
const downloadFn = (mode: 'jpg' | 'raw' = 'raw') => {
if (mode === 'raw') {
window.open(image?.raw?.url, '_blank')
@@ -191,6 +195,20 @@ export const AttachmentReview: IComponent<{
'rounded-lg outline outline-white shadow cursor-pointer': isPop
})}
/>
+ {isPop && (
+ <m.div className='mt-2' onClick={(e) => e.stopPropagation()}>
+ {!!image && (
+ <div className='px-2 py-1 text-xs bg-background/50 border backdrop-blur w-min rounded flex justify-center gap-1'>
+ {image.type === EValueType.Video ? (
+ <PlayCircleIcon className='w-4 h-4' />
+ ) : (
+ <ImageIcon className='w-4 h-4' />
+ )}
+ <code>{image?.type}</code>
+ </div>
+ )}
+ </m.div>
+ )}
</m.div>
{isPop && (
<m.div
@@ -220,21 +238,32 @@ export const AttachmentReview: IComponent<{
<AttachmentTooltipPopup taskId={taskId} active={hoverSync && isHovering} />
</div>
<div className='w-full mt-2 h-10 flex justify-end gap-2'>
- <DropdownMenu>
- <DropdownMenuTrigger asChild className='flex items-center'>
- <Button variant='secondary' className='bg-background/50 backdrop-blur-lg'>
- <code>DOWNLOAD</code> <Download width={16} height={16} className='ml-2' />
- </Button>
- </DropdownMenuTrigger>
- <DropdownMenuContent side='left' className='bg-background/80 backdrop-blur-lg'>
- <DropdownMenuItem onClick={() => downloadFn('jpg')} className='cursor-pointer text-sm'>
- <span>Download compressed JPG</span>
- </DropdownMenuItem>
- <DropdownMenuItem onClick={() => downloadFn()} className='cursor-pointer text-sm'>
- <span>Download Raw</span>
- </DropdownMenuItem>
- </DropdownMenuContent>
- </DropdownMenu>
+ {isVideo ? (
+ <Button
+ variant='secondary'
+ onClick={() => downloadFn()}
+ className='bg-background/50 backdrop-blur-lg'
+ >
+ <code>DOWNLOAD</code> <Download width={16} height={16} className='ml-2' />
+ </Button>
+ ) : (
+ <DropdownMenu>
+ <DropdownMenuTrigger asChild className='flex items-center'>
+ <Button variant='secondary' className='bg-background/50 backdrop-blur-lg'>
+ <code>DOWNLOAD</code> <Download width={16} height={16} className='ml-2' />
+ </Button>
+ </DropdownMenuTrigger>
+ <DropdownMenuContent side='left' className='bg-background/80 backdrop-blur-lg'>
+ <DropdownMenuItem onClick={() => downloadFn('jpg')} className='cursor-pointer text-sm'>
+ <span>Download compressed JPG</span>
+ </DropdownMenuItem>
+
+ <DropdownMenuItem onClick={() => downloadFn()} className='cursor-pointer text-sm'>
+ <span>Download Raw</span>
+ </DropdownMenuItem>
+ </DropdownMenuContent>
+ </DropdownMenu>
+ )}
{!!onPressFavorite && (
<Button
onClick={() => onPressFavorite?.(data?.id!)}
@@ -304,6 +333,19 @@ export const AttachmentReview: IComponent<{
</div>
</div>
)}
+ {!!image && (
+ <div
+ className={cn(
+ 'z-10 group-hover:block absolute bottom-2 left-2 bg-black/50 text-white backdrop-blur rounded p-1'
+ )}
+ >
+ {image.type === EValueType.Video ? (
+ <PlayCircleIcon className='w-4 h-4' />
+ ) : (
+ <ImageIcon className='w-4 h-4' />
+ )}
+ </div>
+ )}
{!!onPressFavorite && (
<div
className={cn('z-10 group-hover:block absolute top-1 left-1', {
diff --git a/components/ImageGallery.tsx b/components/ImageGallery.tsx
index 7785990..dcfd529 100644
--- a/components/ImageGallery.tsx
+++ b/components/ImageGallery.tsx
@@ -102,7 +102,7 @@ export const ImageGallery: IComponent<{
left: `${(virtualRow.lane * 100) / imgPerRow}%`,
width: `calc(100%/${imgPerRow})`,
height: 'fit-content',
- aspectRatio: 'loading' in item ? 1 : item.ratio,
+ aspectRatio: 'loading' in item ? 1 : item.ratio || 1,
transform: `translateY(${virtualRow.start}px)`
}}
>
@@ -129,7 +129,7 @@ export const ImageGallery: IComponent<{
id='bottom'
ref={bottomRef}
className={cn('w-full flex items-center justify-center mt-4 pt-4 pb-24 text-gray-400', {
- 'hidden': items.length === 0
+ hidden: items.length === 0
})}
>
{hasNextPage && <div className='flex'>More data...</div>}
diff --git a/components/dialogs/AddWorkflowDialog/steps/CreateInputNode.tsx b/components/dialogs/AddWorkflowDialog/steps/CreateInputNode.tsx
index 1c031bc..f68c1dd 100644
--- a/components/dialogs/AddWorkflowDialog/steps/CreateInputNode.tsx
+++ b/components/dialogs/AddWorkflowDialog/steps/CreateInputNode.tsx
@@ -24,6 +24,7 @@ import {
LanguageIcon,
ListBulletIcon,
PhotoIcon,
+ PlayCircleIcon,
PuzzlePieceIcon,
SparklesIcon,
VariableIcon
@@ -106,6 +107,9 @@ export const CreateInputNode: IComponent<{
case EValueType.Image:
form.setValue('icon', 'PhotoIcon')
break
+ case EValueType.Video:
+ form.setValue('icon', 'PlayCircleIcon')
+ break
case EValueType.File:
form.setValue('icon', 'DocumentArrowUpIcon')
break
@@ -146,6 +150,8 @@ export const CreateInputNode: IComponent<{
return 'Use for string input like positive, negative, caption,...'
case EValueType.Image:
return 'Use for image input like load_image, load_mask,...'
+ case EValueType.Video:
+ return 'Use for video input like load_video,...'
case EValueType.File:
return 'Use for file input like load_audio, load_text,...'
case EValueType.Boolean:
@@ -288,6 +294,12 @@ export const CreateInputNode: IComponent<{
Image
</div>
</SelectItem>
+ <SelectItem value={EValueType.Video}>
+ <div className='flex items-center'>
+ <PlayCircleIcon className='mr-2 h-4 w-4' />
+ Video
+ </div>
+ </SelectItem>
<SelectItem value={EValueType.File}>
<div className='flex items-center'>
<DocumentArrowUpIcon className='mr-2 h-4 w-4' />
diff --git a/components/dialogs/AddWorkflowDialog/steps/CreateOutputNode.tsx b/components/dialogs/AddWorkflowDialog/steps/CreateOutputNode.tsx
index 05684f0..28b63dc 100644
--- a/components/dialogs/AddWorkflowDialog/steps/CreateOutputNode.tsx
+++ b/components/dialogs/AddWorkflowDialog/steps/CreateOutputNode.tsx
@@ -19,7 +19,7 @@ import {
DocumentArrowUpIcon,
LanguageIcon,
PhotoIcon,
- SparklesIcon,
+ PlayCircleIcon,
VariableIcon
} from '@heroicons/react/24/outline'
import { zodResolver } from '@hookform/resolvers/zod'
@@ -74,6 +74,9 @@ export const CreateOutputNode: IComponent<{
case EValueType.Image:
form.setValue('icon', 'PhotoIcon')
break
+ case EValueType.Video:
+ form.setValue('icon', 'PlayCircleIcon')
+ break
case EValueType.File:
form.setValue('icon', 'DocumentArrowUpIcon')
break
@@ -192,6 +195,12 @@ export const CreateOutputNode: IComponent<{
Image
</div>
</SelectItem>
+ <SelectItem value={EValueType.Video}>
+ <div className='flex items-center'>
+ <PlayCircleIcon className='mr-2 h-4 w-4' />
+ Video
+ </div>
+ </SelectItem>
<SelectItem value={EValueType.File}>
<div className='flex items-center'>
<DocumentArrowUpIcon className='mr-2 h-4 w-4' />
diff --git a/components/dialogs/AddWorkflowDialog/steps/Finalize.tsx b/components/dialogs/AddWorkflowDialog/steps/Finalize.tsx
index fbab2da..1d8b409 100644
--- a/components/dialogs/AddWorkflowDialog/steps/Finalize.tsx
+++ b/components/dialogs/AddWorkflowDialog/steps/Finalize.tsx
@@ -75,7 +75,7 @@ export const FinalizeStep: IComponent = () => {
for (const key of inputKeys) {
const input = workflow?.mapInput?.[key as keyof typeof workflow.mapInput]
if (!input) continue
- if (input.type === EValueType.File || input.type === EValueType.Image) {
+ if (input.type === EValueType.File || input.type === EValueType.Image || input.type === EValueType.Video) {
const files = wfObj[key] as File[]
if (files.length > 0) {
const file = files[0]
@@ -140,7 +140,7 @@ export const FinalizeStep: IComponent = () => {
defaultValue={String(input.default ?? '')}
/>
)}
- {[EValueType.File, EValueType.Image].includes(input.type as EValueType) && (
+ {[EValueType.File, EValueType.Image, EValueType.Video].includes(input.type as EValueType) && (
<DropFileInput
disabled={loading}
defaultFiles={inputWorkflowTest.current[val]}
@@ -240,6 +240,7 @@ export const FinalizeStep: IComponent = () => {
const data = progressEv.data.output[key as keyof typeof progressEv.data]
let items: ReactNode[] = []
switch (data.info.type) {
+ case EValueType.Video:
case EValueType.Image:
const imageURLs = data.data as { url: string }[]
items.push(
diff --git a/components/ui-ext/download-button.tsx b/components/ui-ext/download-button.tsx
index a5a9f15..6dcdba2 100644
--- a/components/ui-ext/download-button.tsx
+++ b/components/ui-ext/download-button.tsx
@@ -7,6 +7,7 @@ import { Button } from '../ui/button'
import { Download } from 'lucide-react'
import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, DropdownMenuTrigger } from '../ui/dropdown-menu'
import { trpc } from '@/utils/trpc'
+import { EValueType } from '@/entities/enum'
const DownloadImagesButton: IComponent<{
workflowTaskId?: string
@@ -19,19 +20,45 @@ const DownloadImagesButton: IComponent<{
refetchOnWindowFocus: false
}
)
+ const haveVideo = !!attachments?.some((a) => a?.type === EValueType.Video)
+ const haveImage = !!attachments?.some((a) => a?.type === EValueType.Image)
+
const { bundleImages, isLoading, progress, error } = useImageBundler()
- const downloadFn = async (mode: 'jpg' | 'raw' = 'jpg') => {
+ const downloadRawOutput = () => {
if (isLoading) return
try {
const images = attachments
?.filter((a) => !!a)
.filter((a) => !!a.raw?.url)
- .map((v) => {
- if (mode === 'jpg') return v.high!.url
- return v.raw!.url
- }) as string[]
- await bundleImages(images)
+ .map((v) => v.raw!.url) as string[]
+ bundleImages(images)
+ } catch (e) {
+ console.error(e)
+ }
+ }
+
+ const downloadCompressedJpg = () => {
+ if (isLoading) return
+ try {
+ const images = attachments
+ ?.filter((a) => !!a && a.type === EValueType.Image)
+ .filter((a) => !!a.high?.url)
+ .map((v) => v.high!.url) as string[]
+ bundleImages(images)
+ } catch (e) {
+ console.error(e)
+ }
+ }
+
+ const downloadVideos = () => {
+ if (isLoading) return
+ try {
+ const videos = attachments
+ ?.filter((a) => !!a && a.type === EValueType.Video)
+ .filter((a) => !!a.raw?.url)
+ .map((v) => v.raw!.url) as string[]
+ bundleImages(videos)
} catch (e) {
console.error(e)
}
@@ -58,11 +85,18 @@ const DownloadImagesButton: IComponent<{
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent side='left' align='center'>
- <DropdownMenuItem onClick={() => downloadFn('jpg')} className='cursor-pointer text-sm'>
- <span>Download compressed JPG</span>
- </DropdownMenuItem>
- <DropdownMenuItem onClick={() => downloadFn()} className='cursor-pointer text-sm'>
- <span>Download Raw</span>
+ {haveVideo && (
+ <DropdownMenuItem onClick={downloadVideos} className='cursor-pointer text-sm'>
+ <span>Download video files</span>
+ </DropdownMenuItem>
+ )}
+ {haveImage && (
+ <DropdownMenuItem onClick={downloadCompressedJpg} className='cursor-pointer text-sm'>
+ <span>Download compressed JPG</span>
+ </DropdownMenuItem>
+ )}
+ <DropdownMenuItem onClick={downloadRawOutput} className='cursor-pointer text-sm'>
+ <span>Download raw output</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
diff --git a/entities/attachment.ts b/entities/attachment.ts
index c7398a6..193c6d7 100644
--- a/entities/attachment.ts
+++ b/entities/attachment.ts
@@ -24,7 +24,7 @@ export class Attachment {
ratio?: number // width / height, only for image
@Property({ type: 'varchar', default: EValueType.Image, nullable: true })
- type?: EValueType.Image | EValueType.File
+ type?: EValueType.Image | EValueType.Video | EValueType.File
@Property({ type: 'varchar', default: EAttachmentStatus.PENDING, index: true })
status!: EAttachmentStatus
diff --git a/entities/enum.ts b/entities/enum.ts
index 2a128f1..c6e6708 100644
--- a/entities/enum.ts
+++ b/entities/enum.ts
@@ -51,6 +51,7 @@ export enum EValueType {
File = 'File',
String = 'String',
Number = 'Number',
+ Video = 'Video',
Image = 'Image',
Boolean = 'Boolean'
}
diff --git a/server/handlers/workflow.ts b/server/handlers/workflow.ts
index 0f67b33..58c4dbd 100644
--- a/server/handlers/workflow.ts
+++ b/server/handlers/workflow.ts
@@ -126,7 +126,11 @@ export const WorkflowPlugin = new Elysia({ prefix: '/workflow', detail: { tags:
throw new Error(`Value of <${key}> is greater than max value (max ${keyConfig.max})`)
}
}
- if (keyConfig.type === EValueType.File || keyConfig.type === EValueType.Image) {
+ if (
+ keyConfig.type === EValueType.File ||
+ keyConfig.type === EValueType.Image ||
+ keyConfig.type === EValueType.Video
+ ) {
const temp = input[key]
if (Array.isArray(temp) && temp.length === 0) {
set.status = 400
diff --git a/server/routers/attachment.ts b/server/routers/attachment.ts
index 3dd9099..2521dd1 100644
--- a/server/routers/attachment.ts
+++ b/server/routers/attachment.ts
@@ -16,6 +16,7 @@ const getAttachmentURL = async (attachment: Attachment, baseUrl = 'http://localh
AttachmentService.getInstance().getFileURL(highName, 3600 * 24, baseUrl)
])
return {
+ type: attachment.type,
raw: imageInfo,
preview: imagePreviewInfo || imageInfo,
high: imageHighInfo || imageInfo
diff --git a/server/routers/snippet.ts b/server/routers/snippet.ts
index b94d47e..e68516a 100644
--- a/server/routers/snippet.ts
+++ b/server/routers/snippet.ts
@@ -22,7 +22,9 @@ export const snippetRouter = router({
const inputBody: Record<string, string | number | boolean> = {}
for (const inputKey in workflow.mapInput) {
if (workflow.mapInput[inputKey].type === EValueUtilityType.Prefixer) continue
- if ([EValueType.File, EValueType.Image].includes(workflow.mapInput[inputKey].type as EValueType)) {
+ if (
+ [EValueType.File, EValueType.Image, EValueType.Video].includes(workflow.mapInput[inputKey].type as EValueType)
+ ) {
inputBody[inputKey] = '<attachment_id>'
needUpload = true
continue
diff --git a/server/routers/workflow.ts b/server/routers/workflow.ts
index fb76a48..309760c 100644
--- a/server/routers/workflow.ts
+++ b/server/routers/workflow.ts
@@ -257,15 +257,16 @@ export const workflowRouter = router({
builder.input(key, String(inputData))
break
case EValueType.File:
+ case EValueType.Video:
case EValueType.Image:
const file = inputData as Attachment
const fileBlob = await AttachmentService.getInstance().getFileBlob(file.fileName)
if (!fileBlob) {
- return subscriber.next({ key: 'failed', detail: 'missing image' })
+ return subscriber.next({ key: 'failed', detail: 'missing file' })
}
const uploadedImg = await api.uploadImage(fileBlob, file.fileName)
if (!uploadedImg) {
- subscriber.next({ key: 'failed', detail: 'failed to upload image' })
+ subscriber.next({ key: 'failed', detail: 'failed to upload file' })
return
}
builder.input(key, uploadedImg.info.filename)
diff --git a/server/routers/workflow_task.ts b/server/routers/workflow_task.ts
index de11725..27a55b5 100644
--- a/server/routers/workflow_task.ts
+++ b/server/routers/workflow_task.ts
@@ -148,13 +148,13 @@ export const workflowTaskRouter = router({
}),
getOutputAttachmentUrls: privateProcedure.input(z.string()).query(async ({ input, ctx }) => {
const task = await ctx.em.findOneOrFail(WorkflowTask, { id: input }, { populate: ['subTasks.events'] })
- let fileNames: string[] = []
+ let fileNames: { filename: string; type?: EValueType }[] = []
if (task.status !== ETaskStatus.Parent) {
const attachments = await ctx.em.find(Attachment, {
task
})
- fileNames = attachments.map((a) => a.fileName)
+ fileNames = attachments.map((a) => ({ filename: a.fileName, type: a.type }))
} else {
const subTaskIds = task.subTasks.map((t) => t.id)
const attachments = await ctx.em.find(Attachment, {
@@ -164,18 +164,20 @@ export const workflowTaskRouter = router({
}
}
})
- fileNames = attachments.map((a) => a.fileName)
+ fileNames = attachments.map((a) => ({ filename: a.fileName, type: a.type }))
}
return Promise.all(
- fileNames.map(async (fileName) => {
- const prevName = `${fileName}_preview.jpg`
- const highName = `${fileName}_high.jpg`
+ fileNames.map(async (info) => {
+ const { filename, type } = info
+ const prevName = `${filename}_preview.jpg`
+ const highName = `${filename}_high.jpg`
const [imageInfo, imagePreviewInfo, imageHighInfo] = await Promise.all([
- AttachmentService.getInstance().getFileURL(fileName, 3600 * 24, ctx.baseUrl),
+ AttachmentService.getInstance().getFileURL(filename, 3600 * 24, ctx.baseUrl),
AttachmentService.getInstance().getFileURL(prevName, 3600 * 24, ctx.baseUrl),
AttachmentService.getInstance().getFileURL(highName, 3600 * 24, ctx.baseUrl)
])
return {
+ type,
raw: imageInfo,
preview: imagePreviewInfo || imageInfo,
high: imageHighInfo || imageInfo
diff --git a/server/schemas/attachment.ts b/server/schemas/attachment.ts
index 9352c04..41b6a17 100644
--- a/server/schemas/attachment.ts
+++ b/server/schemas/attachment.ts
@@ -7,7 +7,7 @@ export const AttachmentSchema = t.Object({
fileName: t.String(),
size: t.Number(),
ratio: t.Number(),
- type: t.UnionEnum([EValueType.File, EValueType.Image]),
+ type: t.UnionEnum([EValueType.File, EValueType.Image, EValueType.Video]),
status: t.Enum(EAttachmentStatus),
storageType: t.String(),
taskEvent: t.Optional(t.Null()),
diff --git a/server/utils/file.ts b/server/utils/file.ts
new file mode 100644
index 0000000..d314f75
--- /dev/null
+++ b/server/utils/file.ts
@@ -0,0 +1,16 @@
+/**
+ * Determine the type of a Blob: Image, Video, or Other.
+ * @param blob - The Blob to classify.
+ * @returns A string indicating the type: 'image', 'video', or 'other'.
+ */
+export function classifyBlob(blob: Blob): 'image' | 'video' | 'other' {
+ const mimeType = blob.type // Get the MIME type from the Blob
+
+ if (mimeType.startsWith('image/')) {
+ return 'image' // It's an image
+ } else if (mimeType.startsWith('video/')) {
+ return 'video' // It's a video
+ } else {
+ return 'other' // Something else
+ }
+}
diff --git a/services/comfyui.service.ts b/services/comfyui.service.ts
index 55d2cf4..66a0364 100644
--- a/services/comfyui.service.ts
+++ b/services/comfyui.service.ts
@@ -27,6 +27,9 @@ import AttachmentService, { EAttachmentType } from './attachment.service'
import { ImageUtil } from '@/server/utils/ImageUtil'
import { delay } from '@/utils/tools'
import { User } from '@/entities/user'
+import { classifyBlob } from '@/server/utils/file'
+import { Workflow } from '@/entities/workflow'
+import mine from 'mime'
const MONITOR_INTERVAL = 5000
@@ -144,6 +147,104 @@ export class ComfyPoolInstance {
await em.flush()
}
+ private handleImageOutput = async (
+ imgBlob: Blob,
+ info: {
+ key: string
+ idx: number
+ task: WorkflowTask
+ workflow: Workflow
+ },
+ attachment: AttachmentService,
+ em: Awaited<ReturnType<Awaited<ReturnType<typeof MikroORMInstance.getInstance>['getEM']>>>
+ ) => {
+ const { key, idx, task, workflow } = info
+ const imgUtil = new ImageUtil(Buffer.from(await imgBlob.arrayBuffer()))
+ const [preview, high, raw] = await Promise.all([
+ // For thumbnail
+ imgUtil
+ .clone()
+ .resizeMax(1024)
+ .intoPreviewJPG()
+ .catch((e) => {
+ this.logger.w('Error while converting to preview', e)
+ return null
+ }),
+ // For on click into thumbnail preview
+ imgUtil
+ .clone()
+ .intoPreviewJPG()
+ .catch((e) => {
+ this.logger.w('Error while converting to preview', e)
+ return null
+ }),
+ // Raw image, use for download
+ imgUtil.intoPNG()
+ ])
+ const tmpName = `${task.id}_${key}_${idx}.png`
+ const [uploaded] = await Promise.all([
+ attachment.uploadFile(raw, `${tmpName}`),
+ preview ? attachment.uploadFile(preview, `${tmpName}_preview.jpg`) : Promise.resolve(false),
+ high ? attachment.uploadFile(high, `${tmpName}_high.jpg`) : Promise.resolve(false)
+ ])
+ if (uploaded) {
+ const fileInfo = await attachment.getFileURL(tmpName)
+ const ratio = await imgUtil.getRatio()
+ const outputAttachment = em.create(
+ Attachment,
+ {
+ fileName: tmpName,
+ size: raw.byteLength,
+ storageType: fileInfo?.type === EAttachmentType.LOCAL ? EStorageType.LOCAL : EStorageType.S3,
+ status: EAttachmentStatus.UPLOADED,
+ ratio,
+ task,
+ workflow
+ },
+ { partial: true }
+ )
+ em.persist(outputAttachment)
+ return outputAttachment.id
+ }
+ }
+
+ private handleVideoOutput = async (
+ videoBlob: Blob,
+ info: {
+ key: string
+ idx: number
+ task: WorkflowTask
+ workflow: Workflow
+ },
+ attachment: AttachmentService,
+ em: Awaited<ReturnType<Awaited<ReturnType<typeof MikroORMInstance.getInstance>['getEM']>>>
+ ) => {
+ const { key, idx, task, workflow } = info
+ const buff = Buffer.from(await videoBlob.arrayBuffer())
+ const extension = mine.getExtension(videoBlob.type) || 'mp4'
+ const tmpName = `${task.id}_${key}_${idx}.${extension}`
+
+ const uploaded = await attachment.uploadFile(buff, `${tmpName}`)
+ if (uploaded) {
+ const fileInfo = await attachment.getFileURL(tmpName)
+ const outputAttachment = em.create(
+ Attachment,
+ {
+ fileName: tmpName,
+ size: buff.byteLength,
+ type: EValueType.Video,
+ storageType: fileInfo?.type === EAttachmentType.LOCAL ? EStorageType.LOCAL : EStorageType.S3,
+ status: EAttachmentStatus.UPLOADED,
+ task,
+ workflow
+ },
+ { partial: true }
+ )
+ em.persist(outputAttachment)
+ return outputAttachment.id
+ }
+ }
+
private async pickingJob() {
const pool = this.pool
const em = await MikroORMInstance.getInstance().getEM()
@@ -215,6 +316,7 @@ export class ComfyPoolInstance {
builder.input(key, String(inputData))
break
case EValueType.File:
+ case EValueType.Video:
case EValueType.Image:
const attachmentId = inputData as string
const file = await em.findOneOrFail(Attachment, { id: attachmentId })
@@ -243,7 +345,7 @@ export class ComfyPoolInstance {
break
}
}
-
+ console.log(JSON.stringify(builder.workflow))
return new CallWrapper(api, builder)
.onPending(async () => {
await this.updateTaskEventFn(task, ETaskStatus.Running, {
@@ -289,56 +391,21 @@ export class ComfyPoolInstance {
if (Array.isArray(tmpOutput[key])) {
tmpOutput[key] = (await Promise.all(
tmpOutput[key].map(async (v, idx) => {
+ console.log('out', key, v)
if (v instanceof Blob) {
- const imgUtil = new ImageUtil(Buffer.from(await v.arrayBuffer()))
- const [preview, high, raw] = await Promise.all([
- // For thumbnail
- imgUtil
- .clone()
- .resizeMax(1024)
- .intoPreviewJPG()
- .catch((e) => {
- this.logger.w('Error while converting to preview', e)
- return null
- }),
- // For on click into thumbnail preview
- imgUtil
- .clone()
- .intoPreviewJPG()
- .catch((e) => {
- this.logger.w('Error while converting to preview', e)
- return null
- }),
- // Raw image, use for download
- imgUtil.intoPNG()
- ])
- const tmpName = `${task.id}_${key}_${idx}.png`
- const [uploaded] = await Promise.all([
- attachment.uploadFile(raw, `${tmpName}`),
- preview
- ? attachment.uploadFile(preview, `${tmpName}_preview.jpg`)
- : Promise.resolve(false),
- high ? attachment.uploadFile(high, `${tmpName}_high.jpg`) : Promise.resolve(false)
- ])
- if (uploaded) {
- const fileInfo = await attachment.getFileURL(tmpName)
- const ratio = await imgUtil.getRatio()
- const outputAttachment = em.create(
- Attachment,
- {
- fileName: tmpName,
- size: raw.byteLength,
- storageType:
- fileInfo?.type === EAttachmentType.LOCAL ? EStorageType.LOCAL : EStorageType.S3,
- status: EAttachmentStatus.UPLOADED,
- ratio,
- task,
- workflow
- },
- { partial: true }
- )
- em.persist(outputAttachment)
- return outputAttachment.id
+ // Check if v is Video, Image or others
+ const blobType = classifyBlob(v)
+ switch (blobType) {
+ case 'image': {
+ await this.handleImageOutput(v, { key, idx, task, workflow }, attachment, em)
+ break
+ }
+ case 'video': {
+ await this.handleVideoOutput(v, { key, idx, task, workflow }, attachment, em)
+ break
+ }
+ default: {
+ }
}
}
return v
diff --git a/utils/workflow.ts b/utils/workflow.ts
index a1cff9c..92b9980 100644
--- a/utils/workflow.ts
+++ b/utils/workflow.ts
@@ -92,10 +92,32 @@ export const parseOutput = async (api: ComfyApi, workflow: Workflow, data: any)
}
break
case EValueType.Image:
- if (!tmp && 'images' in output) {
- const { images } = output
- tmp = await Promise.all(images.map((img: any) => api.getImage(img)))
- break
+ if (tmp) {
+ if (Array.isArray(tmp)) {
+ tmp = await Promise.all(tmp.map((img: any) => api.getImage(img)))
+ } else {
+ tmp = await api.getImage(tmp)
+ }
+ } else {
+ if ('images' in output) {
+ const { images } = output
+ tmp = await Promise.all(images.map((img: any) => api.getImage(img)))
+ break
+ }
+ }
+ case EValueType.Video:
+ if (tmp) {
+ if (Array.isArray(tmp)) {
+ tmp = await Promise.all(tmp.map((img: any) => api.getImage(img)))
+ } else {
+ tmp = await api.getImage(tmp)
+ }
+ } else {
+ if ('gifs' in output) {
+ const { gifs } = output
+ tmp = await Promise.all(gifs.map((img: any) => api.getImage(img)))
+ break
+ }
}
case EValueType.File:
if (tmp) {