From 05e82598de18fb4a98ad1dc60590a18bb3ffddf9 Mon Sep 17 00:00:00 2001 From: Qing Date: Sat, 28 Jan 2023 21:13:21 +0800 Subject: [PATCH] add pix2pix --- lama_cleaner/app/src/adapters/inpainting.ts | 5 + .../app/src/components/Editor/Editor.tsx | 13 +- .../app/src/components/Header/Header.tsx | 4 +- .../components/Settings/ModelSettingBlock.tsx | 20 +- .../src/components/Settings/SettingsModal.tsx | 7 +- .../src/components/SidePanel/P2PSidePanel.tsx | 224 ++++++++++++++++++ lama_cleaner/app/src/components/Workspace.tsx | 4 + lama_cleaner/app/src/store/Atoms.tsx | 41 +++- lama_cleaner/const.py | 3 +- lama_cleaner/model/pix2pix.py | 83 +++++++ lama_cleaner/model_manager.py | 3 +- lama_cleaner/schema.py | 5 + lama_cleaner/server.py | 3 + requirements.txt | 2 +- 14 files changed, 389 insertions(+), 28 deletions(-) create mode 100644 lama_cleaner/app/src/components/SidePanel/P2PSidePanel.tsx create mode 100644 lama_cleaner/model/pix2pix.py diff --git a/lama_cleaner/app/src/adapters/inpainting.ts b/lama_cleaner/app/src/adapters/inpainting.ts index b94aa17..5529a9c 100644 --- a/lama_cleaner/app/src/adapters/inpainting.ts +++ b/lama_cleaner/app/src/adapters/inpainting.ts @@ -82,6 +82,11 @@ export default async function inpaint( fd.append('paintByExampleImage', paintByExampleImage) } + // InstructPix2Pix + fd.append('p2pSteps', settings.p2pSteps.toString()) + fd.append('p2pImageGuidanceScale', settings.p2pImageGuidanceScale.toString()) + fd.append('p2pGuidanceScale', settings.p2pGuidanceScale.toString()) + if (sizeLimit === undefined) { fd.append('sizeLimit', '1080') } else { diff --git a/lama_cleaner/app/src/components/Editor/Editor.tsx b/lama_cleaner/app/src/components/Editor/Editor.tsx index 4b34819..7b34ae8 100644 --- a/lama_cleaner/app/src/components/Editor/Editor.tsx +++ b/lama_cleaner/app/src/components/Editor/Editor.tsx @@ -41,10 +41,10 @@ import { croperState, enableFileManagerState, fileState, - gifImageState, imageHeightState, imageWidthState, interactiveSegClicksState, + isDiffusionModelsState, isInpaintingState, isInteractiveSegRunningState, isInteractiveSegState, @@ -118,8 +118,7 @@ export default function Editor() { const setToastState = useSetRecoilState(toastState) const [isInpainting, setIsInpainting] = useRecoilState(isInpaintingState) const runMannually = useRecoilValue(runManuallyState) - const isSD = useRecoilValue(isSDState) - const isPaintByExample = useRecoilValue(isPaintByExampleState) + const isDiffusionModels = useRecoilValue(isDiffusionModelsState) const [isInteractiveSeg, setIsInteractiveSeg] = useRecoilState( isInteractiveSegState ) @@ -842,7 +841,7 @@ export default function Editor() { } if ( - (isSD || isPaintByExample) && + isDiffusionModels && settings.showCroper && isOutsideCroper(mouseXY(ev)) ) { @@ -1388,7 +1387,7 @@ export default function Editor() { minHeight={Math.min(256, original.naturalHeight)} minWidth={Math.min(256, original.naturalWidth)} scale={scale} - show={(isSD || isPaintByExample) && settings.showCroper} + show={isDiffusionModels && settings.showCroper} /> {isInteractiveSeg ? : <>} @@ -1442,7 +1441,7 @@ export default function Editor() { )}
- {isSD || isPaintByExample || file === undefined ? ( + {isDiffusionModels || file === undefined ? ( <> ) : ( - {settings.runInpaintingManually && !isSD && !isPaintByExample && ( + {settings.runInpaintingManually && !isDiffusionModels && (
- {isSD && file ? : <>} + {(isSD || isPix2Pix) && file ? : <>}
diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx index f5583ba..c0e1b7c 100644 --- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx +++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx @@ -179,28 +179,16 @@ function ModelSettingBlock() { const renderOptionDesc = (): ReactNode => { switch (setting.model) { - case AIModel.LAMA: - return undefined case AIModel.LDM: return renderLDMModelDesc() case AIModel.ZITS: return renderZITSModelDesc() - case AIModel.MAT: - return undefined case AIModel.FCF: return renderFCFModelDesc() - case AIModel.SD15: - return undefined - case AIModel.SD2: - return undefined - case AIModel.PAINT_BY_EXAMPLE: - return undefined - case AIModel.Mange: - return undefined case AIModel.CV2: return renderOpenCV2Desc() default: - return <> + return undefined } } @@ -266,6 +254,12 @@ function ModelSettingBlock() { 'https://arxiv.org/abs/2211.13227', 'https://github.com/Fantasy-Studio/Paint-by-Example' ) + case AIModel.PIX2PIX: + return renderModelDesc( + 'InstructPix2Pix', + 'https://arxiv.org/abs/2211.09800', + 'https://github.com/timothybrooks/instruct-pix2pix' + ) default: return <> } diff --git a/lama_cleaner/app/src/components/Settings/SettingsModal.tsx b/lama_cleaner/app/src/components/Settings/SettingsModal.tsx index 7109a84..aaaf683 100644 --- a/lama_cleaner/app/src/components/Settings/SettingsModal.tsx +++ b/lama_cleaner/app/src/components/Settings/SettingsModal.tsx @@ -1,6 +1,7 @@ import React from 'react' import { useRecoilState, useRecoilValue } from 'recoil' import { + isDiffusionModelsState, isPaintByExampleState, isSDState, settingState, @@ -28,7 +29,7 @@ export default function SettingModal(props: SettingModalProps) { const { onClose } = props const [setting, setSettingState] = useRecoilState(settingState) const isSD = useRecoilValue(isSDState) - const isPaintByExample = useRecoilValue(isPaintByExampleState) + const isDiffusionModels = useRecoilValue(isDiffusionModelsState) const handleOnClose = () => { setSettingState(old => { @@ -56,9 +57,9 @@ export default function SettingModal(props: SettingModalProps) { show={setting.show} > - {isSD || isPaintByExample ? <> : } + {isDiffusionModels ? <> : } - {isSD ? <> : } + {isDiffusionModels ? <> : } ) } diff --git a/lama_cleaner/app/src/components/SidePanel/P2PSidePanel.tsx b/lama_cleaner/app/src/components/SidePanel/P2PSidePanel.tsx new file mode 100644 index 0000000..fd15f39 --- /dev/null +++ b/lama_cleaner/app/src/components/SidePanel/P2PSidePanel.tsx @@ -0,0 +1,224 @@ +import React, { FormEvent } from 'react' +import { useRecoilState, useRecoilValue } from 'recoil' +import * as PopoverPrimitive from '@radix-ui/react-popover' +import { useToggle } from 'react-use' +import { + isInpaintingState, + negativePropmtState, + propmtState, + settingState, +} from '../../store/Atoms' +import NumberInputSetting from '../Settings/NumberInputSetting' +import SettingBlock from '../Settings/SettingBlock' +import { Switch, SwitchThumb } from '../shared/Switch' +import TextAreaInput from '../shared/Textarea' +import emitter, { EVENT_PROMPT } from '../../event' +import ImageResizeScale from './ImageResizeScale' + +const INPUT_WIDTH = 30 + +const P2PSidePanel = () => { + const [open, toggleOpen] = useToggle(true) + const [setting, setSettingState] = useRecoilState(settingState) + const [negativePrompt, setNegativePrompt] = + useRecoilState(negativePropmtState) + const isInpainting = useRecoilValue(isInpaintingState) + const prompt = useRecoilValue(propmtState) + + const handleOnInput = (evt: FormEvent) => { + evt.preventDefault() + evt.stopPropagation() + const target = evt.target as HTMLTextAreaElement + setNegativePrompt(target.value) + } + + const onKeyUp = (e: React.KeyboardEvent) => { + if ( + e.key === 'Enter' && + (e.ctrlKey || e.metaKey) && + prompt.length !== 0 && + !isInpainting + ) { + emitter.emit(EVENT_PROMPT) + } + } + + return ( +
+ + toggleOpen()} + > + Config + + + + { + setSettingState(old => { + return { ...old, showCroper: value } + }) + }} + > + + + } + /> + + + + { + const val = value.length === 0 ? 0 : parseInt(value, 10) + setSettingState(old => { + return { ...old, p2pSteps: val } + }) + }} + /> + + { + const val = value.length === 0 ? 0 : parseFloat(value) + setSettingState(old => { + return { ...old, p2pGuidanceScale: val } + }) + }} + /> + + { + const val = value.length === 0 ? 0 : parseFloat(value) + setSettingState(old => { + return { ...old, p2pImageGuidanceScale: val } + }) + }} + /> + + {/* { + const val = value.length === 0 ? 0 : parseInt(value, 10) + setSettingState(old => { + return { ...old, sdMaskBlur: val } + }) + }} + /> */} + + {/* { + setSettingState(old => { + return { ...old, sdMatchHistograms: value } + }) + }} + > + + + } + /> */} + + {/* { + const sampler = val as SDSampler + setSettingState(old => { + return { ...old, sdSampler: sampler } + }) + }} + /> + } + /> */} + + + { + const val = value.length === 0 ? 0 : parseInt(value, 10) + setSettingState(old => { + return { ...old, sdSeed: val } + }) + }} + /> + { + setSettingState(old => { + return { ...old, sdSeedFixed: value } + }) + }} + style={{ marginLeft: '8px' }} + > + + +
+ } + /> + + + } + /> + + + +
+ ) +} + +export default P2PSidePanel diff --git a/lama_cleaner/app/src/components/Workspace.tsx b/lama_cleaner/app/src/components/Workspace.tsx index 1d3c514..f1bf5b9 100644 --- a/lama_cleaner/app/src/components/Workspace.tsx +++ b/lama_cleaner/app/src/components/Workspace.tsx @@ -8,6 +8,7 @@ import { AIModel, fileState, isPaintByExampleState, + isPix2PixState, isSDState, settingState, showFileManagerState, @@ -22,6 +23,7 @@ import { import SidePanel from './SidePanel/SidePanel' import PESidePanel from './SidePanel/PESidePanel' import FileManager from './FileManager/FileManager' +import P2PSidePanel from './SidePanel/P2PSidePanel' const Workspace = () => { const setFile = useSetRecoilState(fileState) @@ -29,6 +31,7 @@ const Workspace = () => { const [toastVal, setToastState] = useRecoilState(toastState) const isSD = useRecoilValue(isSDState) const isPaintByExample = useRecoilValue(isPaintByExampleState) + const isPix2Pix = useRecoilValue(isPix2PixState) const [showFileManager, setShowFileManager] = useRecoilState(showFileManagerState) @@ -98,6 +101,7 @@ const Workspace = () => { <> {isSD ? : <>} {isPaintByExample ? : <>} + {isPix2Pix ? : <>} ({ @@ -343,6 +344,11 @@ export interface Settings { paintByExampleSeedFixed: boolean paintByExampleMaskBlur: number paintByExampleMatchHistograms: boolean + + // InstructPix2Pix + p2pSteps: number + p2pImageGuidanceScale: number + p2pGuidanceScale: number } const defaultHDSettings: ModelsHDSettings = { @@ -402,6 +408,13 @@ const defaultHDSettings: ModelsHDSettings = { hdStrategyCropMargin: 128, enabled: false, }, + [AIModel.PIX2PIX]: { + hdStrategy: HDStrategy.ORIGINAL, + hdStrategyResizeLimit: 768, + hdStrategyCropTrigerSize: 512, + hdStrategyCropMargin: 128, + enabled: false, + }, [AIModel.Mange]: { hdStrategy: HDStrategy.CROP, hdStrategyResizeLimit: 1280, @@ -471,6 +484,11 @@ export const settingStateDefault: Settings = { paintByExampleMaskBlur: 5, paintByExampleSeedFixed: false, paintByExampleMatchHistograms: false, + + // InstructPix2Pix + p2pSteps: 50, + p2pImageGuidanceScale: 1.5, + p2pGuidanceScale: 7.5, } const localStorageEffect = @@ -567,12 +585,33 @@ export const isPaintByExampleState = selector({ }, }) +export const isPix2PixState = selector({ + key: 'isPix2PixState', + get: ({ get }) => { + const settings = get(settingState) + return settings.model === AIModel.PIX2PIX + }, +}) + export const runManuallyState = selector({ key: 'runManuallyState', get: ({ get }) => { const settings = get(settingState) const isSD = get(isSDState) const isPaintByExample = get(isPaintByExampleState) - return settings.runInpaintingManually || isSD || isPaintByExample + const isPix2Pix = get(isPix2PixState) + return ( + settings.runInpaintingManually || isSD || isPaintByExample || isPix2Pix + ) + }, +}) + +export const isDiffusionModelsState = selector({ + key: 'isDiffusionModelsState', + get: ({ get }) => { + const isSD = get(isSDState) + const isPaintByExample = get(isPaintByExampleState) + const isPix2Pix = get(isPix2PixState) + return isSD || isPaintByExample || isPix2Pix }, }) diff --git a/lama_cleaner/const.py b/lama_cleaner/const.py index 06245fb..48ca4c2 100644 --- a/lama_cleaner/const.py +++ b/lama_cleaner/const.py @@ -11,7 +11,8 @@ AVAILABLE_MODELS = [ "cv2", "manga", "sd2", - "paint_by_example" + "paint_by_example", + "pix2pix", ] AVAILABLE_DEVICES = ["cuda", "cpu", "mps"] diff --git a/lama_cleaner/model/pix2pix.py b/lama_cleaner/model/pix2pix.py new file mode 100644 index 0000000..ddcd5a4 --- /dev/null +++ b/lama_cleaner/model/pix2pix.py @@ -0,0 +1,83 @@ +import PIL.Image +import cv2 +import torch +from loguru import logger + +from lama_cleaner.model.base import DiffusionInpaintModel +from lama_cleaner.model.utils import set_seed +from lama_cleaner.schema import Config + + +class Pix2Pix(DiffusionInpaintModel): + pad_mod = 8 + min_size = 512 + + def init_model(self, device: torch.device, **kwargs): + from diffusers import StableDiffusionInstructPix2PixPipeline + fp16 = not kwargs.get('no_half', False) + + model_kwargs = {"local_files_only": kwargs.get('local_files_only', kwargs['sd_run_local'])} + if kwargs['disable_nsfw'] or kwargs.get('cpu_offload', False): + logger.info("Disable Stable Diffusion Model NSFW checker") + model_kwargs.update(dict( + safety_checker=None, + feature_extractor=None, + requires_safety_checker=False + )) + + use_gpu = device == torch.device('cuda') and torch.cuda.is_available() + torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32 + self.model = StableDiffusionInstructPix2PixPipeline.from_pretrained( + "timbrooks/instruct-pix2pix", + revision="fp16" if use_gpu and fp16 else "main", + torch_dtype=torch_dtype, + **model_kwargs + ) + + self.model.enable_attention_slicing() + if kwargs.get('enable_xformers', False): + self.model.enable_xformers_memory_efficient_attention() + + if kwargs.get('cpu_offload', False) and use_gpu: + logger.info("Enable sequential cpu offload") + self.model.enable_sequential_cpu_offload(gpu_id=0) + else: + self.model = self.model.to(device) + + def forward(self, image, mask, config: Config): + """Input image and output image have same size + image: [H, W, C] RGB + mask: [H, W, 1] 255 means area to repaint + return: BGR IMAGE + edit = pipe(prompt, image=image, num_inference_steps=20, image_guidance_scale=1.5, guidance_scale=7).images[0] + """ + set_seed(config.sd_seed) + + output = self.model( + image=PIL.Image.fromarray(image), + prompt=config.prompt, + negative_prompt=config.negative_prompt, + num_inference_steps=config.p2p_steps, + image_guidance_scale=config.p2p_image_guidance_scale, + guidance_scale=config.p2p_guidance_scale, + output_type="np.array", + ).images[0] + + output = (output * 255).round().astype("uint8") + output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR) + return output + + # + # def forward_post_process(self, result, image, mask, config): + # if config.sd_match_histograms: + # result = self._match_histograms(result, image[:, :, ::-1], mask) + # + # if config.sd_mask_blur != 0: + # k = 2 * config.sd_mask_blur + 1 + # mask = cv2.GaussianBlur(mask, (k, k), 0) + # return result, image, mask + + @staticmethod + def is_downloaded() -> bool: + # model will be downloaded when app start, and can't switch in frontend settings + return True diff --git a/lama_cleaner/model_manager.py b/lama_cleaner/model_manager.py index ab31b1b..554a038 100644 --- a/lama_cleaner/model_manager.py +++ b/lama_cleaner/model_manager.py @@ -7,13 +7,14 @@ from lama_cleaner.model.ldm import LDM from lama_cleaner.model.manga import Manga from lama_cleaner.model.mat import MAT from lama_cleaner.model.paint_by_example import PaintByExample +from lama_cleaner.model.pix2pix import Pix2Pix from lama_cleaner.model.sd import SD15, SD2 from lama_cleaner.model.zits import ZITS from lama_cleaner.model.opencv2 import OpenCV2 from lama_cleaner.schema import Config models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.5": SD15, "cv2": OpenCV2, "manga": Manga, - "sd2": SD2, "paint_by_example": PaintByExample} + "sd2": SD2, "paint_by_example": PaintByExample, "pix2pix": Pix2Pix} class ModelManager: diff --git a/lama_cleaner/schema.py b/lama_cleaner/schema.py index ed2a3e1..96edf60 100644 --- a/lama_cleaner/schema.py +++ b/lama_cleaner/schema.py @@ -88,3 +88,8 @@ class Config(BaseModel): paint_by_example_seed: int = 42 paint_by_example_match_histograms: bool = False paint_by_example_example_image: Image = None + + # InstructPix2Pix + p2p_steps: int = 50 + p2p_image_guidance_scale: float = 1.5 + p2p_guidance_scale: float = 7.5 diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index 2f7be75..a55e1df 100644 --- a/lama_cleaner/server.py +++ b/lama_cleaner/server.py @@ -228,6 +228,9 @@ def process(): paint_by_example_seed=form["paintByExampleSeed"], paint_by_example_match_histograms=form["paintByExampleMatchHistograms"], paint_by_example_example_image=paint_by_example_example_image, + p2p_steps=form["p2pSteps"], + p2p_image_guidance_scale=form["p2pImageGuidanceScale"], + p2p_guidance_scale=form["p2pGuidanceScale"], ) if config.sd_seed == -1: diff --git a/requirements.txt b/requirements.txt index e8bb98d..bf97a7b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,7 +11,7 @@ pytest yacs markupsafe==2.0.1 scikit-image==0.19.3 -diffusers[torch]==0.10.2 +diffusers[torch]==0.12.1 transformers>=4.25.1 watchdog==2.2.1 gradio \ No newline at end of file