import { useEffect, useMemo, useRef, useState } from 'react' import type { ComfyApp } from '@comfyorg/comfyui-frontend-types' import { Eye, EyeOff } from 'lucide-react' import { Button } from '@/components/ui/button' import { useSamplingPreviewFrames, useSamplingPreviewPlayback, type SamplingPreviewFrame, } from '@/hooks/use-sampling-preview-progress' import { LocaleContext, useT } from '@/lib/i18n' export const SAMPLING_PREVIEW_EVENT = 'easy_media.sampling_preview' const EXECUTION_START_EVENT = 'execution_start' export interface SamplingPreviewPayload { node_id: string display_node_id?: string prompt_id: string | null image: string images?: string[] mime: string step: number total: number fps: number | null frame_count: number segment_index: number sampling_pass: string } interface SamplingPreviewWidgetProps { app: ComfyApp nodeId: string | number | (() => string | number) executionRevision?: number onVisibilityChange?: (visible: boolean) => void interactive?: boolean visibilityControlOnly?: boolean } function terminalNodeId(value: string | number | undefined) { return value === undefined ? '' : String(value).split(':').at(-1)?.split('.').at(-1) } function eventTargetsNode( payload: SamplingPreviewPayload, nodeId: string | number | (() => string | number), ) { const currentNodeId = typeof nodeId === 'function' ? nodeId() : nodeId return [payload.display_node_id, payload.node_id] .some((candidate) => terminalNodeId(candidate) === String(currentNodeId)) } function previewBlob(image: string, mime: string) { const binary = atob(image) const bytes = new Uint8Array(binary.length) for (let index = 0; index < binary.length; index += 1) { bytes[index] = binary.charCodeAt(index) } return new Blob([bytes], { type: mime || 'image/jpeg' }) } function PreviewCanvas({ frame, label }: { frame: SamplingPreviewFrame, label: string }) { const canvasRef = useRef(null) useEffect(() => { const canvas = canvasRef.current if (!canvas) return canvas.width = frame.bitmap.width canvas.height = frame.bitmap.height const context = canvas.getContext('2d') if (!context) { console.error('[SamplingPreviewWidget] 2D canvas context is unavailable') return } context.drawImage(frame.bitmap, 0, 0) }, [frame]) return ( ) } function SamplingPreviewWidgetInner({ app, nodeId, executionRevision = 0, onVisibilityChange, interactive = true, visibilityControlOnly = false, }: SamplingPreviewWidgetProps) { const t = useT() const [preview, setPreview] = useState(null) const [pausedPreview, setPausedPreview] = useState(null) const [previewUrl, setPreviewUrl] = useState('') const [legacyPreviewUrls, setLegacyPreviewUrls] = useState([]) const [isPlaying, setIsPlaying] = useState(true) const [isPreviewVisible, setIsPreviewVisible] = useState(true) const [frozenFallbackUrl, setFrozenFallbackUrl] = useState('') const fallbackImageRef = useRef(null) const frozenFallbackCanvasRef = useRef(null) const activePromptIdRef = useRef(null) const playbackInteractive = interactive && !visibilityControlOnly const visibilityInteractive = interactive || visibilityControlOnly const displayedPreview = isPlaying ? preview : (pausedPreview ?? preview) useEffect(() => { const handleExecutionStart = (event: CustomEvent) => { const detail = event.detail activePromptIdRef.current = detail && typeof detail === 'object' && 'prompt_id' in detail && typeof detail.prompt_id === 'string' ? detail.prompt_id : null } const handlePreview = (event: CustomEvent) => { const payload = event.detail as SamplingPreviewPayload | undefined if (!payload?.image || !eventTargetsNode(payload, nodeId)) return if ( activePromptIdRef.current !== null && payload.prompt_id !== activePromptIdRef.current ) return setPreview(payload) } app.api.addCustomEventListener(EXECUTION_START_EVENT, handleExecutionStart) app.api.addCustomEventListener(SAMPLING_PREVIEW_EVENT, handlePreview) return () => { app.api.removeCustomEventListener(EXECUTION_START_EVENT, handleExecutionStart) app.api.removeCustomEventListener(SAMPLING_PREVIEW_EVENT, handlePreview) } }, [app.api, nodeId]) useEffect(() => { setPreview(null) setPausedPreview(null) setIsPlaying(true) setFrozenFallbackUrl('') }, [executionRevision]) useEffect(() => { onVisibilityChange?.(preview !== null) }, [onVisibilityChange, preview]) const previewData = useMemo( () => displayedPreview ? previewBlob(displayedPreview.image, displayedPreview.mime) : null, [displayedPreview], ) useEffect(() => { if (!previewData) { setPreviewUrl('') return } const nextUrl = URL.createObjectURL(previewData) setPreviewUrl(nextUrl) return () => URL.revokeObjectURL(nextUrl) }, [previewData]) useEffect(() => { const legacyImages = displayedPreview?.images ?? [] if (legacyImages.length === 0) { setLegacyPreviewUrls([]) return } const nextUrls = legacyImages.map((image) => URL.createObjectURL( previewBlob(image, 'image/jpeg'), )) setLegacyPreviewUrls(nextUrls) return () => nextUrls.forEach((url) => URL.revokeObjectURL(url)) }, [displayedPreview]) const decodedFrames = useSamplingPreviewFrames( previewData, displayedPreview?.image, playbackInteractive, ) const playbackFrameCount = legacyPreviewUrls.length || displayedPreview?.frame_count || decodedFrames.length const playbackRevision = displayedPreview ? `${displayedPreview.segment_index}:${displayedPreview.sampling_pass}:${displayedPreview.frame_count}:${displayedPreview.fps}` : undefined const { frameIndex, progress: playbackProgress } = useSamplingPreviewPlayback( playbackRevision, playbackFrameCount, displayedPreview?.fps ?? null, isPlaying, ) if (!displayedPreview) return null const passLabel = t(`samplingPreview.${displayedPreview.sampling_pass}`) const decodedFrameIndex = decodedFrames.length <= 1 || playbackFrameCount <= 1 ? 0 : Math.min( decodedFrames.length - 1, Math.floor(frameIndex / playbackFrameCount * decodedFrames.length), ) const decodedFrame = decodedFrames[decodedFrameIndex] const legacyPreviewUrl = legacyPreviewUrls[ Math.min(frameIndex, legacyPreviewUrls.length - 1) ] const usesAnimatedImageFallback = ( displayedPreview.mime === 'image/webp' && legacyPreviewUrls.length === 0 && decodedFrame === undefined ) const freezeAnimatedFallback = () => { if (!usesAnimatedImageFallback) return const image = fallbackImageRef.current const canvas = frozenFallbackCanvasRef.current if (!image?.complete || image.naturalWidth <= 0 || image.naturalHeight <= 0 || !canvas) return canvas.width = image.naturalWidth canvas.height = image.naturalHeight const context = canvas.getContext('2d') if (!context) { console.error('[SamplingPreviewWidget] 2D canvas context is unavailable') return } context.drawImage(image, 0, 0) setFrozenFallbackUrl(previewUrl) } const togglePlayback = () => { if (isPlaying) { freezeAnimatedFallback() setPausedPreview(preview) setIsPlaying(false) } else { setPausedPreview(null) setIsPlaying(true) } } const playbackLabel = t(isPlaying ? 'samplingPreview.pausePreview' : 'samplingPreview.resumePreview') const metaLabel = t('samplingPreview.meta', { segment: displayedPreview.segment_index + 1, pass: passLabel, step: displayedPreview.step, total: displayedPreview.total, }) const progressBar = (
) return (
{isPreviewVisible ? ( playbackInteractive ? ( ) : (
{legacyPreviewUrl ? ( {t('samplingPreview.title')} ) : previewUrl ? ( {t('samplingPreview.title')} ) : null}
) ) : ( ) } export function SamplingPreviewWidget(props: SamplingPreviewWidgetProps) { const locale = props.app?.ui?.settings?.settingsValues?.['Comfy.Locale'] return ( ) }