kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

add pix2pix

Browse Source
This commit is contained in:
Qing 2023-01-28 21:13:21 +08:00
parent 205170e1e5
commit 05e82598de
14 changed files with 389 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lama_cleaner/app/src/adapters/inpainting.ts
View File

@ -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 {

13
lama_cleaner/app/src/components/Editor/Editor.tsx
View File

@ -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 ? <InteractiveSeg /> : <></>}
@ -1442,7 +1441,7 @@ export default function Editor() {
)}
<div className="editor-toolkit-panel">
{isSD || isPaintByExample || file === undefined ? (
{isDiffusionModels || file === undefined ? (
<></>
) : (
<SizeSelector
@ -1545,7 +1544,7 @@ export default function Editor() {
onClick={download}
/>
{settings.runInpaintingManually && !isSD && !isPaintByExample && (
{settings.runInpaintingManually && !isDiffusionModels && (
<Button
toolTip="Run Inpainting"
icon={

4
lama_cleaner/app/src/components/Header/Header.tsx
View File

@ -7,6 +7,7 @@ import {
enableFileManagerState,
fileState,
isInpaintingState,
isPix2PixState,
isSDState,
maskState,
runManuallyState,
@ -30,6 +31,7 @@ const Header = () => {
const [uploadElemId] = useState(`file-upload-${Math.random().toString()}`)
const [maskUploadElemId] = useState(`mask-upload-${Math.random().toString()}`)
const isSD = useRecoilValue(isSDState)
const isPix2Pix = useRecoilValue(isPix2PixState)
const runManually = useRecoilValue(runManuallyState)
const [openMaskPopover, setOpenMaskPopover] = useState(false)
const [showFileManager, setShowFileManager] =
@ -172,7 +174,7 @@ const Header = () => {
</div>
</div>
{isSD && file ? <PromptInput /> : <></>}
{(isSD || isPix2Pix) && file ? <PromptInput /> : <></>}
<div className="header-icons-wrapper">
<CoffeeIcon />

20
lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
View File

@ -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 <></>
}

7
lama_cleaner/app/src/components/Settings/SettingsModal.tsx
View File

@ -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}
>
<DownloadMaskSettingBlock />
{isSD || isPaintByExample ? <></> : <ManualRunInpaintingSettingBlock />}
{isDiffusionModels ? <></> : <ManualRunInpaintingSettingBlock />}
<ModelSettingBlock />
{isSD ? <></> : <HDSettingBlock />}
{isDiffusionModels ? <></> : <HDSettingBlock />}
</Modal>
)
}

224
lama_cleaner/app/src/components/SidePanel/P2PSidePanel.tsx Normal file
View File

@ -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<HTMLTextAreaElement>) => {
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 (
<div className="side-panel">
<PopoverPrimitive.Root open={open}>
<PopoverPrimitive.Trigger
className="btn-primary side-panel-trigger"
onClick={() => toggleOpen()}
>
Config
</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>
}
/>
<ImageResizeScale />
<NumberInputSetting
title="Steps"
width={INPUT_WIDTH}
value={`${setting.p2pSteps}`}
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, p2pSteps: val }
})
}}
/>
<NumberInputSetting
title="Guidance Scale"
width={INPUT_WIDTH}
allowFloat
value={`${setting.p2pGuidanceScale}`}
desc="Higher guidance scale encourages to generate images that are closely linked to the text prompt, usually at the expense of lower image quality."
onValue={value => {
const val = value.length === 0 ? 0 : parseFloat(value)
setSettingState(old => {
return { ...old, p2pGuidanceScale: val }
})
}}
/>
<NumberInputSetting
title="Image Guidance Scale"
width={INPUT_WIDTH}
allowFloat
value={`${setting.p2pImageGuidanceScale}`}
desc=""
onValue={value => {
const val = value.length === 0 ? 0 : parseFloat(value)
setSettingState(old => {
return { ...old, p2pImageGuidanceScale: val }
})
}}
/>
{/* <NumberInputSetting
title="Mask Blur"
width={INPUT_WIDTH}
value={`${setting.sdMaskBlur}`}
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, sdMaskBlur: 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.sdMatchHistograms}
onCheckedChange={value => {
setSettingState(old => {
return { ...old, sdMatchHistograms: value }
})
}}
>
<SwitchThumb />
</Switch>
}
/> */}
{/* <SettingBlock
className="sub-setting-block"
title="Sampler"
input={
<Selector
width={80}
value={setting.sdSampler as string}
options={Object.values(SDSampler)}
onChange={val => {
const sampler = val as SDSampler
setSettingState(old => {
return { ...old, sdSampler: sampler }
})
}}
/>
}
/> */}
<SettingBlock
title="Seed"
input={
<div
style={{
display: 'flex',
gap: 0,
justifyContent: 'center',
alignItems: 'center',
}}
>
<NumberInputSetting
title=""
width={80}
value={`${setting.sdSeed}`}
desc=""
disable={!setting.sdSeedFixed}
onValue={value => {
const val = value.length === 0 ? 0 : parseInt(value, 10)
setSettingState(old => {
return { ...old, sdSeed: val }
})
}}
/>
<Switch
checked={setting.sdSeedFixed}
onCheckedChange={value => {
setSettingState(old => {
return { ...old, sdSeedFixed: value }
})
}}
style={{ marginLeft: '8px' }}
>
<SwitchThumb />
</Switch>
</div>
}
/>
<SettingBlock
className="sub-setting-block"
title="Negative prompt"
layout="v"
input={
<TextAreaInput
className="negative-prompt"
value={negativePrompt}
onInput={handleOnInput}
onKeyUp={onKeyUp}
placeholder=""
/>
}
/>
</PopoverPrimitive.Content>
</PopoverPrimitive.Portal>
</PopoverPrimitive.Root>
</div>
)
}
export default P2PSidePanel

4
lama_cleaner/app/src/components/Workspace.tsx
View File

@ -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 ? <SidePanel /> : <></>}
{isPaintByExample ? <PESidePanel /> : <></>}
{isPix2Pix ? <P2PSidePanel /> : <></>}
<FileManager
photoWidth={256}
show={showFileManager}

41
lama_cleaner/app/src/store/Atoms.tsx
View File

@ -14,6 +14,7 @@ export enum AIModel {
CV2 = 'cv2',
Mange = 'manga',
PAINT_BY_EXAMPLE = 'paint_by_example',
PIX2PIX = 'pix2pix',
}
export const maskState = atom<File | undefined>({
@ -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
},
})

3
lama_cleaner/const.py
View File

@ -11,7 +11,8 @@ AVAILABLE_MODELS = [
"cv2",
"manga",
"sd2",
"paint_by_example"
"paint_by_example",
"pix2pix",
]
AVAILABLE_DEVICES = ["cuda", "cpu", "mps"]

83
lama_cleaner/model/pix2pix.py Normal file
View File

@ -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

3
lama_cleaner/model_manager.py
View File

@ -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:

5
lama_cleaner/schema.py
View File

@ -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

3
lama_cleaner/server.py
View File

@ -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:

2
requirements.txt
View File

@ -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