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

Merge branch 'make_gif_share'

Browse Source
This commit is contained in:
Qing 2023-02-07 21:47:11 +08:00
commit f837e4be8a
49 changed files with 15954 additions and 12250 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -7,3 +7,4 @@ build
!lama_cleaner/app/build
dist/
lama_cleaner.egg-info/
venv/

6
README.md
View File

@ -59,6 +59,6 @@ Only needed if you plan to modify the frontend and recompile yourself.
Frontend code are modified from [cleanup.pictures](https://github.com/initml/cleanup.pictures), You can experience their
great online services [here](https://cleanup.pictures/).
- Install dependencies:`cd lama_cleaner/app/ && yarn`
- Start development server: `yarn start`
- Build: `yarn build`
- Install dependencies:`cd lama_cleaner/app/ && pnpm install`
- Start development server: `pnpm start`
- Build: `pnpm build`

9
lama_cleaner/app/build/asset-manifest.json
View File

@ -1,7 +1,8 @@
{
"files": {
"main.css": "/static/css/main.c28d98ca.css",
"main.js": "/static/js/main.1bd455bc.js",
"main.css": "/static/css/main.e24c9a9b.css",
"main.js": "/static/js/main.23732b19.js",
"static/media/coffee-machine-lineal.gif": "/static/media/coffee-machine-lineal.ee32631219cc3986f861.gif",
"static/media/WorkSans-SemiBold.ttf": "/static/media/WorkSans-SemiBold.1e98db4eb705b586728e.ttf",
"static/media/WorkSans-Bold.ttf": "/static/media/WorkSans-Bold.2bea7a7f7d052c74da25.ttf",
"static/media/WorkSans-Regular.ttf": "/static/media/WorkSans-Regular.bb287b894b27372d8ea7.ttf",
@ -9,7 +10,7 @@
"index.html": "/index.html"
},
"entrypoints": [
"static/css/main.c28d98ca.css",
"static/js/main.1bd455bc.js"
"static/css/main.e24c9a9b.css",
"static/js/main.23732b19.js"
]
}

2
lama_cleaner/app/build/index.html
View File

@ -1 +1 @@
<!doctype html><html lang="en"><head><meta http-equiv="Cache-Control" content="no-cache, no-store, must-revalidate"/><meta http-equiv="Pragma" content="no-cache"/><meta http-equiv="Expires" content="0"/><meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1,maximum-scale=1,user-scalable=0"/><meta name="theme-color" content="#ffffff"/><title>lama-cleaner - Image inpainting powered by SOTA AI model</title><script defer="defer" src="/static/js/main.1bd455bc.js"></script><link href="/static/css/main.c28d98ca.css" rel="stylesheet"></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root"></div></body></html>
<!doctype html><html lang="en"><head><meta http-equiv="Cache-Control" content="no-cache, no-store, must-revalidate"/><meta http-equiv="Pragma" content="no-cache"/><meta http-equiv="Expires" content="0"/><meta charset="utf-8"/><meta name="viewport" content="width=device-width,initial-scale=1,maximum-scale=1,user-scalable=0"/><meta name="theme-color" content="#ffffff"/><title>lama-cleaner - Image inpainting powered by SOTA AI model</title><script defer="defer" src="/static/js/main.23732b19.js"></script><link href="/static/css/main.e24c9a9b.css" rel="stylesheet"></head><body><noscript>You need to enable JavaScript to run this app.</noscript><div id="root"></div></body></html>

1
lama_cleaner/app/build/static/css/main.c28d98ca.css
View File

File diff suppressed because one or more lines are too long

1
lama_cleaner/app/build/static/css/main.e24c9a9b.css Normal file
View File

File diff suppressed because one or more lines are too long

2
lama_cleaner/app/build/static/js/main.1bd455bc.js
View File

File diff suppressed because one or more lines are too long

2
lama_cleaner/app/build/static/js/main.23732b19.js Normal file
View File

File diff suppressed because one or more lines are too long

0
lama_cleaner/app/build/static/js/main.1bd455bc.js.LICENSE.txt → lama_cleaner/app/build/static/js/main.23732b19.js.LICENSE.txt
View File

BIN
lama_cleaner/app/build/static/media/coffee-machine-lineal.ee32631219cc3986f861.gif Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 422 KiB

2
lama_cleaner/app/package.json
View File

@ -2,7 +2,7 @@
"name": "lama-cleaner",
"version": "0.1.0",
"private": true,
"proxy": "http://localhost:8080",
"proxy": "http://127.0.0.1:8080",
"dependencies": {
"@babel/core": "^7.16.0",
"@heroicons/react": "^2.0.0",

14605
lama_cleaner/app/pnpm-lock.yaml Normal file
View File

File diff suppressed because it is too large Load Diff

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

@ -1,5 +1,5 @@
import { Rect, Settings } from '../store/Atoms'
import { dataURItoBlob, srcToFile } from '../utils'
import { dataURItoBlob, loadImage, srcToFile } from '../utils'
export const API_ENDPOINT = `${process.env.REACT_APP_INPAINTING_URL}`
@ -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 {
@ -223,3 +228,34 @@ export async function downloadToOutput(
throw new Error(`Something went wrong: ${error}`)
}
}
export async function makeGif(
originFile: File,
cleanImage: HTMLImageElement,
filename: string,
mimeType: string
) {
const cleanFile = await srcToFile(cleanImage.src, filename, mimeType)
const fd = new FormData()
fd.append('origin_img', originFile)
fd.append('clean_img', cleanFile)
fd.append('filename', filename)
try {
const res = await fetch(`${API_ENDPOINT}/make_gif`, {
method: 'POST',
body: fd,
})
if (!res.ok) {
const errMsg = await res.text()
throw new Error(errMsg)
}
const blob = await res.blob()
const newImage = new Image()
await loadImage(newImage, URL.createObjectURL(blob))
return newImage
} catch (error) {
throw new Error(`Something went wrong: ${error}`)
}
}

22
lama_cleaner/app/src/components/CoffeeIcon/CoffeeIcon.tsx
View File

@ -2,6 +2,7 @@ import React, { useState } from 'react'
import { Coffee } from 'react-feather'
import Button from '../shared/Button'
import Modal from '../shared/Modal'
import CoffeeMachineGif from '../../media/coffee-machine-lineal.gif'
const CoffeeIcon = () => {
const [show, setShow] = useState(false)
@ -23,11 +24,27 @@ const CoffeeIcon = () => {
className="modal-setting"
show={show}
showCloseIcon={false}
>
<div
style={{
display: 'flex',
flexDirection: 'column',
}}
>
<h4 style={{ lineHeight: '24px' }}>
Hi there, If you found my project is useful, and want to help keep it
alive please consider donating! Thank you for your support!
Hi, if you found my project is useful, please conside buy me a
coffee to support my work. Thanks!
</h4>
<img
src={CoffeeMachineGif}
alt="coffee machine"
style={{
height: 150,
objectFit: 'contain',
}}
/>
</div>
<div
style={{
display: 'flex',
@ -53,7 +70,6 @@ const CoffeeIcon = () => {
}}
>
Sure
<Coffee />
</div>
</Button>
</a>

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

@ -55,10 +55,9 @@
position: fixed;
bottom: 0.5rem;
border-radius: 3rem;
padding: 0.6rem 3rem;
display: grid;
grid-template-areas: 'toolkit-size-selector toolkit-brush-slider toolkit-btns';
column-gap: 2rem;
padding: 0.6rem 32px;
display: flex;
gap: 16px;
align-items: center;
justify-content: center;
backdrop-filter: blur(12px);
@ -96,10 +95,8 @@
}
.editor-toolkit-btns {
grid-area: toolkit-btns;
display: grid;
grid-auto-flow: column;
column-gap: 1rem;
display: flex;
gap: 12px;
}
.brush-shape {

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

@ -16,10 +16,11 @@ import {
TransformComponent,
TransformWrapper,
} from 'react-zoom-pan-pinch'
import { useRecoilState, useRecoilValue } from 'recoil'
import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
import { useWindowSize, useKey, useKeyPressEvent } from 'react-use'
import inpaint, {
downloadToOutput,
makeGif,
postInteractiveSeg,
} from '../../adapters/inpainting'
import Button from '../shared/Button'
@ -43,10 +44,12 @@ import {
imageHeightState,
imageWidthState,
interactiveSegClicksState,
isDiffusionModelsState,
isInpaintingState,
isInteractiveSegRunningState,
isInteractiveSegState,
isPaintByExampleState,
isPix2PixState,
isSDState,
negativePropmtState,
propmtState,
@ -66,6 +69,7 @@ import FileSelect from '../FileSelect/FileSelect'
import InteractiveSeg from '../InteractiveSeg/InteractiveSeg'
import InteractiveSegConfirmActions from '../InteractiveSeg/ConfirmActions'
import InteractiveSegReplaceModal from '../InteractiveSeg/ReplaceModal'
import MakeGIF from './MakeGIF'
const TOOLBAR_SIZE = 200
const MIN_BRUSH_SIZE = 10
@ -112,11 +116,11 @@ export default function Editor() {
const settings = useRecoilValue(settingState)
const [seedVal, setSeed] = useRecoilState(seedState)
const croperRect = useRecoilValue(croperState)
const [toastVal, setToastState] = useRecoilState(toastState)
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 isPix2Pix = useRecoilValue(isPix2PixState)
const [isInteractiveSeg, setIsInteractiveSeg] = useRecoilState(
isInteractiveSegState
)
@ -181,8 +185,8 @@ export default function Editor() {
const [redoLineGroups, setRedoLineGroups] = useState<LineGroup[]>([])
const enableFileManager = useRecoilValue(enableFileManagerState)
const [imageWidth, setImageWidth] = useRecoilState(imageWidthState)
const [imageHeight, setImageHeight] = useRecoilState(imageHeightState)
const setImageWidth = useSetRecoilState(imageWidthState)
const setImageHeight = useSetRecoilState(imageHeightState)
const app = useRecoilValue(appState)
const draw = useCallback(
@ -253,13 +257,40 @@ export default function Editor() {
_lineGroups.forEach(lineGroup => {
drawLines(ctx, lineGroup, 'white')
})
if (
(maskImage === undefined || maskImage === null) &&
_lineGroups.length === 1 &&
_lineGroups[0].length === 0 &&
isPix2Pix
) {
// For InstructPix2Pix without mask
drawLines(
ctx,
[
{
size: 9999999999,
pts: [
{ x: 0, y: 0 },
{ x: original.naturalWidth, y: 0 },
{ x: original.naturalWidth, y: original.naturalHeight },
{ x: 0, y: original.naturalHeight },
],
},
[context, maskCanvas]
],
'white'
)
}
},
[context, maskCanvas, isPix2Pix]
)
const hadDrawSomething = useCallback(() => {
if (isPix2Pix) {
return true
}
return curLineGroup.length !== 0
}, [curLineGroup])
}, [curLineGroup, isPix2Pix])
const drawOnCurrentRender = useCallback(
(lineGroup: LineGroup) => {
@ -424,6 +455,8 @@ export default function Editor() {
} else if (prevInteractiveSegMask) {
// 使用上一次 IS 的 mask 生成
runInpainting(false, undefined, prevInteractiveSegMask)
} else if (isPix2Pix) {
runInpainting(false, undefined, null)
} else {
setToastState({
open: true,
@ -839,7 +872,7 @@ export default function Editor() {
}
if (
(isSD || isPaintByExample) &&
isDiffusionModels &&
settings.showCroper &&
isOutsideCroper(mouseXY(ev))
) {
@ -1385,7 +1418,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 /> : <></>}
@ -1439,7 +1472,7 @@ export default function Editor() {
)}
<div className="editor-toolkit-panel">
{isSD || isPaintByExample || file === undefined ? (
{isDiffusionModels || file === undefined ? (
<></>
) : (
<SizeSelector
@ -1534,6 +1567,7 @@ export default function Editor() {
}}
disabled={renders.length === 0}
/>
<MakeGIF renders={renders} />
<Button
toolTip="Save Image"
icon={<ArrowDownTrayIcon />}
@ -1541,7 +1575,7 @@ export default function Editor() {
onClick={download}
/>
{settings.runInpaintingManually && !isSD && !isPaintByExample && (
{settings.runInpaintingManually && !isDiffusionModels && (
<Button
toolTip="Run Inpainting"
icon={

114
lama_cleaner/app/src/components/Editor/MakeGIF.tsx Normal file
View File

@ -0,0 +1,114 @@
import React, { useState } from 'react'
import { GifIcon } from '@heroicons/react/24/outline'
import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
import Button from '../shared/Button'
import { fileState, gifImageState, toastState } from '../../store/Atoms'
import { makeGif } from '../../adapters/inpainting'
import Modal from '../shared/Modal'
import { LoadingIcon } from '../shared/Toast'
import { downloadImage } from '../../utils'
interface Props {
renders: HTMLImageElement[]
}
const MakeGIF = (props: Props) => {
const { renders } = props
const [gifImg, setGifImg] = useRecoilState(gifImageState)
const file = useRecoilValue(fileState)
const setToastState = useSetRecoilState(toastState)
const [show, setShow] = useState(false)
const handleOnClose = () => {
setShow(false)
}
const handleDownload = () => {
if (gifImg) {
const name = file.name.replace(/\.[^/.]+$/, '.gif')
downloadImage(gifImg.src, name)
}
}
return (
<div>
<Button
toolTip="Make Gif"
icon={<GifIcon />}
disabled={!renders.length}
onClick={async () => {
setShow(true)
setGifImg(null)
try {
const gif = await makeGif(
file,
renders[renders.length - 1],
file.name,
file.type
)
if (gif) {
setGifImg(gif)
}
} catch (e: any) {
setToastState({
open: true,
desc: e.message ? e.message : e.toString(),
state: 'error',
duration: 2000,
})
}
}}
/>
<Modal
onClose={handleOnClose}
title="GIF"
className="modal-setting"
show={show}
>
<div
style={{
display: 'flex',
alignItems: 'center',
justifyContent: 'center',
flexDirection: 'column',
gap: 16,
}}
>
{gifImg ? (
<img src={gifImg.src} style={{ borderRadius: 8 }} alt="gif" />
) : (
<div
style={{
display: 'flex',
alignItems: 'center',
justifyContent: 'center',
gap: 8,
}}
>
<LoadingIcon />
Generating GIF...
</div>
)}
{gifImg && (
<div
style={{
display: 'flex',
width: '100%',
justifyContent: 'flex-end',
alignItems: 'center',
gap: '12px',
}}
>
<Button onClick={handleDownload} border>
Download
</Button>
</div>
)}
</div>
</Modal>
</div>
)
}
export default MakeGIF

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

25
lama_cleaner/app/src/components/Settings/SettingIcon.tsx
View File

@ -1,5 +1,6 @@
import React from 'react'
import { useRecoilState } from 'recoil'
import { Cog6ToothIcon } from '@heroicons/react/24/outline'
import { settingState } from '../../store/Atoms'
import Button from '../shared/Button'
@ -16,29 +17,7 @@ const SettingIcon = () => {
onClick={onClick}
toolTip="Settings"
style={{ border: 0 }}
icon={
<svg
xmlns="http://www.w3.org/2000/svg"
fill="none"
role="img"
width="28"
height="28"
viewBox="0 0 24 24"
stroke="currentColor"
strokeWidth="2"
>
<path
strokeLinecap="round"
strokeLinejoin="round"
d="M10.325 4.317c.426-1.756 2.924-1.756 3.35 0a1.724 1.724 0 002.573 1.066c1.543-.94 3.31.826 2.37 2.37a1.724 1.724 0 001.065 2.572c1.756.426 1.756 2.924 0 3.35a1.724 1.724 0 00-1.066 2.573c.94 1.543-.826 3.31-2.37 2.37a1.724 1.724 0 00-2.572 1.065c-.426 1.756-2.924 1.756-3.35 0a1.724 1.724 0 00-2.573-1.066c-1.543.94-3.31-.826-2.37-2.37a1.724 1.724 0 00-1.065-2.572c-1.756-.426-1.756-2.924 0-3.35a1.724 1.724 0 001.066-2.573c-.94-1.543.826-3.31 2.37-2.37.996.608 2.296.07 2.572-1.065z"
/>
<path
strokeLinecap="round"
strokeLinejoin="round"
d="M15 12a3 3 0 11-6 0 3 3 0 016 0z"
/>
</svg>
}
icon={<Cog6ToothIcon />}
/>
</div>
)

1
lama_cleaner/app/src/components/Settings/Settings.scss
View File

@ -8,7 +8,6 @@
color: var(--modal-text-color);
box-shadow: 0px 0px 20px rgb(0, 0, 40, 0.2);
width: 680px;
min-height: 420px;
@include mobile {
display: grid;

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>
)
}

14
lama_cleaner/app/src/components/Shortcuts/Shortcuts.scss
View File

@ -15,13 +15,14 @@
}
.shortcut-options {
display: grid;
row-gap: 1rem;
display: flex;
gap: 48px;
flex-direction: row;
.shortcut-option {
display: grid;
grid-template-columns: repeat(2, auto);
column-gap: 6rem;
column-gap: 2rem;
align-items: center;
@include mobile {
@ -67,3 +68,10 @@
}
}
}
.shortcut-options-column {
display: flex;
flex-direction: column;
gap: 12px;
width: 320px;
}

35
lama_cleaner/app/src/components/Shortcuts/ShortcutsModal.tsx
View File

@ -50,25 +50,38 @@ export default function ShortcutsModal() {
show={shortcutsShow}
>
<div className="shortcut-options">
<ShortCut
content="Multi-Stroke Mask Drawing"
keys={[`Hold ${CmdOrCtrl}`]}
/>
<ShortCut content="Cancel Mask Drawing" keys={['Esc']} />
<ShortCut content="Run Inpainting Manually" keys={['Shift', 'R']} />
<ShortCut content="Interactive Segmentation" keys={['I']} />
<ShortCut content="Undo Inpainting" keys={[CmdOrCtrl, 'Z']} />
<ShortCut content="Redo Inpainting" keys={[CmdOrCtrl, 'Shift', 'Z']} />
<ShortCut content="View Original Image" keys={['Hold Tab']} />
<div className="shortcut-options-column">
<ShortCut content="Pan" keys={['Space + Drag']} />
<ShortCut content="Reset Zoom/Pan" keys={['Esc']} />
<ShortCut content="Decrease Brush Size" keys={['[']} />
<ShortCut content="Increase Brush Size" keys={[']']} />
<ShortCut content="Toggle Dark Mode" keys={['Shift', 'D']} />
<ShortCut content="View Original Image" keys={['Hold Tab']} />
<ShortCut
content="Multi-Stroke Drawing"
keys={[`Hold ${CmdOrCtrl}`]}
/>
<ShortCut content="Cancel Drawing" keys={['Esc']} />
</div>
<div className="shortcut-options-column">
<ShortCut content="Undo" keys={[CmdOrCtrl, 'Z']} />
<ShortCut content="Redo" keys={[CmdOrCtrl, 'Shift', 'Z']} />
<ShortCut content="Copy Result" keys={[CmdOrCtrl, 'C']} />
<ShortCut content="Paste Image" keys={[CmdOrCtrl, 'V']} />
<ShortCut
content="Trigger Manually Inpainting"
keys={['Shift', 'R']}
/>
<ShortCut content="Trigger Interactive Segmentation" keys={['I']} />
</div>
<div className="shortcut-options-column">
<ShortCut content="Switch Theme" keys={['Shift', 'D']} />
<ShortCut content="Toggle Hotkeys Dialog" keys={['H']} />
<ShortCut content="Toggle Settings Dialog" keys={['S']} />
<ShortCut content="Toggle File Manager" keys={['F']} />
</div>
</div>
</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}

2
lama_cleaner/app/src/components/shared/Toast.tsx
View File

@ -3,7 +3,7 @@ import * as ToastPrimitive from '@radix-ui/react-toast'
import { ToastProps } from '@radix-ui/react-toast'
import { CheckIcon, ExclamationCircleIcon } from '@heroicons/react/24/outline'
const LoadingIcon = () => {
export const LoadingIcon = () => {
return (
<span className="loading-icon">
<svg

BIN
lama_cleaner/app/src/media/coffee-machine-lineal.gif Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 422 KiB

55
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 = 'instruct_pix2pix',
}
export const maskState = atom<File | undefined>({
@ -45,6 +46,7 @@ interface AppState {
interactiveSegClicks: number[][]
showFileManager: boolean
enableFileManager: boolean
gifImage: HTMLImageElement | undefined
}
export const appState = atom<AppState>({
@ -61,6 +63,7 @@ export const appState = atom<AppState>({
interactiveSegClicks: [],
showFileManager: false,
enableFileManager: false,
gifImage: undefined,
},
})
@ -134,6 +137,18 @@ export const enableFileManagerState = selector({
},
})
export const gifImageState = selector({
key: 'gifImageState',
get: ({ get }) => {
const app = get(appState)
return app.gifImage
},
set: ({ get, set }, newValue: any) => {
const app = get(appState)
set(appState, { ...app, gifImage: newValue })
},
})
export const fileState = selector({
key: 'fileState',
get: ({ get }) => {
@ -329,6 +344,11 @@ export interface Settings {
paintByExampleSeedFixed: boolean
paintByExampleMaskBlur: number
paintByExampleMatchHistograms: boolean
// InstructPix2Pix
p2pSteps: number
p2pImageGuidanceScale: number
p2pGuidanceScale: number
}
const defaultHDSettings: ModelsHDSettings = {
@ -388,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,
@ -457,6 +484,11 @@ export const settingStateDefault: Settings = {
paintByExampleMaskBlur: 5,
paintByExampleSeedFixed: false,
paintByExampleMatchHistograms: false,
// InstructPix2Pix
p2pSteps: 50,
p2pImageGuidanceScale: 1.5,
p2pGuidanceScale: 7.5,
}
const localStorageEffect =
@ -553,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
},
})

11861
lama_cleaner/app/yarn.lock
View File

File diff suppressed because it is too large Load Diff

10
lama_cleaner/const.py
View File

@ -1,5 +1,12 @@
import os
MPS_SUPPORT_MODELS = [
"instruct_pix2pix",
"sd1.5",
"sd2",
"paint_by_example"
]
DEFAULT_MODEL = "lama"
AVAILABLE_MODELS = [
"lama",
@ -11,7 +18,8 @@ AVAILABLE_MODELS = [
"cv2",
"manga",
"sd2",
"paint_by_example"
"paint_by_example",
"instruct_pix2pix",
]
AVAILABLE_DEVICES = ["cuda", "cpu", "mps"]

27
lama_cleaner/helper.py
View File

@ -62,7 +62,7 @@ def load_model(model: torch.nn.Module, url_or_path, device):
model_path = download_model(url_or_path)
try:
state_dict = torch.load(model_path, map_location='cpu')
state_dict = torch.load(model_path, map_location="cpu")
model.load_state_dict(state_dict, strict=True)
model.to(device)
logger.info(f"Load model from: {model_path}")
@ -85,26 +85,43 @@ def numpy_to_bytes(image_numpy: np.ndarray, ext: str) -> bytes:
return image_bytes
def load_img(img_bytes, gray: bool = False):
def pil_to_bytes(pil_img, ext: str, exif=None) -> bytes:
with io.BytesIO() as output:
pil_img.save(output, format=ext, exif=exif, quality=95)
image_bytes = output.getvalue()
return image_bytes
def load_img(img_bytes, gray: bool = False, return_exif: bool = False):
alpha_channel = None
image = Image.open(io.BytesIO(img_bytes))
try:
if return_exif:
exif = image.getexif()
except:
exif = None
logger.error("Failed to extract exif from image")
try:
image = ImageOps.exif_transpose(image)
except:
pass
if gray:
image = image.convert('L')
image = image.convert("L")
np_img = np.array(image)
else:
if image.mode == 'RGBA':
if image.mode == "RGBA":
np_img = np.array(image)
alpha_channel = np_img[:, :, -1]
np_img = cv2.cvtColor(np_img, cv2.COLOR_RGBA2RGB)
else:
image = image.convert('RGB')
image = image.convert("RGB")
np_img = np.array(image)
if return_exif:
return np_img, alpha_channel, exif
return np_img, alpha_channel

125
lama_cleaner/make_gif.py Normal file
View File

@ -0,0 +1,125 @@
import io
import math
from pathlib import Path
from PIL import Image, ImageDraw
def keep_ratio_resize(img, size, resample=Image.BILINEAR):
if img.width > img.height:
w = size
h = int(img.height * size / img.width)
else:
h = size
w = int(img.width * size / img.height)
return img.resize((w, h), resample)
def cubic_bezier(p1, p2, duration: int, frames: int):
"""
Args:
p1:
p2:
duration: Total duration of the curve
frames:
Returns:
"""
x0, y0 = (0, 0)
x1, y1 = p1
x2, y2 = p2
x3, y3 = (1, 1)
def cal_y(t):
return math.pow(1 - t, 3) * y0 + \
3 * math.pow(1 - t, 2) * t * y1 + \
3 * (1 - t) * math.pow(t, 2) * y2 + \
math.pow(t, 3) * y3
def cal_x(t):
return math.pow(1 - t, 3) * x0 + \
3 * math.pow(1 - t, 2) * t * x1 + \
3 * (1 - t) * math.pow(t, 2) * x2 + \
math.pow(t, 3) * x3
res = []
for t in range(0, 1 * frames, duration):
t = t / frames
res.append((cal_x(t), cal_y(t)))
res.append((1, 0))
return res
def make_compare_gif(
clean_img: Image.Image,
src_img: Image.Image,
max_side_length: int = 600,
splitter_width: int = 5,
splitter_color=(255, 203, 0, int(255 * 0.73))
):
if clean_img.size != src_img.size:
clean_img = clean_img.resize(src_img.size, Image.BILINEAR)
duration_per_frame = 20
num_frames = 50
# erase-in-out
cubic_bezier_points = cubic_bezier((0.33, 0), (0.66, 1), 1, num_frames)
cubic_bezier_points.reverse()
max_side_length = min(max_side_length, max(clean_img.size))
src_img = keep_ratio_resize(src_img, max_side_length)
clean_img = keep_ratio_resize(clean_img, max_side_length)
width, height = src_img.size
# Generate images to make Gif from right to left
images = []
for i in range(num_frames):
new_frame = Image.new('RGB', (width, height))
new_frame.paste(clean_img, (0, 0))
left = int(cubic_bezier_points[i][0] * width)
cropped_src_img = src_img.crop((left, 0, width, height))
new_frame.paste(cropped_src_img, (left, 0, width, height))
if i != num_frames - 1:
# draw a yellow splitter on the edge of the cropped image
draw = ImageDraw.Draw(new_frame)
draw.line([(left, 0), (left, height)], width=splitter_width, fill=splitter_color)
images.append(new_frame)
for i in range(10):
images.append(src_img)
cubic_bezier_points.reverse()
# Generate images to make Gif from left to right
for i in range(num_frames):
new_frame = Image.new('RGB', (width, height))
new_frame.paste(src_img, (0, 0))
right = int(cubic_bezier_points[i][0] * width)
cropped_src_img = clean_img.crop((0, 0, right, height))
new_frame.paste(cropped_src_img, (0, 0, right, height))
if i != num_frames - 1:
# draw a yellow splitter on the edge of the cropped image
draw = ImageDraw.Draw(new_frame)
draw.line([(right, 0), (right, height)], width=splitter_width, fill=splitter_color)
images.append(new_frame)
images.append(clean_img)
img_byte_arr = io.BytesIO()
clean_img.save(
img_byte_arr,
format='GIF',
save_all=True,
include_color_table=True,
append_images=images,
optimize=False,
duration=duration_per_frame,
loop=0
)
return img_byte_arr.getvalue()

39
lama_cleaner/model/base.py
View File

@ -245,3 +245,42 @@ class InpaintModel:
crop_img, crop_mask, [l, t, r, b] = self._crop_box(image, mask, box, config)
return self._pad_forward(crop_img, crop_mask, config), [l, t, r, b]
class DiffusionInpaintModel(InpaintModel):
@torch.no_grad()
def __call__(self, image, mask, config: Config):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
return: BGR IMAGE
"""
# boxes = boxes_from_mask(mask)
if config.use_croper:
crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config)
crop_image = self._scaled_pad_forward(crop_img, crop_mask, config)
inpaint_result = image[:, :, ::-1]
inpaint_result[t:b, l:r, :] = crop_image
else:
inpaint_result = self._scaled_pad_forward(image, mask, config)
return inpaint_result
def _scaled_pad_forward(self, image, mask, config: Config):
longer_side_length = int(config.sd_scale * max(image.shape[:2]))
origin_size = image.shape[:2]
downsize_image = resize_max_size(image, size_limit=longer_side_length)
downsize_mask = resize_max_size(mask, size_limit=longer_side_length)
logger.info(
f"Resize image to do sd inpainting: {image.shape} -> {downsize_image.shape}"
)
inpaint_result = self._pad_forward(downsize_image, downsize_mask, config)
# only paste masked area result
inpaint_result = cv2.resize(
inpaint_result,
(origin_size[1], origin_size[0]),
interpolation=cv2.INTER_CUBIC,
)
original_pixel_indices = mask < 127
inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices]
return inpaint_result

83
lama_cleaner/model/instruct_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 InstructPix2Pix(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', False)}
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

51
lama_cleaner/model/paint_by_example.py
View File

@ -1,19 +1,16 @@
import random
import PIL
import PIL.Image
import cv2
import numpy as np
import torch
from diffusers import DiffusionPipeline
from loguru import logger
from lama_cleaner.helper import resize_max_size
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.model.utils import set_seed
from lama_cleaner.schema import Config
class PaintByExample(InpaintModel):
class PaintByExample(DiffusionInpaintModel):
pad_mod = 8
min_size = 512
@ -53,11 +50,7 @@ class PaintByExample(InpaintModel):
mask: [H, W, 1] 255 means area to repaint
return: BGR IMAGE
"""
seed = config.paint_by_example_seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
set_seed(config.paint_by_example_seed)
output = self.model(
image=PIL.Image.fromarray(image),
@ -71,42 +64,6 @@ class PaintByExample(InpaintModel):
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
return output
def _scaled_pad_forward(self, image, mask, config: Config):
longer_side_length = int(config.sd_scale * max(image.shape[:2]))
origin_size = image.shape[:2]
downsize_image = resize_max_size(image, size_limit=longer_side_length)
downsize_mask = resize_max_size(mask, size_limit=longer_side_length)
logger.info(
f"Resize image to do paint_by_example: {image.shape} -> {downsize_image.shape}"
)
inpaint_result = self._pad_forward(downsize_image, downsize_mask, config)
# only paste masked area result
inpaint_result = cv2.resize(
inpaint_result,
(origin_size[1], origin_size[0]),
interpolation=cv2.INTER_CUBIC,
)
original_pixel_indices = mask < 127
inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices]
return inpaint_result
@torch.no_grad()
def __call__(self, image, mask, config: Config):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
return: BGR IMAGE
"""
if config.use_croper:
crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config)
crop_image = self._scaled_pad_forward(crop_img, crop_mask, config)
inpaint_result = image[:, :, ::-1]
inpaint_result[t:b, l:r, :] = crop_image
else:
inpaint_result = self._scaled_pad_forward(image, mask, config)
return inpaint_result
def forward_post_process(self, result, image, mask, config):
if config.paint_by_example_match_histograms:
result = self._match_histograms(result, image[:, :, ::-1], mask)

101
lama_cleaner/model/sd.py
View File

@ -4,20 +4,25 @@ import PIL.Image
import cv2
import numpy as np
import torch
from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, \
EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
from diffusers import (
PNDMScheduler,
DDIMScheduler,
LMSDiscreteScheduler,
EulerDiscreteScheduler,
EulerAncestralDiscreteScheduler,
DPMSolverMultistepScheduler,
)
from loguru import logger
from lama_cleaner.helper import resize_max_size
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.model.utils import torch_gc
from lama_cleaner.model.base import DiffusionInpaintModel
from lama_cleaner.model.utils import torch_gc, set_seed
from lama_cleaner.schema import Config, SDSampler
class CPUTextEncoderWrapper:
def __init__(self, text_encoder, torch_dtype):
self.config = text_encoder.config
self.text_encoder = text_encoder.to(torch.device('cpu'), non_blocking=True)
self.text_encoder = text_encoder.to(torch.device("cpu"), non_blocking=True)
self.text_encoder = self.text_encoder.to(torch.float32, non_blocking=True)
self.torch_dtype = torch_dtype
del text_encoder
@ -25,27 +30,40 @@ class CPUTextEncoderWrapper:
def __call__(self, x, **kwargs):
input_device = x.device
return [self.text_encoder(x.to(self.text_encoder.device), **kwargs)[0].to(input_device).to(self.torch_dtype)]
return [
self.text_encoder(x.to(self.text_encoder.device), **kwargs)[0]
.to(input_device)
.to(self.torch_dtype)
]
@property
def dtype(self):
return self.torch_dtype
class SD(InpaintModel):
class SD(DiffusionInpaintModel):
pad_mod = 8
min_size = 512
def init_model(self, device: torch.device, **kwargs):
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
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):
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(
model_kwargs.update(
dict(
safety_checker=None,
feature_extractor=None,
requires_safety_checker=False
))
requires_safety_checker=False,
)
)
use_gpu = device == torch.device('cuda') and torch.cuda.is_available()
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 = StableDiffusionInpaintPipeline.from_pretrained(
self.model_id_or_path,
@ -58,40 +76,23 @@ class SD(InpaintModel):
# https://huggingface.co/docs/diffusers/v0.7.0/en/api/pipelines/stable_diffusion#diffusers.StableDiffusionInpaintPipeline.enable_attention_slicing
self.model.enable_attention_slicing()
# https://huggingface.co/docs/diffusers/v0.7.0/en/optimization/fp16#memory-efficient-attention
if kwargs.get('enable_xformers', False):
if kwargs.get("enable_xformers", False):
self.model.enable_xformers_memory_efficient_attention()
if kwargs.get('cpu_offload', False) and use_gpu:
if kwargs.get("cpu_offload", False) and use_gpu:
# TODO: gpu_id
logger.info("Enable sequential cpu offload")
self.model.enable_sequential_cpu_offload(gpu_id=0)
else:
self.model = self.model.to(device)
if kwargs['sd_cpu_textencoder']:
if kwargs["sd_cpu_textencoder"]:
logger.info("Run Stable Diffusion TextEncoder on CPU")
self.model.text_encoder = CPUTextEncoderWrapper(self.model.text_encoder, torch_dtype)
self.model.text_encoder = CPUTextEncoderWrapper(
self.model.text_encoder, torch_dtype
)
self.callback = kwargs.pop("callback", None)
def _scaled_pad_forward(self, image, mask, config: Config):
longer_side_length = int(config.sd_scale * max(image.shape[:2]))
origin_size = image.shape[:2]
downsize_image = resize_max_size(image, size_limit=longer_side_length)
downsize_mask = resize_max_size(mask, size_limit=longer_side_length)
logger.info(
f"Resize image to do sd inpainting: {image.shape} -> {downsize_image.shape}"
)
inpaint_result = self._pad_forward(downsize_image, downsize_mask, config)
# only paste masked area result
inpaint_result = cv2.resize(
inpaint_result,
(origin_size[1], origin_size[0]),
interpolation=cv2.INTER_CUBIC,
)
original_pixel_indices = mask < 127
inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices]
return inpaint_result
def forward(self, image, mask, config: Config):
"""Input image and output image have same size
image: [H, W, C] RGB
@ -118,11 +119,7 @@ class SD(InpaintModel):
self.model.scheduler = scheduler
seed = config.sd_seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
set_seed(config.sd_seed)
if config.sd_mask_blur != 0:
k = 2 * config.sd_mask_blur + 1
@ -147,24 +144,6 @@ class SD(InpaintModel):
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
return output
@torch.no_grad()
def __call__(self, image, mask, config: Config):
"""
images: [H, W, C] RGB, not normalized
masks: [H, W]
return: BGR IMAGE
"""
# boxes = boxes_from_mask(mask)
if config.use_croper:
crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config)
crop_image = self._scaled_pad_forward(crop_img, crop_mask, config)
inpaint_result = image[:, :, ::-1]
inpaint_result[t:b, l:r, :] = crop_image
else:
inpaint_result = self._scaled_pad_forward(image, mask, config)
return inpaint_result
def forward_post_process(self, result, image, mask, config):
if config.sd_match_histograms:
result = self._match_histograms(result, image[:, :, ::-1], mask)

8
lama_cleaner/model/utils.py
View File

@ -1,4 +1,5 @@
import math
import random
from typing import Any
import torch
@ -713,3 +714,10 @@ def torch_gc():
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def set_seed(seed: int):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

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.instruct_pix2pix import InstructPix2Pix
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, "instruct_pix2pix": InstructPix2Pix}
class ModelManager:

98
lama_cleaner/parse_args.py
View File

@ -5,9 +5,25 @@ from pathlib import Path
from loguru import logger
from lama_cleaner.const import AVAILABLE_MODELS, NO_HALF_HELP, CPU_OFFLOAD_HELP, DISABLE_NSFW_HELP, \
SD_CPU_TEXTENCODER_HELP, LOCAL_FILES_ONLY_HELP, AVAILABLE_DEVICES, ENABLE_XFORMERS_HELP, MODEL_DIR_HELP, \
OUTPUT_DIR_HELP, INPUT_HELP, GUI_HELP, DEFAULT_DEVICE, NO_GUI_AUTO_CLOSE_HELP, DEFAULT_MODEL_DIR
from lama_cleaner.const import (
AVAILABLE_MODELS,
NO_HALF_HELP,
CPU_OFFLOAD_HELP,
DISABLE_NSFW_HELP,
SD_CPU_TEXTENCODER_HELP,
LOCAL_FILES_ONLY_HELP,
AVAILABLE_DEVICES,
ENABLE_XFORMERS_HELP,
MODEL_DIR_HELP,
OUTPUT_DIR_HELP,
INPUT_HELP,
GUI_HELP,
DEFAULT_DEVICE,
NO_GUI_AUTO_CLOSE_HELP,
DEFAULT_MODEL_DIR,
DEFAULT_MODEL,
MPS_SUPPORT_MODELS,
)
from lama_cleaner.runtime import dump_environment_info
@ -16,22 +32,40 @@ def parse_args():
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", default=8080, type=int)
parser.add_argument("--config-installer", action="store_true",
help="Open config web page, mainly for windows installer")
parser.add_argument("--load-installer-config", action="store_true",
help="Load all cmd args from installer config file")
parser.add_argument("--installer-config", default=None, help="Config file for windows installer")
parser.add_argument(
"--config-installer",
action="store_true",
help="Open config web page, mainly for windows installer",
)
parser.add_argument(
"--load-installer-config",
action="store_true",
help="Load all cmd args from installer config file",
)
parser.add_argument(
"--installer-config", default=None, help="Config file for windows installer"
)
parser.add_argument("--model", default="lama", choices=AVAILABLE_MODELS)
parser.add_argument("--model", default=DEFAULT_MODEL, choices=AVAILABLE_MODELS)
parser.add_argument("--no-half", action="store_true", help=NO_HALF_HELP)
parser.add_argument("--cpu-offload", action="store_true", help=CPU_OFFLOAD_HELP)
parser.add_argument("--disable-nsfw", action="store_true", help=DISABLE_NSFW_HELP)
parser.add_argument("--sd-cpu-textencoder", action="store_true", help=SD_CPU_TEXTENCODER_HELP)
parser.add_argument("--local-files-only", action="store_true", help=LOCAL_FILES_ONLY_HELP)
parser.add_argument("--enable-xformers", action="store_true", help=ENABLE_XFORMERS_HELP)
parser.add_argument("--device", default=DEFAULT_DEVICE, type=str, choices=AVAILABLE_DEVICES)
parser.add_argument(
"--sd-cpu-textencoder", action="store_true", help=SD_CPU_TEXTENCODER_HELP
)
parser.add_argument(
"--local-files-only", action="store_true", help=LOCAL_FILES_ONLY_HELP
)
parser.add_argument(
"--enable-xformers", action="store_true", help=ENABLE_XFORMERS_HELP
)
parser.add_argument(
"--device", default=DEFAULT_DEVICE, type=str, choices=AVAILABLE_DEVICES
)
parser.add_argument("--gui", action="store_true", help=GUI_HELP)
parser.add_argument("--no-gui-auto-close", action="store_true", help=NO_GUI_AUTO_CLOSE_HELP)
parser.add_argument(
"--no-gui-auto-close", action="store_true", help=NO_GUI_AUTO_CLOSE_HELP
)
parser.add_argument(
"--gui-size",
default=[1600, 1000],
@ -41,8 +75,14 @@ def parse_args():
)
parser.add_argument("--input", type=str, default=None, help=INPUT_HELP)
parser.add_argument("--output-dir", type=str, default=None, help=OUTPUT_DIR_HELP)
parser.add_argument("--model-dir", type=str, default=DEFAULT_MODEL_DIR, help=MODEL_DIR_HELP)
parser.add_argument("--disable-model-switch", action="store_true", help="Disable model switch in frontend")
parser.add_argument(
"--model-dir", type=str, default=DEFAULT_MODEL_DIR, help=MODEL_DIR_HELP
)
parser.add_argument(
"--disable-model-switch",
action="store_true",
help="Disable model switch in frontend",
)
parser.add_argument("--debug", action="store_true")
# useless args
@ -64,7 +104,7 @@ def parse_args():
parser.add_argument(
"--sd-enable-xformers",
action="store_true",
help="Enable xFormers optimizations. Requires that xformers package has been installed. See: https://github.com/facebookresearch/xformers"
help="Enable xFormers optimizations. Requires that xformers package has been installed. See: https://github.com/facebookresearch/xformers",
)
args = parser.parse_args()
@ -74,14 +114,18 @@ def parse_args():
if args.config_installer:
if args.installer_config is None:
parser.error(f"args.config_installer==True, must set args.installer_config to store config file")
parser.error(
f"args.config_installer==True, must set args.installer_config to store config file"
)
from lama_cleaner.web_config import main
logger.info(f"Launching installer web config page")
main(args.installer_config)
exit()
if args.load_installer_config:
from lama_cleaner.web_config import load_config
if args.installer_config and not os.path.exists(args.installer_config):
parser.error(f"args.installer_config={args.installer_config} not exists")
@ -93,9 +137,17 @@ def parse_args():
if args.device == "cuda":
import torch
if torch.cuda.is_available() is False:
parser.error(
"torch.cuda.is_available() is False, please use --device cpu or check your pytorch installation")
"torch.cuda.is_available() is False, please use --device cpu or check your pytorch installation"
)
if args.device == "mps":
if args.model not in MPS_SUPPORT_MODELS:
parser.error(
f"mps only support: {MPS_SUPPORT_MODELS}, but got {args.model}"
)
if args.model_dir and args.model_dir is not None:
if os.path.isfile(args.model_dir):
@ -115,7 +167,9 @@ def parse_args():
parser.error(f"invalid --input: {args.input} is not a valid image file")
else:
if args.output_dir is None:
parser.error(f"invalid --input: {args.input} is a directory, --output-dir is required")
parser.error(
f"invalid --input: {args.input} is a directory, --output-dir is required"
)
else:
output_dir = Path(args.output_dir)
if not output_dir.exists():
@ -123,6 +177,8 @@ def parse_args():
output_dir.mkdir(parents=True)
else:
if not output_dir.is_dir():
parser.error(f"invalid --output-dir: {output_dir} is not a directory")
parser.error(
f"invalid --output-dir: {output_dir} is not a directory"
)
return args

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

104
lama_cleaner/server.py
View File

@ -19,6 +19,7 @@ from loguru import logger
from watchdog.events import FileSystemEventHandler
from lama_cleaner.interactive_seg import InteractiveSeg, Click
from lama_cleaner.make_gif import make_compare_gif
from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import Config
from lama_cleaner.file_manager import FileManager
@ -31,7 +32,15 @@ try:
except:
pass
from flask import Flask, request, send_file, cli, make_response, send_from_directory, jsonify
from flask import (
Flask,
request,
send_file,
cli,
make_response,
send_from_directory,
jsonify,
)
# Disable ability for Flask to display warning about using a development server in a production environment.
# https://gist.github.com/jerblack/735b9953ba1ab6234abb43174210d356
@ -42,6 +51,7 @@ from lama_cleaner.helper import (
load_img,
numpy_to_bytes,
resize_max_size,
pil_to_bytes,
)
NUM_THREADS = str(multiprocessing.cpu_count())
@ -93,6 +103,25 @@ def diffuser_callback(i, t, latents):
# socketio.emit('diffusion_step', {'diffusion_step': step})
@app.route("/make_gif", methods=["POST"])
def make_gif():
input = request.files
filename = request.form["filename"]
origin_image_bytes = input["origin_img"].read()
clean_image_bytes = input["clean_img"].read()
origin_image, _ = load_img(origin_image_bytes)
clean_image, _ = load_img(clean_image_bytes)
gif_bytes = make_compare_gif(
Image.fromarray(origin_image), Image.fromarray(clean_image)
)
return send_file(
io.BytesIO(gif_bytes),
mimetype="image/gif",
as_attachment=True,
attachment_filename=filename,
)
@app.route("/save_image", methods=["POST"])
def save_image():
# all image in output directory
@ -100,12 +129,12 @@ def save_image():
origin_image_bytes = input["image"].read() # RGB
image, _ = load_img(origin_image_bytes)
thumb.save_to_output_directory(image, request.form["filename"])
return 'ok', 200
return "ok", 200
@app.route("/medias/<tab>")
def medias(tab):
if tab == 'image':
if tab == "image":
response = make_response(jsonify(thumb.media_names), 200)
else:
response = make_response(jsonify(thumb.output_media_names), 200)
@ -116,18 +145,18 @@ def medias(tab):
return response
@app.route('/media/<tab>/<filename>')
@app.route("/media/<tab>/<filename>")
def media_file(tab, filename):
if tab == 'image':
if tab == "image":
return send_from_directory(thumb.root_directory, filename)
return send_from_directory(thumb.output_dir, filename)
@app.route('/media_thumbnail/<tab>/<filename>')
@app.route("/media_thumbnail/<tab>/<filename>")
def media_thumbnail_file(tab, filename):
args = request.args
width = args.get('width')
height = args.get('height')
width = args.get("width")
height = args.get("height")
if width is None and height is None:
width = 256
if width:
@ -136,9 +165,11 @@ def media_thumbnail_file(tab, filename):
height = int(float(height))
directory = thumb.root_directory
if tab == 'output':
if tab == "output":
directory = thumb.output_dir
thumb_filename, (width, height) = thumb.get_thumbnail(directory, filename, width, height)
thumb_filename, (width, height) = thumb.get_thumbnail(
directory, filename, width, height
)
thumb_filepath = f"{app.config['THUMBNAIL_MEDIA_THUMBNAIL_ROOT']}{thumb_filename}"
response = make_response(send_file(thumb_filepath))
@ -152,13 +183,16 @@ def process():
input = request.files
# RGB
origin_image_bytes = input["image"].read()
image, alpha_channel = load_img(origin_image_bytes)
image, alpha_channel, exif = load_img(origin_image_bytes, return_exif=True)
mask, _ = load_img(input["mask"].read(), gray=True)
mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
if image.shape[:2] != mask.shape[:2]:
return f"Mask shape{mask.shape[:2]} not queal to Image shape{image.shape[:2]}", 400
return (
f"Mask shape{mask.shape[:2]} not queal to Image shape{image.shape[:2]}",
400,
)
original_shape = image.shape
interpolation = cv2.INTER_CUBIC
@ -171,7 +205,9 @@ def process():
size_limit = int(size_limit)
if "paintByExampleImage" in input:
paint_by_example_example_image, _ = load_img(input["paintByExampleImage"].read())
paint_by_example_example_image, _ = load_img(
input["paintByExampleImage"].read()
)
paint_by_example_example_image = Image.fromarray(paint_by_example_example_image)
else:
paint_by_example_example_image = None
@ -200,13 +236,16 @@ def process():
sd_seed=form["sdSeed"],
sd_match_histograms=form["sdMatchHistograms"],
cv2_flag=form["cv2Flag"],
cv2_radius=form['cv2Radius'],
cv2_radius=form["cv2Radius"],
paint_by_example_steps=form["paintByExampleSteps"],
paint_by_example_guidance_scale=form["paintByExampleGuidanceScale"],
paint_by_example_mask_blur=form["paintByExampleMaskBlur"],
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:
@ -235,6 +274,7 @@ def process():
logger.info(f"process time: {(time.time() - start) * 1000}ms")
torch.cuda.empty_cache()
res_np_img = cv2.cvtColor(res_np_img.astype(np.uint8), cv2.COLOR_BGR2RGB)
if alpha_channel is not None:
if alpha_channel.shape[:2] != res_np_img.shape[:2]:
alpha_channel = cv2.resize(
@ -246,9 +286,15 @@ def process():
ext = get_image_ext(origin_image_bytes)
if exif is not None:
bytes_io = io.BytesIO(pil_to_bytes(Image.fromarray(res_np_img), ext, exif=exif))
else:
bytes_io = io.BytesIO(pil_to_bytes(Image.fromarray(res_np_img), ext))
response = make_response(
send_file(
io.BytesIO(numpy_to_bytes(res_np_img, ext)),
# io.BytesIO(numpy_to_bytes(res_np_img, ext)),
bytes_io,
mimetype=f"image/{ext}",
)
)
@ -261,7 +307,7 @@ def interactive_seg():
input = request.files
origin_image_bytes = input["image"].read() # RGB
image, _ = load_img(origin_image_bytes)
if 'mask' in input:
if "mask" in input:
mask, _ = load_img(input["mask"].read(), gray=True)
else:
mask = None
@ -269,14 +315,16 @@ def interactive_seg():
_clicks = json.loads(request.form["clicks"])
clicks = []
for i, click in enumerate(_clicks):
clicks.append(Click(coords=(click[1], click[0]), indx=i, is_positive=click[2] == 1))
clicks.append(
Click(coords=(click[1], click[0]), indx=i, is_positive=click[2] == 1)
)
start = time.time()
new_mask = interactive_seg_model(image, clicks=clicks, prev_mask=mask)
logger.info(f"interactive seg process time: {(time.time() - start) * 1000}ms")
response = make_response(
send_file(
io.BytesIO(numpy_to_bytes(new_mask, 'png')),
io.BytesIO(numpy_to_bytes(new_mask, "png")),
mimetype=f"image/png",
)
)
@ -290,13 +338,13 @@ def current_model():
@app.route("/is_disable_model_switch")
def get_is_disable_model_switch():
res = 'true' if is_disable_model_switch else 'false'
res = "true" if is_disable_model_switch else "false"
return res, 200
@app.route("/is_enable_file_manager")
def get_is_enable_file_manager():
res = 'true' if is_enable_file_manager else 'false'
res = "true" if is_enable_file_manager else "false"
return res, 200
@ -365,14 +413,18 @@ def main(args):
is_disable_model_switch = args.disable_model_switch
is_desktop = args.gui
if is_disable_model_switch:
logger.info(f"Start with --disable-model-switch, model switch on frontend is disable")
logger.info(
f"Start with --disable-model-switch, model switch on frontend is disable"
)
if args.input and os.path.isdir(args.input):
logger.info(f"Initialize file manager")
thumb = FileManager(app)
is_enable_file_manager = True
app.config["THUMBNAIL_MEDIA_ROOT"] = args.input
app.config["THUMBNAIL_MEDIA_THUMBNAIL_ROOT"] = os.path.join(args.output_dir, 'lama_cleaner_thumbnails')
app.config["THUMBNAIL_MEDIA_THUMBNAIL_ROOT"] = os.path.join(
args.output_dir, "lama_cleaner_thumbnails"
)
thumb.output_dir = Path(args.output_dir)
# thumb.start()
# try:
@ -408,8 +460,12 @@ def main(args):
from flaskwebgui import FlaskUI
ui = FlaskUI(
app, width=app_width, height=app_height, host=args.host, port=args.port,
close_server_on_exit=not args.no_gui_auto_close
app,
width=app_width,
height=app_height,
host=args.host,
port=args.port,
close_server_on_exit=not args.no_gui_auto_close,
)
ui.run()
else:

0
lama_cleaner/tests/overture-creations-5sI6fQgYIuo_all_mask.png Normal file
View File

62
lama_cleaner/tests/test_instruct_pix2pix.py Normal file
View File

@ -0,0 +1,62 @@
from pathlib import Path
import pytest
import torch
from lama_cleaner.model_manager import ModelManager
from lama_cleaner.tests.test_model import get_config, assert_equal
from lama_cleaner.schema import HDStrategy
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / 'result'
save_dir.mkdir(exist_ok=True, parents=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
@pytest.mark.parametrize("disable_nsfw", [True, False])
@pytest.mark.parametrize("cpu_offload", [False, True])
def test_instruct_pix2pix(disable_nsfw, cpu_offload):
sd_steps = 50 if device == 'cuda' else 1
model = ModelManager(name="instruct_pix2pix",
device=torch.device(device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=False,
cpu_offload=cpu_offload)
cfg = get_config(strategy=HDStrategy.ORIGINAL, prompt='What if it were snowing?', p2p_steps=sd_steps, sd_scale=1.1)
name = f"device_{device}_disnsfw_{disable_nsfw}_cpu_offload_{cpu_offload}"
assert_equal(
model,
cfg,
f"instruct_pix2pix_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.3
)
@pytest.mark.parametrize("disable_nsfw", [False])
@pytest.mark.parametrize("cpu_offload", [False])
def test_instruct_pix2pix_snow(disable_nsfw, cpu_offload):
sd_steps = 50 if device == 'cuda' else 1
model = ModelManager(name="instruct_pix2pix",
device=torch.device(device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=False,
cpu_offload=cpu_offload)
cfg = get_config(strategy=HDStrategy.ORIGINAL, prompt='What if it were snowing?', p2p_steps=sd_steps)
name = f"snow"
assert_equal(
model,
cfg,
f"instruct_pix2pix_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)

42
lama_cleaner/tests/test_save_exif.py Normal file
View File

@ -0,0 +1,42 @@
import io
from pathlib import Path
from PIL import Image
from lama_cleaner.helper import pil_to_bytes
current_dir = Path(__file__).parent.absolute().resolve()
png_img_p = current_dir / "image.png"
jpg_img_p = current_dir / "bunny.jpeg"
def print_exif(exif):
for k, v in exif.items():
print(f"{k}: {v}")
def test_png():
img = Image.open(png_img_p)
exif = img.getexif()
print_exif(exif)
pil_bytes = pil_to_bytes(img, ext="png", exif=exif)
res_img = Image.open(io.BytesIO(pil_bytes))
res_exif = res_img.getexif()
assert dict(exif) == dict(res_exif)
def test_jpeg():
img = Image.open(jpg_img_p)
exif = img.getexif()
print_exif(exif)
pil_bytes = pil_to_bytes(img, ext="jpeg", exif=exif)
res_img = Image.open(io.BytesIO(pil_bytes))
res_exif = res_img.getexif()
assert dict(exif) == dict(res_exif)

128
lama_cleaner/tests/test_sd_model.py
View File

@ -8,33 +8,37 @@ from lama_cleaner.schema import HDStrategy, SDSampler
from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / 'result'
save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
@pytest.mark.parametrize("sd_device", ['cuda'])
@pytest.mark.parametrize("sd_device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
@pytest.mark.parametrize("cpu_textencoder", [True, False])
@pytest.mark.parametrize("disable_nsfw", [True, False])
def test_runway_sd_1_5_ddim(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
def test_runway_sd_1_5_ddim(
sd_device, strategy, sampler, cpu_textencoder, disable_nsfw
):
def callback(i, t, latents):
print(f"sd_step_{i}")
pass
if sd_device == 'cuda' and not torch.cuda.is_available():
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == 'cuda' else 1
model = ModelManager(name="sd1.5",
sd_steps = 50 if sd_device == "cuda" else 1
model = ModelManager(
name="sd1.5",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=cpu_textencoder,
callback=callback)
cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=sd_steps)
callback=callback,
)
cfg = get_config(strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
@ -45,31 +49,35 @@ def test_runway_sd_1_5_ddim(sd_device, strategy, sampler, cpu_textencoder, disab
f"runway_sd_{strategy.capitalize()}_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.3
fx=1.3,
)
@pytest.mark.parametrize("sd_device", ['cuda'])
@pytest.mark.parametrize("sd_device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.pndm, SDSampler.k_lms, SDSampler.k_euler, SDSampler.k_euler_a])
@pytest.mark.parametrize(
"sampler", [SDSampler.pndm, SDSampler.k_lms, SDSampler.k_euler, SDSampler.k_euler_a]
)
@pytest.mark.parametrize("cpu_textencoder", [False])
@pytest.mark.parametrize("disable_nsfw", [True])
def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
def callback(i, t, latents):
print(f"sd_step_{i}")
if sd_device == 'cuda' and not torch.cuda.is_available():
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == 'cuda' else 1
model = ModelManager(name="sd1.5",
sd_steps = 50 if sd_device == "cuda" else 1
model = ModelManager(
name="sd1.5",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=cpu_textencoder,
callback=callback)
cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=sd_steps)
callback=callback,
)
cfg = get_config(strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
@ -80,35 +88,37 @@ def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_ns
f"runway_sd_{strategy.capitalize()}_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.3
fx=1.3,
)
@pytest.mark.parametrize("sd_device", ['cuda'])
@pytest.mark.parametrize("sd_device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
def callback(i, t, latents):
pass
if sd_device == 'cuda' and not torch.cuda.is_available():
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == 'cuda' else 1
model = ModelManager(name="sd1.5",
sd_steps = 50 if sd_device == "cuda" else 1
model = ModelManager(
name="sd1.5",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=False,
sd_cpu_textencoder=False,
callback=callback)
callback=callback,
)
cfg = get_config(
strategy,
sd_steps=sd_steps,
prompt='Face of a fox, high resolution, sitting on a park bench',
negative_prompt='orange, yellow, small',
prompt="Face of a fox, high resolution, sitting on a park bench",
negative_prompt="orange, yellow, small",
sd_sampler=sampler,
sd_match_histograms=True
sd_match_histograms=True,
)
name = f"{sampler}_negative_prompt"
@ -119,27 +129,33 @@ def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
f"runway_sd_{strategy.capitalize()}_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1
fx=1,
)
@pytest.mark.parametrize("sd_device", ['cuda'])
@pytest.mark.parametrize("sd_device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
@pytest.mark.parametrize("cpu_textencoder", [False])
@pytest.mark.parametrize("disable_nsfw", [False])
def test_runway_sd_1_5_sd_scale(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
if sd_device == 'cuda' and not torch.cuda.is_available():
def test_runway_sd_1_5_sd_scale(
sd_device, strategy, sampler, cpu_textencoder, disable_nsfw
):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == 'cuda' else 1
model = ModelManager(name="sd1.5",
sd_steps = 50 if sd_device == "cuda" else 1
model = ModelManager(
name="sd1.5",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=cpu_textencoder)
cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=sd_steps, sd_scale=0.85)
sd_cpu_textencoder=cpu_textencoder,
)
cfg = get_config(
strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps, sd_scale=0.85
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
@ -150,26 +166,30 @@ def test_runway_sd_1_5_sd_scale(sd_device, strategy, sampler, cpu_textencoder, d
f"runway_sd_{strategy.capitalize()}_{name}_sdscale.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.3
fx=1.3,
)
@pytest.mark.parametrize("sd_device", ['cuda'])
@pytest.mark.parametrize("sd_device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
if sd_device == 'cuda' and not torch.cuda.is_available():
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == 'cuda' else 1
model = ModelManager(name="sd1.5",
sd_steps = 50 if sd_device == "cuda" else 1
model = ModelManager(
name="sd1.5",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=True,
sd_cpu_textencoder=False,
cpu_offload=True)
cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=sd_steps, sd_scale=0.85)
cpu_offload=True,
)
cfg = get_config(
strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps, sd_scale=0.85
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
@ -182,27 +202,3 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("sd_device", ['cpu'])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
def test_runway_sd_1_5_cpu_offload_cpu_device(sd_device, strategy, sampler):
model = ModelManager(name="sd1.5",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=False,
sd_cpu_textencoder=False,
cpu_offload=True)
cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=1, sd_scale=0.85)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
assert_equal(
model,
cfg,
f"runway_sd_{strategy.capitalize()}_{name}_cpu_offload_cpu_device.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)

82
lama_cleaner/web_config.py
View File

@ -6,9 +6,24 @@ import gradio as gr
from loguru import logger
from pydantic import BaseModel
from lama_cleaner.const import AVAILABLE_MODELS, AVAILABLE_DEVICES, CPU_OFFLOAD_HELP, NO_HALF_HELP, DISABLE_NSFW_HELP, \
SD_CPU_TEXTENCODER_HELP, LOCAL_FILES_ONLY_HELP, ENABLE_XFORMERS_HELP, MODEL_DIR_HELP, OUTPUT_DIR_HELP, INPUT_HELP, \
GUI_HELP, DEFAULT_MODEL, DEFAULT_DEVICE, NO_GUI_AUTO_CLOSE_HELP, DEFAULT_MODEL_DIR
from lama_cleaner.const import (
AVAILABLE_MODELS,
AVAILABLE_DEVICES,
CPU_OFFLOAD_HELP,
NO_HALF_HELP,
DISABLE_NSFW_HELP,
SD_CPU_TEXTENCODER_HELP,
LOCAL_FILES_ONLY_HELP,
ENABLE_XFORMERS_HELP,
MODEL_DIR_HELP,
OUTPUT_DIR_HELP,
INPUT_HELP,
GUI_HELP,
DEFAULT_MODEL,
DEFAULT_DEVICE,
NO_GUI_AUTO_CLOSE_HELP,
DEFAULT_MODEL_DIR, MPS_SUPPORT_MODELS,
)
_config_file = None
@ -33,16 +48,28 @@ class Config(BaseModel):
def load_config(installer_config: str):
if os.path.exists(installer_config):
with open(installer_config, "r", encoding='utf-8') as f:
with open(installer_config, "r", encoding="utf-8") as f:
return Config(**json.load(f))
else:
return Config()
def save_config(
host, port, model, device, gui, no_gui_auto_close, no_half, cpu_offload,
disable_nsfw, sd_cpu_textencoder, enable_xformers, local_files_only,
model_dir, input, output_dir
host,
port,
model,
device,
gui,
no_gui_auto_close,
no_half,
cpu_offload,
disable_nsfw,
sd_cpu_textencoder,
enable_xformers,
local_files_only,
model_dir,
input,
output_dir,
):
config = Config(**locals())
print(config)
@ -63,6 +90,7 @@ def save_config(
def close_server(*args):
# TODO: make close both browser and server works
import os, signal
pid = os.getpid()
os.kill(pid, signal.SIGUSR1)
@ -86,19 +114,37 @@ def main(config_file: str):
port = gr.Number(init_config.port, label="Port", precision=0)
with gr.Row():
model = gr.Radio(AVAILABLE_MODELS, label="Model", value=init_config.model)
device = gr.Radio(AVAILABLE_DEVICES, label="Device", value=init_config.device)
device = gr.Radio(
AVAILABLE_DEVICES, label=f"Device(mps supports {MPS_SUPPORT_MODELS})", value=init_config.device
)
gui = gr.Checkbox(init_config.gui, label=f"{GUI_HELP}")
no_gui_auto_close = gr.Checkbox(init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}")
no_gui_auto_close = gr.Checkbox(
init_config.no_gui_auto_close, label=f"{NO_GUI_AUTO_CLOSE_HELP}"
)
no_half = gr.Checkbox(init_config.no_half, label=f"{NO_HALF_HELP}")
cpu_offload = gr.Checkbox(init_config.cpu_offload, label=f"{CPU_OFFLOAD_HELP}")
disable_nsfw = gr.Checkbox(init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}")
sd_cpu_textencoder = gr.Checkbox(init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}")
enable_xformers = gr.Checkbox(init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}")
local_files_only = gr.Checkbox(init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}")
disable_nsfw = gr.Checkbox(
init_config.disable_nsfw, label=f"{DISABLE_NSFW_HELP}"
)
sd_cpu_textencoder = gr.Checkbox(
init_config.sd_cpu_textencoder, label=f"{SD_CPU_TEXTENCODER_HELP}"
)
enable_xformers = gr.Checkbox(
init_config.enable_xformers, label=f"{ENABLE_XFORMERS_HELP}"
)
local_files_only = gr.Checkbox(
init_config.local_files_only, label=f"{LOCAL_FILES_ONLY_HELP}"
)
model_dir = gr.Textbox(init_config.model_dir, label=f"{MODEL_DIR_HELP}")
input = gr.Textbox(init_config.input, label=f"Input file or directory. {INPUT_HELP}")
output_dir = gr.Textbox(init_config.output_dir, label=f"Output directory. {OUTPUT_DIR_HELP}")
save_btn.click(save_config, [
input = gr.Textbox(
init_config.input, label=f"Input file or directory. {INPUT_HELP}"
)
output_dir = gr.Textbox(
init_config.output_dir, label=f"Output directory. {OUTPUT_DIR_HELP}"
)
save_btn.click(
save_config,
[
host,
port,
model,
@ -114,5 +160,7 @@ def main(config_file: str):
model_dir,
input,
output_dir,
], message)
],
message,
)
demo.launch(inbrowser=True, show_api=False)

4
requirements.txt
View File

@ -1,6 +1,7 @@
torch>=1.9.0
opencv-python
flask_cors
Jinja2==2.11.3
flask==1.1.4
flaskwebgui==0.3.5
tqdm
@ -11,7 +12,8 @@ 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
piexif==1.1.3

2
setup.py
View File

@ -21,7 +21,7 @@ def load_requirements():
# https://setuptools.readthedocs.io/en/latest/setuptools.html#including-data-files
setuptools.setup(
name="lama-cleaner",
version="0.34.0",
version="0.35.0",
author="PanicByte",
author_email="cwq1913@gmail.com",
description="Image inpainting tool powered by SOTA AI Model",