feat: add ws demo

.env.example

@@ -3,3 +3,5 @@ COMFY_API_TOKEN=""
 COMFY_DEPLOYMENT_ID=""
 COMFY_DEPLOYMENT_ID_IMG_2_IMG=""
 COMFY_DEPLOYMENT_ID_CONTROLNET=""
+
+COMFY_DEPLOYMENT_WS=""

package.json

@@ -17,14 +17,17 @@
     "@radix-ui/react-tabs": "^1.0.4",
     "class-variance-authority": "^0.7.0",
     "clsx": "^2.1.0",
-    "comfydeploy": "^0.0.11",
+    "comfydeploy": "0.0.12",
     "lucide-react": "^0.309.0",
     "next": "14.0.3",
     "react": "^18",
     "react-dom": "^18",
     "react-hook-form": "^7.49.3",
+    "react-use-websocket": "^4.7.0",
+    "swr": "^2.2.5",
     "tailwind-merge": "^2.2.0",
     "tailwindcss-animate": "^1.0.7",
+    "use-debounce": "^10.0.0",
     "zod": "^3.22.4"
   },
   "devDependencies": {

src/app/page.tsx

@@ -26,16 +26,21 @@ import {
 
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { ImageGenerationResult } from "@/components/ImageGenerationResult";
+import { WebsocketDemo } from "@/components/WebsocketDemo";
 
 export default function Page() {
   return (
     <main className="flex min-h-screen flex-col items-center justify-between mt-10">
-      <Tabs defaultValue="txt2img" className="w-full max-w-[600px]">
-        <TabsList className="grid w-full grid-cols-3">
+      <Tabs defaultValue="ws" className="w-full max-w-[600px]">
+        <TabsList className="grid w-full grid-cols-4">
+          <TabsTrigger value="ws">Realtime</TabsTrigger>
           <TabsTrigger value="txt2img">txt2img</TabsTrigger>
           <TabsTrigger value="img2img">img2img</TabsTrigger>
           <TabsTrigger value="controlpose">Controlpose</TabsTrigger>
         </TabsList>
+        <TabsContent value="ws">
+          <WebsocketDemo />
+        </TabsContent>
         <TabsContent value="txt2img">
           <Txt2img />
         </TabsContent>

src/components/WebsocketDemo.tsx

@@ -0,0 +1,173 @@
+"use client"
+
+import { getWebsocketUrl } from '@/server/generate'
+import { useCallback, useEffect, useRef, useState } from 'react'
+import useSWR from 'swr'
+import useWebSocket, { ReadyState } from 'react-use-websocket';
+import { useDebounce } from "use-debounce";
+import { Input } from './ui/input';
+import { Badge } from './ui/badge';
+
+export function WebsocketDemo() {
+    const { data } = useSWR("ws", getWebsocketUrl, {
+        revalidateOnFocus: false,
+    })
+    const [ws, setWs] = useState<WebSocket>()
+
+    const [status, setStatus] = useState("not-connected")
+    const [prompt, setPrompt] = useState('A anime cat');
+    const [debouncedPrompt] = useDebounce(prompt, 200);
+
+    const [currentLog, setCurrentLog] = useState<string>();
+
+    const [reconnectCounter, setReconnectCounter] = useState(0)
+
+    const canvasRef = useRef<HTMLCanvasElement>(null); // Reference to the canvas element
+
+    const sendInput = useCallback(() => {
+        if (status == "reconnecting")
+            return
+
+        if (ws?.readyState == ws?.CLOSED) {
+            setStatus('reconnecting')
+            setReconnectCounter(x => x + 1)
+            return
+        }
+
+        if (status != "ready")
+            return
+
+        ws?.send(JSON.stringify(
+            {
+                "event": "input",
+                "inputs": {
+                    "input_text": debouncedPrompt
+                }
+            }
+        ))
+    }, [ws, debouncedPrompt, status])
+
+    const preStatus = useRef(status)
+
+    useEffect(() => {
+        if (preStatus.current != status && status == "ready")
+            sendInput();
+        preStatus.current = status
+    }, [status, sendInput])
+
+    useEffect(() => {
+        sendInput();
+    }, [debouncedPrompt])
+
+    const connectWS = useCallback((data: NonNullable<Awaited<ReturnType<typeof getWebsocketUrl>>>) => {
+        const websocket = new WebSocket(data.ws_connection_url);
+        websocket.binaryType = "arraybuffer";
+        websocket.onopen = () => {
+            setStatus("connected");
+        };
+        websocket.onmessage = (event) => {
+            if (typeof event.data === "string") {
+                const message = JSON.parse(event.data);
+                if (message?.event == "status" && message?.data?.sid) {
+                    setStatus("ready");
+                }
+                if (message?.event) {
+                    if (message?.event == "executing" && message?.data?.node == null)
+                        setCurrentLog("done")
+                    else if (message?.event == "live_status")
+                    setCurrentLog(`running - ${message.data?.current_node} ${(message.data.progress * 100).toFixed(2)}%`)
+                    else if (message?.event == "elapsed_time")
+                        setCurrentLog(`elapsed time: ${Math.ceil(message.data?.elapsed_time * 100) / 100}s`)
+                }
+                console.log("Received message:", message);
+            }
+            if (event.data instanceof ArrayBuffer) {
+                console.log("Received binary message:");
+                drawImage(event.data);
+            }
+        };
+        websocket.onclose = () => setStatus("closed");
+        websocket.onerror = () => setStatus("error");
+
+        setWs(websocket);
+
+        return () => {
+            websocket.close();
+        };
+    }, [data])
+
+    const drawImage = useCallback((arrayBuffer: ArrayBuffer) => {
+        const view = new DataView(arrayBuffer);
+        const eventType = view.getUint32(0);
+        const buffer = arrayBuffer.slice(4);
+        switch (eventType) {
+            case 1:
+                const view2 = new DataView(arrayBuffer);
+                const imageType = view2.getUint32(0)
+                let imageMime
+                switch (imageType) {
+                    case 1:
+                    default:
+                        imageMime = "image/jpeg";
+                        break;
+                    case 2:
+                        imageMime = "image/png"
+                }
+                const blob = new Blob([buffer.slice(4)], { type: imageMime });
+
+                // const blob = new Blob([arrayBuffer], { type: 'image/png' }); // Assuming the image is a JPEG
+                const url = URL.createObjectURL(blob);
+
+                const canvas = canvasRef.current;
+                const ctx = canvas?.getContext('2d');
+
+                if (ctx) {
+                    console.log("drawing");
+
+                    const img = new Image();
+                    img.onload = () => {
+                        if (canvas) {
+                            ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
+                        }
+                        URL.revokeObjectURL(url); // Clean up
+                    };
+                    img.src = url;
+                }
+                // this.dispatchEvent(new CustomEvent("b_preview", { detail: imageBlob }));
+                break;
+            default:
+                throw new Error(`Unknown binary websocket message of type ${eventType}`);
+        }
+    }, []);
+
+    useEffect(() => {
+        if (!data) {
+            setStatus("not-connected");
+            return;
+        }
+
+        return connectWS(data)
+    }, [connectWS, reconnectCounter])
+
+    return (
+        <div className='flex flex-col gap-2'>
+            <div className='flex gap-2'>
+                <Badge variant={'outline'} className='w-fit'>Status: {status}</Badge>
+                <Badge variant={'outline'} className='w-fit'>
+                    {currentLog}
+                    {status == "connected" &&!currentLog && "stating comfy ui"}
+                    {status == "ready" &&!currentLog && " running"}
+                </Badge>
+            </div>
+            <canvas ref={canvasRef} className='rounded-lg' width="1024" height="1024"></canvas>
+
+
+            <Input
+                id="picture"
+                type="text"
+                value={prompt}
+                onChange={(e) => setPrompt(e.target.value)}
+            />
+        </div>
+    )
+}

src/components/ui/badge.tsx

@@ -0,0 +1,36 @@
+import * as React from "react"
+import { cva, type VariantProps } from "class-variance-authority"
+
+import { cn } from "@/lib/utils"
+
+const badgeVariants = cva(
+  "inline-flex items-center rounded-full border px-2.5 py-0.5 text-xs font-semibold transition-colors focus:outline-none focus:ring-2 focus:ring-ring focus:ring-offset-2",
+  {
+    variants: {
+      variant: {
+        default:
+          "border-transparent bg-primary text-primary-foreground hover:bg-primary/80",
+        secondary:
+          "border-transparent bg-secondary text-secondary-foreground hover:bg-secondary/80",
+        destructive:
+          "border-transparent bg-destructive text-destructive-foreground hover:bg-destructive/80",
+        outline: "text-foreground",
+      },
+    },
+    defaultVariants: {
+      variant: "default",
+    },
+  }
+)
+
+export interface BadgeProps
+  extends React.HTMLAttributes<HTMLDivElement>,
+    VariantProps<typeof badgeVariants> {}
+
+function Badge({ className, variant, ...props }: BadgeProps) {
+  return (
+    <div className={cn(badgeVariants({ variant }), className)} {...props} />
+  )
+}
+
+export { Badge, badgeVariants }

src/server/generate.tsx

@@ -7,7 +7,7 @@ const client = new ComfyDeployClient({
     apiToken: process.env.COMFY_API_TOKEN!,
 })
 
-export async function generate(prompt: string){
+export async function generate(prompt: string) {
     return await client.run({
         deployment_id: process.env.COMFY_DEPLOYMENT_ID!,
         inputs: {
@@ -16,7 +16,7 @@ export async function generate(prompt: string){
     })
 }
 
-export async function generate_img(input_image: string){
+export async function generate_img(input_image: string) {
     return await client.run({
         deployment_id: process.env.COMFY_DEPLOYMENT_ID_IMG_2_IMG!,
         inputs: {
@@ -25,7 +25,7 @@ export async function generate_img(input_image: string){
     })
 }
 
-export async function generate_img_with_controlnet(input_openpose_url: string, prompt: string){
+export async function generate_img_with_controlnet(input_openpose_url: string, prompt: string) {
     return await client.run({
         deployment_id: process.env.COMFY_DEPLOYMENT_ID_CONTROLNET!,
         inputs: {
@@ -35,14 +35,20 @@ export async function generate_img_with_controlnet(input_openpose_url: string, p
     })
 }
 
-export async function checkStatus(run_id: string){
+export async function checkStatus(run_id: string) {
     return await client.getRun(run_id)
 }
 
-export async function getUploadUrl(type: string, file_size: number){
+export async function getUploadUrl(type: string, file_size: number) {
     try {
         return await client.getUploadUrl(type, file_size)
     } catch (error) {
         console.log(error)
     }
 }
+
+export async function getWebsocketUrl() {
+    return await client.getWebsocketUrl({
+        deployment_id: process.env.COMFY_DEPLOYMENT_WS!,
+    })
+}