Browse Source

优化画笔

feature/handleReconstruction
yun 2 days ago
parent
commit
ece94a5896
  1. 283
      src/components/media/MediaInpaintingOverlay.tsx
  2. 6
      src/components/media/MediaInpaintingToolbar.tsx
  3. 18
      src/components/media/image-edit/InpaintingEditSession.tsx
  4. 6
      src/components/nodes/GenerateImageNode.tsx
  5. 6
      src/components/nodes/ImageInputNode.tsx

283
src/components/media/MediaInpaintingOverlay.tsx

@ -1,6 +1,6 @@
"use client"; "use client";
import { useCallback, useEffect, useRef, useState } from "react"; import { forwardRef, useCallback, useEffect, useImperativeHandle, useRef, useState } from "react";
import type { RefObject } from "react"; import type { RefObject } from "react";
export interface MediaInpaintingFrameRect { export interface MediaInpaintingFrameRect {
@ -16,6 +16,10 @@ export interface MediaInpaintingMaskPayload {
height: number; height: number;
} }
export interface MediaInpaintingOverlayHandle {
exportMask: () => Promise<MediaInpaintingMaskPayload | null>;
}
export type MediaInpaintingDrawTool = "brush" | "rectangle"; export type MediaInpaintingDrawTool = "brush" | "rectangle";
interface DisplayRect { interface DisplayRect {
@ -36,12 +40,16 @@ interface MediaInpaintingOverlayProps {
clearSignal?: number; clearSignal?: number;
undoSignal?: number; undoSignal?: number;
redoSignal?: number; redoSignal?: number;
onMaskChange?: (payload: MediaInpaintingMaskPayload | null) => void;
onFrameChange?: (rect: MediaInpaintingFrameRect | null) => void; onFrameChange?: (rect: MediaInpaintingFrameRect | null) => void;
} }
const MAX_HISTORY = 30; const MAX_HISTORY = 30;
interface MaskHistoryEntry {
visible: ImageData;
mask: ImageData;
}
function clamp(value: number, min: number, max: number) { function clamp(value: number, min: number, max: number) {
return Math.min(max, Math.max(min, value)); return Math.min(max, Math.max(min, value));
} }
@ -69,16 +77,6 @@ async function canvasToPngBlob(canvas: HTMLCanvasElement): Promise<Blob> {
}); });
} }
function canvasHasVisiblePixels(canvas: HTMLCanvasElement): boolean {
const context = canvas.getContext("2d");
if (!context) return false;
const pixels = context.getImageData(0, 0, canvas.width, canvas.height).data;
for (let index = 3; index < pixels.length; index += 4) {
if (pixels[index] > 0) return true;
}
return false;
}
function createCheckerPattern(context: CanvasRenderingContext2D) { function createCheckerPattern(context: CanvasRenderingContext2D) {
const size = 40; const size = 40;
const tile = document.createElement("canvas"); const tile = document.createElement("canvas");
@ -95,7 +93,7 @@ function createCheckerPattern(context: CanvasRenderingContext2D) {
return context.createPattern(tile, "repeat"); return context.createPattern(tile, "repeat");
} }
export function MediaInpaintingOverlay({ export const MediaInpaintingOverlay = forwardRef<MediaInpaintingOverlayHandle, MediaInpaintingOverlayProps>(function MediaInpaintingOverlay({
mediaElementRef, mediaElementRef,
naturalWidth, naturalWidth,
naturalHeight, naturalHeight,
@ -106,26 +104,48 @@ export function MediaInpaintingOverlay({
clearSignal = 0, clearSignal = 0,
undoSignal = 0, undoSignal = 0,
redoSignal = 0, redoSignal = 0,
onMaskChange,
onFrameChange, onFrameChange,
}: MediaInpaintingOverlayProps) { }, ref) {
const overlayRef = useRef<HTMLDivElement>(null); const overlayRef = useRef<HTMLDivElement>(null);
const canvasRef = useRef<HTMLCanvasElement>(null); const canvasRef = useRef<HTMLCanvasElement>(null);
const previewCanvasRef = useRef<HTMLCanvasElement>(null); const maskCanvasRef = useRef<HTMLCanvasElement | null>(null);
const cursorRef = useRef<HTMLDivElement>(null);
const drawingRef = useRef(false); const drawingRef = useRef(false);
const lastPointRef = useRef<{ x: number; y: number } | null>(null); const lastPointRef = useRef<{ x: number; y: number } | null>(null);
const rectangleStartRef = useRef<{ x: number; y: number } | null>(null); const rectangleStartRef = useRef<{ x: number; y: number } | null>(null);
const historyRef = useRef<ImageData[]>([]); const historyRef = useRef<MaskHistoryEntry[]>([]);
const redoRef = useRef<ImageData[]>([]); const redoRef = useRef<MaskHistoryEntry[]>([]);
const checkerPatternRef = useRef<CanvasPattern | null>(null);
const [bounds, setBounds] = useState<DisplayRect | null>(null); const [bounds, setBounds] = useState<DisplayRect | null>(null);
const [cursorPoint, setCursorPoint] = useState<{ x: number; y: number } | null>(null);
const [rectanglePreview, setRectanglePreview] = useState<{ start: { x: number; y: number }; end: { x: number; y: number } } | null>(null); const [rectanglePreview, setRectanglePreview] = useState<{ start: { x: number; y: number }; end: { x: number; y: number } } | null>(null);
const [naturalDimensions, setNaturalDimensions] = useState<{ width: number; height: number } | null>(() => const [naturalDimensions, setNaturalDimensions] = useState<{ width: number; height: number } | null>(() =>
validDimensions(naturalWidth, naturalHeight) validDimensions(naturalWidth, naturalHeight)
); );
const [hasMask, setHasMask] = useState(false);
const activeDrawTool = eraseMode ? "brush" : drawTool; const activeDrawTool = eraseMode ? "brush" : drawTool;
const resetCanvases = useCallback(() => {
const canvas = canvasRef.current;
if (!canvas || !naturalDimensions) return;
canvas.width = naturalDimensions.width;
canvas.height = naturalDimensions.height;
canvas.getContext("2d")?.clearRect(0, 0, canvas.width, canvas.height);
const maskCanvas = maskCanvasRef.current ?? document.createElement("canvas");
maskCanvas.width = naturalDimensions.width;
maskCanvas.height = naturalDimensions.height;
maskCanvas.getContext("2d")?.clearRect(0, 0, maskCanvas.width, maskCanvas.height);
maskCanvasRef.current = maskCanvas;
historyRef.current = [];
redoRef.current = [];
checkerPatternRef.current = null;
}, [naturalDimensions]);
const setCanvasElement = useCallback((canvas: HTMLCanvasElement | null) => {
canvasRef.current = canvas;
if (canvas) resetCanvases();
}, [resetCanvases]);
useEffect(() => { useEffect(() => {
const explicitDimensions = validDimensions(naturalWidth, naturalHeight); const explicitDimensions = validDimensions(naturalWidth, naturalHeight);
if (explicitDimensions) { if (explicitDimensions) {
@ -184,87 +204,45 @@ export function MediaInpaintingOverlay({
}, [measureBounds, onFrameChange]); }, [measureBounds, onFrameChange]);
useEffect(() => { useEffect(() => {
const canvas = canvasRef.current; resetCanvases();
if (!canvas || !naturalDimensions) return; }, [resetCanvases]);
canvas.width = naturalDimensions.width;
canvas.height = naturalDimensions.height;
const context = canvas.getContext("2d");
context?.clearRect(0, 0, canvas.width, canvas.height);
const previewCanvas = previewCanvasRef.current;
if (previewCanvas) {
previewCanvas.width = naturalDimensions.width;
previewCanvas.height = naturalDimensions.height;
previewCanvas.getContext("2d")?.clearRect(0, 0, previewCanvas.width, previewCanvas.height);
}
historyRef.current = [];
redoRef.current = [];
setHasMask(false);
onMaskChange?.(null);
}, [naturalDimensions, onMaskChange]);
const renderPreview = useCallback(() => {
const maskCanvas = canvasRef.current;
const previewCanvas = previewCanvasRef.current;
const context = previewCanvas?.getContext("2d");
if (!maskCanvas || !previewCanvas || !context) return;
if (previewCanvas.width !== maskCanvas.width || previewCanvas.height !== maskCanvas.height) {
previewCanvas.width = maskCanvas.width;
previewCanvas.height = maskCanvas.height;
}
context.clearRect(0, 0, previewCanvas.width, previewCanvas.height);
if (!canvasHasVisiblePixels(maskCanvas)) return;
const pattern = createCheckerPattern(context);
context.save();
context.fillStyle = pattern ?? "rgba(230,233,238,0.9)";
context.fillRect(0, 0, previewCanvas.width, previewCanvas.height);
context.globalCompositeOperation = "destination-in";
context.drawImage(maskCanvas, 0, 0);
context.globalCompositeOperation = "source-atop";
context.fillStyle = "rgba(255,255,255,0.28)";
context.fillRect(0, 0, previewCanvas.width, previewCanvas.height);
context.restore();
}, []);
const exportMask = useCallback(async () => { const exportMask = useCallback(async () => {
const sourceCanvas = canvasRef.current; const sourceCanvas = maskCanvasRef.current;
if (!sourceCanvas || !naturalDimensions || !canvasHasVisiblePixels(sourceCanvas)) { if (!sourceCanvas || !naturalDimensions) return null;
setHasMask(false);
onMaskChange?.(null);
return;
}
const exportCanvas = document.createElement("canvas"); const exportCanvas = document.createElement("canvas");
exportCanvas.width = naturalDimensions.width; exportCanvas.width = naturalDimensions.width;
exportCanvas.height = naturalDimensions.height; exportCanvas.height = naturalDimensions.height;
const context = exportCanvas.getContext("2d"); const context = exportCanvas.getContext("2d");
if (!context) { if (!context) return null;
onMaskChange?.(null);
return;
}
context.fillStyle = "#000"; context.fillStyle = "#000";
context.fillRect(0, 0, exportCanvas.width, exportCanvas.height); context.fillRect(0, 0, exportCanvas.width, exportCanvas.height);
context.drawImage(sourceCanvas, 0, 0); context.drawImage(sourceCanvas, 0, 0);
const blob = await canvasToPngBlob(exportCanvas); const blob = await canvasToPngBlob(exportCanvas);
setHasMask(true); return {
onMaskChange?.({
blob, blob,
width: exportCanvas.width, width: exportCanvas.width,
height: exportCanvas.height, height: exportCanvas.height,
}); };
}, [naturalDimensions, onMaskChange]); }, [naturalDimensions]);
useEffect(() => { useImperativeHandle(ref, () => ({ exportMask }), [exportMask]);
void exportMask();
}, [exportMask]);
const pushHistory = useCallback(() => { const pushHistory = useCallback(() => {
const canvas = canvasRef.current; const canvas = canvasRef.current;
const maskCanvas = maskCanvasRef.current;
const context = canvas?.getContext("2d"); const context = canvas?.getContext("2d");
if (!canvas || !context) return; const maskContext = maskCanvas?.getContext("2d");
historyRef.current = [...historyRef.current.slice(-(MAX_HISTORY - 1)), context.getImageData(0, 0, canvas.width, canvas.height)]; if (!canvas || !maskCanvas || !context || !maskContext) return;
historyRef.current = [
...historyRef.current.slice(-(MAX_HISTORY - 1)),
{
visible: context.getImageData(0, 0, canvas.width, canvas.height),
mask: maskContext.getImageData(0, 0, maskCanvas.width, maskCanvas.height),
},
];
redoRef.current = []; redoRef.current = [];
}, []); }, []);
@ -295,22 +273,48 @@ export function MediaInpaintingOverlay({
}; };
}, [bounds]); }, [bounds]);
const withBrushComposite = useCallback((context: CanvasRenderingContext2D, draw: () => void) => { const updateCursor = useCallback((point: { x: number; y: number } | null) => {
const cursor = cursorRef.current;
if (!cursor) return;
if (!point) {
cursor.style.display = "none";
return;
}
cursor.style.display = "block";
cursor.style.transform = `translate(${point.x}px, ${point.y}px) translate(-50%, -50%)`;
}, []);
const withBrushComposite = useCallback((draw: (context: CanvasRenderingContext2D) => void) => {
const canvas = canvasRef.current;
const maskCanvas = maskCanvasRef.current;
const context = canvas?.getContext("2d");
const maskContext = maskCanvas?.getContext("2d");
if (!canvas || !maskCanvas || !context || !maskContext) return;
context.save(); context.save();
context.globalCompositeOperation = eraseMode ? "destination-out" : "source-over"; context.globalCompositeOperation = eraseMode ? "destination-out" : "source-over";
context.fillStyle = "#fff"; if (eraseMode) {
draw(); draw(context);
} else {
if (!checkerPatternRef.current) {
checkerPatternRef.current = createCheckerPattern(context);
}
context.fillStyle = checkerPatternRef.current ?? "rgba(230,233,238,0.9)";
draw(context);
context.fillStyle = "rgba(255,255,255,0.28)";
draw(context);
}
context.restore(); context.restore();
setHasMask(true);
renderPreview();
}, [eraseMode, renderPreview]);
const stampBrush = useCallback((point: { x: number; y: number }) => { maskContext.save();
const canvas = canvasRef.current; maskContext.globalCompositeOperation = eraseMode ? "destination-out" : "source-over";
const context = canvas?.getContext("2d"); maskContext.fillStyle = "#fff";
if (!canvas || !context) return; draw(maskContext);
maskContext.restore();
}, [eraseMode]);
withBrushComposite(context, () => { const stampBrush = useCallback((point: { x: number; y: number }) => {
withBrushComposite((context) => {
const halfSize = brushSize / 2; const halfSize = brushSize / 2;
context.beginPath(); context.beginPath();
context.arc(point.x, point.y, halfSize, 0, Math.PI * 2); context.arc(point.x, point.y, halfSize, 0, Math.PI * 2);
@ -319,17 +323,13 @@ export function MediaInpaintingOverlay({
}, [brushSize, withBrushComposite]); }, [brushSize, withBrushComposite]);
const fillRectangle = useCallback((start: { x: number; y: number }, end: { x: number; y: number }) => { const fillRectangle = useCallback((start: { x: number; y: number }, end: { x: number; y: number }) => {
const canvas = canvasRef.current;
const context = canvas?.getContext("2d");
if (!canvas || !context) return;
const left = Math.min(start.x, end.x); const left = Math.min(start.x, end.x);
const top = Math.min(start.y, end.y); const top = Math.min(start.y, end.y);
const width = Math.abs(end.x - start.x); const width = Math.abs(end.x - start.x);
const height = Math.abs(end.y - start.y); const height = Math.abs(end.y - start.y);
if (width <= 0 || height <= 0) return; if (width <= 0 || height <= 0) return;
withBrushComposite(context, () => { withBrushComposite((context) => {
context.fillRect(left, top, width, height); context.fillRect(left, top, width, height);
}); });
}, [withBrushComposite]); }, [withBrushComposite]);
@ -355,7 +355,7 @@ export function MediaInpaintingOverlay({
const point = pointFromEvent(event); const point = pointFromEvent(event);
if (!point) return; if (!point) return;
const displayPoint = displayPointFromEvent(event); const displayPoint = displayPointFromEvent(event);
if (displayPoint) setCursorPoint(displayPoint); if (displayPoint) updateCursor(displayPoint);
event.preventDefault(); event.preventDefault();
event.stopPropagation(); event.stopPropagation();
pushHistory(); pushHistory();
@ -371,7 +371,7 @@ export function MediaInpaintingOverlay({
const handlePointerMove = useCallback((event: React.PointerEvent) => { const handlePointerMove = useCallback((event: React.PointerEvent) => {
const displayPoint = displayPointFromEvent(event); const displayPoint = displayPointFromEvent(event);
if (displayPoint) setCursorPoint(displayPoint); if (displayPoint) updateCursor(displayPoint);
if (!drawingRef.current || busy) return; if (!drawingRef.current || busy) return;
const point = pointFromEvent(event); const point = pointFromEvent(event);
if (!point) return; if (!point) return;
@ -383,17 +383,17 @@ export function MediaInpaintingOverlay({
return; return;
} }
drawTo(point); drawTo(point);
}, [activeDrawTool, busy, displayPointFromEvent, drawTo, pointFromEvent]); }, [activeDrawTool, busy, displayPointFromEvent, drawTo, pointFromEvent, updateCursor]);
const handlePointerEnter = useCallback((event: React.PointerEvent) => { const handlePointerEnter = useCallback((event: React.PointerEvent) => {
const displayPoint = displayPointFromEvent(event); const displayPoint = displayPointFromEvent(event);
if (displayPoint) setCursorPoint(displayPoint); if (displayPoint) updateCursor(displayPoint);
}, [displayPointFromEvent]); }, [displayPointFromEvent, updateCursor]);
const handlePointerLeave = useCallback(() => { const handlePointerLeave = useCallback(() => {
if (drawingRef.current) return; if (drawingRef.current) return;
setCursorPoint(null); updateCursor(null);
}, []); }, [updateCursor]);
const stopDrawing = useCallback((event: React.PointerEvent) => { const stopDrawing = useCallback((event: React.PointerEvent) => {
if (!drawingRef.current) return; if (!drawingRef.current) return;
@ -408,9 +408,8 @@ export function MediaInpaintingOverlay({
lastPointRef.current = null; lastPointRef.current = null;
rectangleStartRef.current = null; rectangleStartRef.current = null;
const displayPoint = displayPointFromEvent(event); const displayPoint = displayPointFromEvent(event);
if (displayPoint) setCursorPoint(displayPoint); if (displayPoint) updateCursor(displayPoint);
void exportMask(); }, [activeDrawTool, displayPointFromEvent, fillRectangle, pointFromEvent, updateCursor]);
}, [activeDrawTool, displayPointFromEvent, exportMask, fillRectangle, pointFromEvent]);
const displayBrushSize = bounds && naturalDimensions const displayBrushSize = bounds && naturalDimensions
? Math.max(8, brushSize * (bounds.width / naturalDimensions.width)) ? Math.max(8, brushSize * (bounds.width / naturalDimensions.width))
@ -427,39 +426,46 @@ export function MediaInpaintingOverlay({
useEffect(() => { useEffect(() => {
if (clearSignal <= 0) return; if (clearSignal <= 0) return;
const canvas = canvasRef.current; const canvas = canvasRef.current;
const maskCanvas = maskCanvasRef.current;
const context = canvas?.getContext("2d"); const context = canvas?.getContext("2d");
if (!canvas || !context) return; const maskContext = maskCanvas?.getContext("2d");
if (!canvas || !maskCanvas || !context || !maskContext) return;
pushHistory(); pushHistory();
context.clearRect(0, 0, canvas.width, canvas.height); context.clearRect(0, 0, canvas.width, canvas.height);
renderPreview(); maskContext.clearRect(0, 0, maskCanvas.width, maskCanvas.height);
setHasMask(false); }, [clearSignal, pushHistory]);
onMaskChange?.(null);
}, [clearSignal, onMaskChange, pushHistory, renderPreview]);
useEffect(() => { useEffect(() => {
if (undoSignal <= 0) return; if (undoSignal <= 0) return;
const canvas = canvasRef.current; const canvas = canvasRef.current;
const maskCanvas = maskCanvasRef.current;
const context = canvas?.getContext("2d"); const context = canvas?.getContext("2d");
const maskContext = maskCanvas?.getContext("2d");
const previous = historyRef.current.pop(); const previous = historyRef.current.pop();
if (!canvas || !context || !previous) return; if (!canvas || !maskCanvas || !context || !maskContext || !previous) return;
redoRef.current.push(context.getImageData(0, 0, canvas.width, canvas.height)); redoRef.current.push({
context.putImageData(previous, 0, 0); visible: context.getImageData(0, 0, canvas.width, canvas.height),
renderPreview(); mask: maskContext.getImageData(0, 0, maskCanvas.width, maskCanvas.height),
void exportMask(); });
}, [exportMask, renderPreview, undoSignal]); context.putImageData(previous.visible, 0, 0);
maskContext.putImageData(previous.mask, 0, 0);
}, [undoSignal]);
useEffect(() => { useEffect(() => {
if (redoSignal <= 0) return; if (redoSignal <= 0) return;
const canvas = canvasRef.current; const canvas = canvasRef.current;
const maskCanvas = maskCanvasRef.current;
const context = canvas?.getContext("2d"); const context = canvas?.getContext("2d");
const maskContext = maskCanvas?.getContext("2d");
const next = redoRef.current.pop(); const next = redoRef.current.pop();
if (!canvas || !context || !next) return; if (!canvas || !maskCanvas || !context || !maskContext || !next) return;
historyRef.current.push(context.getImageData(0, 0, canvas.width, canvas.height)); historyRef.current.push({
context.putImageData(next, 0, 0); visible: context.getImageData(0, 0, canvas.width, canvas.height),
renderPreview(); mask: maskContext.getImageData(0, 0, maskCanvas.width, maskCanvas.height),
setHasMask(true); });
void exportMask(); context.putImageData(next.visible, 0, 0);
}, [exportMask, redoSignal, renderPreview]); maskContext.putImageData(next.mask, 0, 0);
}, [redoSignal]);
return ( return (
<div <div
@ -469,8 +475,8 @@ export function MediaInpaintingOverlay({
> >
{bounds && naturalDimensions && ( {bounds && naturalDimensions && (
<canvas <canvas
ref={canvasRef} ref={setCanvasElement}
className={`absolute rounded-[2px] opacity-0 ${activeDrawTool === "rectangle" ? "cursor-default" : "cursor-none"}`} className={`absolute rounded-[2px] ${activeDrawTool === "rectangle" ? "cursor-default" : "cursor-none"}`}
style={{ style={{
left: bounds.left, left: bounds.left,
top: bounds.top, top: bounds.top,
@ -485,27 +491,16 @@ export function MediaInpaintingOverlay({
onPointerCancel={stopDrawing} onPointerCancel={stopDrawing}
/> />
)} )}
{bounds && naturalDimensions && ( {bounds && naturalDimensions && !busy && activeDrawTool === "brush" ? (
<canvas
ref={previewCanvasRef}
className="pointer-events-none absolute rounded-[2px]"
style={{
left: bounds.left,
top: bounds.top,
width: bounds.width,
height: bounds.height,
}}
/>
)}
{cursorPoint && bounds && naturalDimensions && !busy && activeDrawTool === "brush" ? (
<div <div
ref={cursorRef}
className="pointer-events-none absolute rounded-full border border-white shadow-[0_0_0_1px_rgba(0,0,0,0.55)]" className="pointer-events-none absolute rounded-full border border-white shadow-[0_0_0_1px_rgba(0,0,0,0.55)]"
style={{ style={{
left: cursorPoint.x, left: 0,
top: cursorPoint.y, top: 0,
display: "none",
width: displayBrushSize, width: displayBrushSize,
height: displayBrushSize, height: displayBrushSize,
transform: "translate(-50%, -50%)",
background: "transparent", background: "transparent",
}} }}
/> />
@ -518,4 +513,4 @@ export function MediaInpaintingOverlay({
) : null} ) : null}
</div> </div>
); );
} });

6
src/components/media/MediaInpaintingToolbar.tsx

@ -11,12 +11,11 @@ import {
UndoOutlined, UndoOutlined,
} from "@ant-design/icons"; } from "@ant-design/icons";
import { GenerateActionIcon } from "./MediaToolbarIcons"; import { GenerateActionIcon } from "./MediaToolbarIcons";
import type { MediaInpaintingDrawTool, MediaInpaintingMaskPayload } from "./MediaInpaintingOverlay"; import type { MediaInpaintingDrawTool } from "./MediaInpaintingOverlay";
export type InpaintingBatchSize = 1 | 2 | 4; export type InpaintingBatchSize = 1 | 2 | 4;
interface MediaInpaintingToolbarProps { interface MediaInpaintingToolbarProps {
mask: MediaInpaintingMaskPayload | null;
modelName?: string; modelName?: string;
prompt: string; prompt: string;
batchSize: InpaintingBatchSize; batchSize: InpaintingBatchSize;
@ -44,7 +43,6 @@ const BATCH_SIZE_OPTIONS = [
]; ];
export function MediaInpaintingToolbar({ export function MediaInpaintingToolbar({
mask,
modelName, modelName,
prompt, prompt,
batchSize, batchSize,
@ -64,7 +62,7 @@ export function MediaInpaintingToolbar({
onCancel, onCancel,
onConfirm, onConfirm,
}: MediaInpaintingToolbarProps) { }: MediaInpaintingToolbarProps) {
const canConfirm = Boolean(mask) && !busy; const canConfirm = !busy;
if (!frameRect) return null; if (!frameRect) return null;
const centerX = frameRect.left + frameRect.width / 2; const centerX = frameRect.left + frameRect.width / 2;

18
src/components/media/image-edit/InpaintingEditSession.tsx

@ -1,9 +1,9 @@
"use client"; "use client";
import { useState } from "react"; import { useRef, useState } from "react";
import type { ProviderModel } from "@/lib/providers/types"; import type { ProviderModel } from "@/lib/providers/types";
import { uploadBlobToGatewayMedia } from "@/utils/gatewayMediaUpload"; import { uploadBlobToGatewayMedia } from "@/utils/gatewayMediaUpload";
import { MediaInpaintingOverlay, type MediaInpaintingDrawTool, type MediaInpaintingFrameRect, type MediaInpaintingMaskPayload } from "@/components/media/MediaInpaintingOverlay"; import { MediaInpaintingOverlay, type MediaInpaintingDrawTool, type MediaInpaintingFrameRect, type MediaInpaintingOverlayHandle } from "@/components/media/MediaInpaintingOverlay";
import { MediaInpaintingToolbar, type InpaintingBatchSize } from "@/components/media/MediaInpaintingToolbar"; import { MediaInpaintingToolbar, type InpaintingBatchSize } from "@/components/media/MediaInpaintingToolbar";
import { useImageDerivedNodeActions } from "./useImageDerivedNodeActions"; import { useImageDerivedNodeActions } from "./useImageDerivedNodeActions";
import type { ImageEditSessionCommonProps } from "./types"; import type { ImageEditSessionCommonProps } from "./types";
@ -20,7 +20,7 @@ export function InpaintingEditSession({
model, model,
onClose, onClose,
}: InpaintingEditSessionProps) { }: InpaintingEditSessionProps) {
const [mask, setMask] = useState<MediaInpaintingMaskPayload | null>(null); const overlayRef = useRef<MediaInpaintingOverlayHandle>(null);
const [prompt, setPrompt] = useState(""); const [prompt, setPrompt] = useState("");
const [batchSize, setBatchSize] = useState<InpaintingBatchSize>(1); const [batchSize, setBatchSize] = useState<InpaintingBatchSize>(1);
const [brushSize, setBrushSize] = useState(90); const [brushSize, setBrushSize] = useState(90);
@ -37,10 +37,15 @@ export function InpaintingEditSession({
}); });
const handleConfirm = () => { const handleConfirm = () => {
if (!model || !mask || busy) return; if (!model || busy) return;
setBusy(true); setBusy(true);
onClose();
void (async () => { void (async () => {
const mask = await overlayRef.current?.exportMask();
if (!mask) {
setBusy(false);
return;
}
onClose();
const uploadedMask = await uploadBlobToGatewayMedia( const uploadedMask = await uploadBlobToGatewayMedia(
mask.blob, mask.blob,
`inpainting-mask-${Date.now()}.png`, `inpainting-mask-${Date.now()}.png`,
@ -61,6 +66,7 @@ export function InpaintingEditSession({
return ( return (
<> <>
<MediaInpaintingOverlay <MediaInpaintingOverlay
ref={overlayRef}
mediaElementRef={imageRef} mediaElementRef={imageRef}
naturalWidth={dimensions?.width} naturalWidth={dimensions?.width}
naturalHeight={dimensions?.height} naturalHeight={dimensions?.height}
@ -71,11 +77,9 @@ export function InpaintingEditSession({
clearSignal={clearSignal} clearSignal={clearSignal}
undoSignal={undoSignal} undoSignal={undoSignal}
redoSignal={redoSignal} redoSignal={redoSignal}
onMaskChange={setMask}
onFrameChange={setFrameRect} onFrameChange={setFrameRect}
/> />
<MediaInpaintingToolbar <MediaInpaintingToolbar
mask={mask}
modelName={model?.name} modelName={model?.name}
prompt={prompt} prompt={prompt}
batchSize={batchSize} batchSize={batchSize}

6
src/components/nodes/GenerateImageNode.tsx

@ -610,6 +610,12 @@ export function GenerateImageNodeView({ id, data, selected, renderInputHandle =
label: "扩图", label: "扩图",
disabled: !outpaintingModel, disabled: !outpaintingModel,
}, },
{
key: "inpainting",
icon: <GenerateImageEditMenuIcon tool="inpainting" className="h-4 w-4" />,
label: "擦除",
disabled: !inpaintingModel,
},
{ {
key: "crop", key: "crop",
icon: <GenerateImageEditMenuIcon tool="crop" className="h-4 w-4" />, icon: <GenerateImageEditMenuIcon tool="crop" className="h-4 w-4" />,

6
src/components/nodes/ImageInputNode.tsx

@ -414,6 +414,12 @@ export function ImageInputNodeView({ id, data, selected, renderInputHandle = fal
label: "扩图", label: "扩图",
disabled: !outpaintingModel, disabled: !outpaintingModel,
}, },
{
key: "inpainting",
icon: <span className="[&_svg]:h-4 [&_svg]:w-4"><ImageToolbarIcon kind="inpainting" /></span>,
label: "擦除",
disabled: !inpaintingModel,
},
{ {
key: "crop", key: "crop",
icon: <span className="[&_svg]:h-4 [&_svg]:w-4"><ImageToolbarIcon kind="crop" /></span>, icon: <span className="[&_svg]:h-4 [&_svg]:w-4"><ImageToolbarIcon kind="crop" /></span>,

Loading…
Cancel
Save