From eec41734c370a923dab58cbb0b5a191954e8b1aa Mon Sep 17 00:00:00 2001 From: Qing Date: Sun, 13 Nov 2022 23:31:11 +0800 Subject: [PATCH] add custom mask upload, WIP, need more test better handle server error --- lama_cleaner/app/src/adapters/inpainting.ts | 18 +++++--- .../app/src/components/Editor/Editor.tsx | 42 +++++++++++------ .../app/src/components/Header/Header.tsx | 46 +++++++++++++++++-- lama_cleaner/app/src/event.ts | 4 ++ lama_cleaner/server.py | 18 ++++++-- 5 files changed, 99 insertions(+), 29 deletions(-) diff --git a/lama_cleaner/app/src/adapters/inpainting.ts b/lama_cleaner/app/src/adapters/inpainting.ts index 2cb6206..a8a81d8 100644 --- a/lama_cleaner/app/src/adapters/inpainting.ts +++ b/lama_cleaner/app/src/adapters/inpainting.ts @@ -5,19 +5,23 @@ export const API_ENDPOINT = `${process.env.REACT_APP_INPAINTING_URL}` export default async function inpaint( imageFile: File, - maskBase64: string, settings: Settings, croperRect: Rect, prompt?: string, negativePrompt?: string, sizeLimit?: string, - seed?: number + seed?: number, + maskBase64?: string, + customMask?: File ) { // 1080, 2000, Original const fd = new FormData() fd.append('image', imageFile) - const mask = dataURItoBlob(maskBase64) - fd.append('mask', mask) + if (maskBase64 !== undefined) { + fd.append('mask', dataURItoBlob(maskBase64)) + } else if (customMask !== undefined) { + fd.append('mask', customMask) + } const hdSettings = settings.hdSettings[settings.model] fd.append('ldmSteps', settings.ldmSteps.toString()) @@ -70,8 +74,10 @@ export default async function inpaint( const newSeed = res.headers.get('x-seed') return { blob: URL.createObjectURL(blob), seed: newSeed } } - } catch { - throw new Error('Something went wrong on server side.') + const errMsg = await res.text() + throw new Error(errMsg) + } catch (error) { + throw new Error(`Something went wrong: ${error}`) } } diff --git a/lama_cleaner/app/src/components/Editor/Editor.tsx b/lama_cleaner/app/src/components/Editor/Editor.tsx index 1653258..22da920 100644 --- a/lama_cleaner/app/src/components/Editor/Editor.tsx +++ b/lama_cleaner/app/src/components/Editor/Editor.tsx @@ -45,7 +45,11 @@ import { } from '../../store/Atoms' import useHotKey from '../../hooks/useHotkey' import Croper from '../Croper/Croper' -import emitter, { EVENT_PROMPT } from '../../event' +import emitter, { + EVENT_PROMPT, + EVENT_CUSTOM_MASK, + CustomMaskEventData, +} from '../../event' import FileSelect from '../FileSelect/FileSelect' const TOOLBAR_SIZE = 200 @@ -195,22 +199,23 @@ export default function Editor() { ) const runInpainting = useCallback( - async (prompt?: string, useLastLineGroup?: boolean) => { + async (useLastLineGroup?: boolean, customMask?: File) => { if (file === undefined) { return } + const useCustomMask = customMask !== undefined // useLastLineGroup 的影响 // 1. 使用上一次的 mask // 2. 结果替换当前 render console.log('runInpainting') - let maskLineGroup = [] + let maskLineGroup: LineGroup = [] if (useLastLineGroup === true) { if (lastLineGroup.length === 0) { return } maskLineGroup = lastLineGroup - } else { + } else if (!useCustomMask) { if (!hadDrawSomething()) { return } @@ -256,23 +261,23 @@ export default function Editor() { const sdSeed = settings.sdSeedFixed ? settings.sdSeed : -1 + console.log({ useCustomMask }) try { const res = await inpaint( targetFile, - maskCanvas.toDataURL(), settings, croperRect, - prompt, + promptVal, negativePromptVal, sizeLimit.toString(), - sdSeed + sdSeed, + useCustomMask ? undefined : maskCanvas.toDataURL(), + useCustomMask ? customMask : undefined ) if (!res) { - throw new Error('empty response') + throw new Error('Something went wrong on server side.') } const { blob, seed } = res - console.log(seed) - console.log(settings.sdSeedFixed) if (seed && !settings.sdSeedFixed) { setSeed(parseInt(seed, 10)) } @@ -324,9 +329,9 @@ export default function Editor() { useEffect(() => { emitter.on(EVENT_PROMPT, () => { if (hadDrawSomething()) { - runInpainting(promptVal) + runInpainting() } else if (lastLineGroup.length !== 0) { - runInpainting(promptVal, true) + runInpainting(true) } else { setToastState({ open: true, @@ -336,10 +341,21 @@ export default function Editor() { }) } }) + return () => { emitter.off(EVENT_PROMPT) } - }, [hadDrawSomething, runInpainting, prompt]) + }, [hadDrawSomething, runInpainting, promptVal]) + + useEffect(() => { + emitter.on(EVENT_CUSTOM_MASK, (data: any) => { + runInpainting(false, data.mask) + }) + + return () => { + emitter.off(EVENT_CUSTOM_MASK) + } + }, [runInpainting]) const hadRunInpainting = () => { return renders.length !== 0 diff --git a/lama_cleaner/app/src/components/Header/Header.tsx b/lama_cleaner/app/src/components/Header/Header.tsx index 1fd0b03..3b9ae56 100644 --- a/lama_cleaner/app/src/components/Header/Header.tsx +++ b/lama_cleaner/app/src/components/Header/Header.tsx @@ -1,25 +1,33 @@ import { ArrowLeftIcon, UploadIcon } from '@heroicons/react/outline' import React, { useState } from 'react' import { useRecoilState, useRecoilValue } from 'recoil' -import { fileState, isSDState } from '../../store/Atoms' +import { fileState, isInpaintingState, isSDState } from '../../store/Atoms' import Button from '../shared/Button' import Shortcuts from '../Shortcuts/Shortcuts' -import useResolution from '../../hooks/useResolution' import { ThemeChanger } from './ThemeChanger' import SettingIcon from '../Settings/SettingIcon' import PromptInput from './PromptInput' import CoffeeIcon from '../CoffeeIcon/CoffeeIcon' +import emitter, { EVENT_CUSTOM_MASK } from '../../event' const Header = () => { + const isInpainting = useRecoilValue(isInpaintingState) const [file, setFile] = useRecoilState(fileState) - const resolution = useResolution() const [uploadElemId] = useState(`file-upload-${Math.random().toString()}`) + const [maskUploadElemId] = useState(`mask-upload-${Math.random().toString()}`) const isSD = useRecoilValue(isSDState) const renderHeader = () => { return (
-
+
+ +
diff --git a/lama_cleaner/app/src/event.ts b/lama_cleaner/app/src/event.ts index e4344fb..cdfc70a 100644 --- a/lama_cleaner/app/src/event.ts +++ b/lama_cleaner/app/src/event.ts @@ -1,6 +1,10 @@ import mitt from 'mitt' export const EVENT_PROMPT = 'prompt' +export const EVENT_CUSTOM_MASK = 'custom_mask' +export interface CustomMaskEventData { + mask: File +} const emitter = mitt() diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index 6c763f6..f7014b0 100644 --- a/lama_cleaner/server.py +++ b/lama_cleaner/server.py @@ -94,6 +94,10 @@ def process(): origin_image_bytes = input["image"].read() image, alpha_channel = load_img(origin_image_bytes) + mask, _ = load_img(input["mask"].read(), gray=True) + if image.shape[:2] != mask.shape[:2]: + return f"Mask shape{mask.shape[:2]} not queal to Image shape{image.shape[:2]}", 400 + original_shape = image.shape interpolation = cv2.INTER_CUBIC @@ -136,14 +140,18 @@ def process(): image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation) logger.info(f"Resized image shape: {image.shape}") - mask, _ = load_img(input["mask"].read(), gray=True) mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation) start = time.time() - res_np_img = model(image, mask, config) - logger.info(f"process time: {(time.time() - start) * 1000}ms") - - torch.cuda.empty_cache() + try: + res_np_img = model(image, mask, config) + except RuntimeError as e: + # NOTE: the string may change? + if "CUDA out of memory. " in str(e): + return "CUDA out of memory", 500 + finally: + logger.info(f"process time: {(time.time() - start) * 1000}ms") + torch.cuda.empty_cache() if alpha_channel is not None: if alpha_channel.shape[:2] != res_np_img.shape[:2]: