add paint by example

This commit is contained in:
Qing
2022-12-10 22:06:15 +08:00
parent 6e9d3d8442
commit 203f2bc9c7
18 changed files with 572 additions and 82 deletions

View File

@@ -12,7 +12,8 @@ export default async function inpaint(
sizeLimit?: string,
seed?: number,
maskBase64?: string,
customMask?: File
customMask?: File,
paintByExampleImage?: File
) {
// 1080, 2000, Original
const fd = new FormData()
@@ -48,6 +49,7 @@ export default async function inpaint(
fd.append('croperHeight', croperRect.height.toString())
fd.append('croperWidth', croperRect.width.toString())
fd.append('useCroper', settings.showCroper ? 'true' : 'false')
fd.append('sdMaskBlur', settings.sdMaskBlur.toString())
fd.append('sdStrength', settings.sdStrength.toString())
fd.append('sdSteps', settings.sdSteps.toString())
@@ -59,6 +61,26 @@ export default async function inpaint(
fd.append('cv2Radius', settings.cv2Radius.toString())
fd.append('cv2Flag', settings.cv2Flag.toString())
fd.append('paintByExampleSteps', settings.paintByExampleSteps.toString())
fd.append(
'paintByExampleGuidanceScale',
settings.paintByExampleGuidanceScale.toString()
)
fd.append('paintByExampleSeed', seed ? seed.toString() : '-1')
fd.append(
'paintByExampleMaskBlur',
settings.paintByExampleMaskBlur.toString()
)
fd.append(
'paintByExampleMatchHistograms',
settings.paintByExampleMatchHistograms ? 'true' : 'false'
)
// TODO: resize image's shortest_edge to 224 before pass to backend, save network time?
// https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPImageProcessor
if (paintByExampleImage) {
fd.append('paintByExampleImage', paintByExampleImage)
}
if (sizeLimit === undefined) {
fd.append('sizeLimit', '1080')
} else {

View File

@@ -39,6 +39,7 @@ import {
isInpaintingState,
isInteractiveSegRunningState,
isInteractiveSegState,
isPaintByExampleState,
isSDState,
negativePropmtState,
propmtState,
@@ -53,6 +54,7 @@ import emitter, {
EVENT_PROMPT,
EVENT_CUSTOM_MASK,
CustomMaskEventData,
EVENT_PAINT_BY_EXAMPLE,
} from '../../event'
import FileSelect from '../FileSelect/FileSelect'
import InteractiveSeg from '../InteractiveSeg/InteractiveSeg'
@@ -108,6 +110,7 @@ export default function Editor() {
const [isInpainting, setIsInpainting] = useRecoilState(isInpaintingState)
const runMannually = useRecoilValue(runManuallyState)
const isSD = useRecoilValue(isSDState)
const isPaintByExample = useRecoilValue(isPaintByExampleState)
const [isInteractiveSeg, setIsInteractiveSeg] = useRecoilState(
isInteractiveSegState
)
@@ -262,8 +265,11 @@ export default function Editor() {
async (
useLastLineGroup?: boolean,
customMask?: File,
maskImage?: HTMLImageElement | null
maskImage?: HTMLImageElement | null,
paintByExampleImage?: File
) => {
// customMask: mask uploaded by user
// maskImage: mask from interactive segmentation
if (file === undefined) {
return
}
@@ -328,9 +334,6 @@ export default function Editor() {
}
}
const sdSeed = settings.sdSeedFixed ? settings.sdSeed : -1
console.log({ useCustomMask })
try {
const res = await inpaint(
targetFile,
@@ -339,15 +342,16 @@ export default function Editor() {
promptVal,
negativePromptVal,
sizeLimit.toString(),
sdSeed,
seedVal,
useCustomMask ? undefined : maskCanvas.toDataURL(),
useCustomMask ? customMask : undefined
useCustomMask ? customMask : undefined,
paintByExampleImage
)
if (!res) {
throw new Error('Something went wrong on server side.')
}
const { blob, seed } = res
if (seed && !settings.sdSeedFixed) {
if (seed) {
setSeed(parseInt(seed, 10))
}
const newRender = new Image()
@@ -395,6 +399,7 @@ export default function Editor() {
drawOnCurrentRender,
hadDrawSomething,
drawLinesOnMask,
seedVal,
]
)
@@ -439,6 +444,31 @@ export default function Editor() {
}
}, [runInpainting])
useEffect(() => {
emitter.on(EVENT_PAINT_BY_EXAMPLE, (data: any) => {
if (hadDrawSomething() || interactiveSegMask) {
runInpainting(false, undefined, interactiveSegMask, data.image)
} else if (lastLineGroup.length !== 0) {
// 使用上一次手绘的 mask 生成
runInpainting(true, undefined, prevInteractiveSegMask, data.image)
} else if (prevInteractiveSegMask) {
// 使用上一次 IS 的 mask 生成
runInpainting(false, undefined, prevInteractiveSegMask, data.image)
} else {
setToastState({
open: true,
desc: 'Please draw mask on picture',
state: 'error',
duration: 1500,
})
}
})
return () => {
emitter.off(EVENT_PAINT_BY_EXAMPLE)
}
}, [runInpainting])
const hadRunInpainting = () => {
return renders.length !== 0
}
@@ -793,7 +823,11 @@ export default function Editor() {
return
}
if (isSD && settings.showCroper && isOutsideCroper(mouseXY(ev))) {
if (
(isSD || isPaintByExample) &&
settings.showCroper &&
isOutsideCroper(mouseXY(ev))
) {
return
}
@@ -876,7 +910,12 @@ export default function Editor() {
return false
}
useKey(undoPredicate, undo, undefined, [undoStroke, undoRender, isSD])
useKey(undoPredicate, undo, undefined, [
undoStroke,
undoRender,
runMannually,
curLineGroup,
])
const disableUndo = () => {
if (isInteractiveSeg) {
@@ -955,7 +994,12 @@ export default function Editor() {
return false
}
useKey(redoPredicate, redo, undefined, [redoStroke, redoRender, isSD])
useKey(redoPredicate, redo, undefined, [
redoStroke,
redoRender,
runMannually,
redoCurLines,
])
const disableRedo = () => {
if (isInteractiveSeg) {
@@ -1295,7 +1339,7 @@ export default function Editor() {
</div>
</div>
{isSD && settings.showCroper ? (
{(isSD || isPaintByExample) && settings.showCroper ? (
<Croper
maxHeight={original.naturalHeight}
maxWidth={original.naturalWidth}
@@ -1358,7 +1402,7 @@ export default function Editor() {
)}
<div className="editor-toolkit-panel">
{isSD || file === undefined ? (
{isSD || isPaintByExample || file === undefined ? (
<></>
) : (
<SizeSelector
@@ -1466,7 +1510,7 @@ export default function Editor() {
onClick={download}
/>
{settings.runInpaintingManually && !isSD && (
{settings.runInpaintingManually && !isSD && !isPaintByExample && (
<Button
toolTip="Run Inpainting"
tooltipPosition="top"

View File

@@ -193,6 +193,8 @@ function ModelSettingBlock() {
return undefined
case AIModel.SD2:
return undefined
case AIModel.PAINT_BY_EXAMPLE:
return undefined
case AIModel.Mange:
return undefined
case AIModel.CV2:
@@ -258,6 +260,12 @@ function ModelSettingBlock() {
'https://docs.opencv.org/4.6.0/df/d3d/tutorial_py_inpainting.html',
'https://docs.opencv.org/4.6.0/df/d3d/tutorial_py_inpainting.html'
)
case AIModel.PAINT_BY_EXAMPLE:
return renderModelDesc(
'Paint by Example',
'https://arxiv.org/abs/2211.13227',
'https://github.com/Fantasy-Studio/Paint-by-Example'
)
default:
return <></>
}
@@ -270,7 +278,6 @@ function ModelSettingBlock() {
titleSuffix={renderPaperCodeBadge()}
input={
<Selector
width={80}
value={setting.model as string}
options={Object.values(AIModel)}
onChange={val => onModelChange(val as AIModel)}

View File

@@ -0,0 +1,231 @@
import React, { useState } from 'react'
import { useRecoilState, useRecoilValue } from 'recoil'
import * as PopoverPrimitive from '@radix-ui/react-popover'
import { useToggle } from 'react-use'
import { UploadIcon } from '@radix-ui/react-icons'
import {
isInpaintingState,
paintByExampleImageState,
settingState,
} from '../../store/Atoms'
import NumberInputSetting from '../Settings/NumberInputSetting'
import SettingBlock from '../Settings/SettingBlock'
import { Switch, SwitchThumb } from '../shared/Switch'
import Button from '../shared/Button'
import emitter, { EVENT_PAINT_BY_EXAMPLE } from '../../event'
import { useImage } from '../../utils'
const INPUT_WIDTH = 30
const PESidePanel = () => {
const [open, toggleOpen] = useToggle(true)
const [setting, setSettingState] = useRecoilState(settingState)
const [paintByExampleImage, setPaintByExampleImage] = useRecoilState(
paintByExampleImageState
)
const [uploadElemId] = useState(
`example-file-upload-${Math.random().toString()}`
)
const [exampleImage, isExampleImageLoaded] = useImage(paintByExampleImage)
const isInpainting = useRecoilValue(isInpaintingState)
const renderUploadIcon = () => {
return (
<label htmlFor={uploadElemId}>
<Button
border
toolTip="Upload example image"
tooltipPosition="top"
icon={<UploadIcon />}
style={{ padding: '0.3rem', gap: 0 }}
>
<input
style={{ display: 'none' }}
id={uploadElemId}
name={uploadElemId}
type="file"
onChange={ev => {
const newFile = ev.currentTarget.files?.[0]
if (newFile) {
setPaintByExampleImage(newFile)
}
}}
accept="image/png, image/jpeg"
/>
</Button>
</label>
)
}
return (
<div className="side-panel">
<PopoverPrimitive.Root open={open}>
<PopoverPrimitive.Trigger
className="btn-primary side-panel-trigger"
onClick={() => toggleOpen()}
>
Configurations
</PopoverPrimitive.Trigger>
<PopoverPrimitive.Portal>
<PopoverPrimitive.Content className="side-panel-content">
<SettingBlock
title="Croper"
input={
<Switch
checked={setting.showCroper}
onCheckedChange={value => {
setSettingState(old => {
return { ...old, showCroper: value }
})
}}
>
<SwitchThumb />
</Switch>
}
/>
<NumberInputSetting
title="Steps"
width={INPUT_WIDTH}
value={`${setting.paintByExampleSteps}`}
desc="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference."
onValue={value => {
const val = value.length === 0 ? 0 : parseInt(value, 10)
setSettingState(old => {
return { ...old, paintByExampleSteps: val }
})
}}
/>
<NumberInputSetting
title="Guidance Scale"
width={INPUT_WIDTH}
allowFloat
value={`${setting.paintByExampleGuidanceScale}`}
desc="Higher guidance scale encourages to generate images that are close to the example image"
onValue={value => {
const val = value.length === 0 ? 0 : parseFloat(value)
setSettingState(old => {
return { ...old, paintByExampleGuidanceScale: val }
})
}}
/>
<NumberInputSetting
title="Mask Blur"
width={INPUT_WIDTH}
value={`${setting.paintByExampleMaskBlur}`}
desc="Blur the edge of mask area. The higher the number the smoother blend with the original image"
onValue={value => {
const val = value.length === 0 ? 0 : parseInt(value, 10)
setSettingState(old => {
return { ...old, paintByExampleMaskBlur: val }
})
}}
/>
<SettingBlock
title="Match Histograms"
desc="Match the inpainting result histogram to the source image histogram, will improves the inpainting quality for some images."
input={
<Switch
checked={setting.paintByExampleMatchHistograms}
onCheckedChange={value => {
setSettingState(old => {
return { ...old, paintByExampleMatchHistograms: value }
})
}}
>
<SwitchThumb />
</Switch>
}
/>
<SettingBlock
title="Seed"
input={
<div
style={{
display: 'flex',
gap: 0,
justifyContent: 'center',
alignItems: 'center',
}}
>
{/* 每次会从服务器返回更新该值 */}
<NumberInputSetting
title=""
width={80}
value={`${setting.paintByExampleSeed}`}
desc=""
disable={!setting.paintByExampleSeedFixed}
onValue={value => {
const val = value.length === 0 ? 0 : parseInt(value, 10)
setSettingState(old => {
return { ...old, paintByExampleSeed: val }
})
}}
/>
<Switch
checked={setting.paintByExampleSeedFixed}
onCheckedChange={value => {
setSettingState(old => {
return { ...old, paintByExampleSeedFixed: value }
})
}}
style={{ marginLeft: '8px' }}
>
<SwitchThumb />
</Switch>
</div>
}
/>
<div style={{ display: 'flex', flexDirection: 'column' }}>
<SettingBlock title="Example Image" input={renderUploadIcon()} />
{paintByExampleImage ? (
<div
style={{
display: 'flex',
justifyContent: 'center',
alignItems: 'center',
}}
>
<img
src={exampleImage.src}
alt="example"
style={{
maxWidth: 200,
maxHeight: 200,
margin: 12,
}}
/>
</div>
) : (
<></>
)}
</div>
<Button
border
disabled={!isExampleImageLoaded || isInpainting}
style={{ width: '100%' }}
onClick={() => {
if (isExampleImageLoaded) {
emitter.emit(EVENT_PAINT_BY_EXAMPLE, {
image: paintByExampleImage,
})
}
}}
>
Paint
</Button>
</PopoverPrimitive.Content>
</PopoverPrimitive.Portal>
</PopoverPrimitive.Root>
</div>
)
}
export default PESidePanel

View File

@@ -7,6 +7,7 @@ import Toast from './shared/Toast'
import {
AIModel,
fileState,
isPaintByExampleState,
isSDState,
settingState,
toastState,
@@ -17,12 +18,14 @@ import {
switchModel,
} from '../adapters/inpainting'
import SidePanel from './SidePanel/SidePanel'
import PESidePanel from './SidePanel/PESidePanel'
const Workspace = () => {
const [file, setFile] = useRecoilState(fileState)
const [settings, setSettingState] = useRecoilState(settingState)
const [toastVal, setToastState] = useRecoilState(toastState)
const isSD = useRecoilValue(isSDState)
const isPaintByExample = useRecoilValue(isPaintByExampleState)
const onSettingClose = async () => {
const curModel = await currentModel().then(res => res.text())
@@ -88,6 +91,7 @@ const Workspace = () => {
return (
<>
{isSD ? <SidePanel /> : <></>}
{isPaintByExample ? <PESidePanel /> : <></>}
<Editor />
<SettingModal onClose={onSettingClose} />
<ShortcutsModal />

View File

@@ -1,11 +1,17 @@
import mitt from 'mitt'
export const EVENT_PROMPT = 'prompt'
export const EVENT_CUSTOM_MASK = 'custom_mask'
export interface CustomMaskEventData {
mask: File
}
export const EVENT_PAINT_BY_EXAMPLE = 'paint_by_example'
export interface PaintByExampleEventData {
image: File
}
const emitter = mitt()
export default emitter

View File

@@ -13,6 +13,7 @@ export enum AIModel {
SD2 = 'sd2',
CV2 = 'cv2',
Mange = 'manga',
PAINT_BY_EXAMPLE = 'paint_by_example',
}
export const maskState = atom<File | undefined>({
@@ -20,6 +21,11 @@ export const maskState = atom<File | undefined>({
default: undefined,
})
export const paintByExampleImageState = atom<File | undefined>({
key: 'paintByExampleImageState',
default: undefined,
})
export interface Rect {
x: number
y: number
@@ -252,6 +258,14 @@ export interface Settings {
// For OpenCV2
cv2Radius: number
cv2Flag: CV2Flag
// Paint by Example
paintByExampleSteps: number
paintByExampleGuidanceScale: number
paintByExampleSeed: number
paintByExampleSeedFixed: boolean
paintByExampleMaskBlur: number
paintByExampleMatchHistograms: boolean
}
const defaultHDSettings: ModelsHDSettings = {
@@ -304,6 +318,13 @@ const defaultHDSettings: ModelsHDSettings = {
hdStrategyCropMargin: 128,
enabled: false,
},
[AIModel.PAINT_BY_EXAMPLE]: {
hdStrategy: HDStrategy.ORIGINAL,
hdStrategyResizeLimit: 768,
hdStrategyCropTrigerSize: 512,
hdStrategyCropMargin: 128,
enabled: false,
},
[AIModel.Mange]: {
hdStrategy: HDStrategy.CROP,
hdStrategyResizeLimit: 1280,
@@ -364,6 +385,14 @@ export const settingStateDefault: Settings = {
// CV2
cv2Radius: 5,
cv2Flag: CV2Flag.INPAINT_NS,
// Paint by Example
paintByExampleSteps: 50,
paintByExampleGuidanceScale: 7.5,
paintByExampleSeed: 42,
paintByExampleMaskBlur: 5,
paintByExampleSeedFixed: false,
paintByExampleMatchHistograms: false,
}
const localStorageEffect =
@@ -401,11 +430,28 @@ export const seedState = selector({
key: 'seed',
get: ({ get }) => {
const settings = get(settingState)
return settings.sdSeed
switch (settings.model) {
case AIModel.PAINT_BY_EXAMPLE:
return settings.paintByExampleSeedFixed
? settings.paintByExampleSeed
: -1
default:
return settings.sdSeedFixed ? settings.sdSeed : -1
}
},
set: ({ get, set }, newValue: any) => {
const settings = get(settingState)
set(settingState, { ...settings, sdSeed: newValue })
switch (settings.model) {
case AIModel.PAINT_BY_EXAMPLE:
if (!settings.paintByExampleSeedFixed) {
set(settingState, { ...settings, paintByExampleSeed: newValue })
}
break
default:
if (!settings.sdSeedFixed) {
set(settingState, { ...settings, sdSeed: newValue })
}
}
},
})
@@ -435,11 +481,20 @@ export const isSDState = selector({
},
})
export const isPaintByExampleState = selector({
key: 'isPaintByExampleState',
get: ({ get }) => {
const settings = get(settingState)
return settings.model === AIModel.PAINT_BY_EXAMPLE
},
})
export const runManuallyState = selector({
key: 'runManuallyState',
get: ({ get }) => {
const settings = get(settingState)
const isSD = get(isSDState)
return settings.runInpaintingManually || isSD
const isPaintByExample = get(isPaintByExampleState)
return settings.runInpaintingManually || isSD || isPaintByExample
},
})