add paint by example

lama_cleaner/app/src/adapters/inpainting.ts

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

lama_cleaner/app/src/components/Editor/Editor.tsx

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

lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx

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

lama_cleaner/app/src/components/SidePanel/PESidePanel.tsx

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

lama_cleaner/app/src/components/Workspace.tsx

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

lama_cleaner/app/src/event.ts

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

lama_cleaner/app/src/store/Atoms.tsx

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

lama_cleaner/model/base.py

@@ -211,6 +211,26 @@ class InpaintModel:
 
         return result
 
+    def _apply_cropper(self, image, mask, config: Config):
+        img_h, img_w = image.shape[:2]
+        l, t, w, h = (
+            config.croper_x,
+            config.croper_y,
+            config.croper_width,
+            config.croper_height,
+        )
+        r = l + w
+        b = t + h
+
+        l = max(l, 0)
+        r = min(r, img_w)
+        t = max(t, 0)
+        b = min(b, img_h)
+
+        crop_img = image[t:b, l:r, :]
+        crop_mask = mask[t:b, l:r]
+        return crop_img, crop_mask, (l, t, r, b)
+
     def _run_box(self, image, mask, box, config: Config):
         """
 

lama_cleaner/model/paint_by_example.py

@@ -0,0 +1,80 @@
+import random
+
+import PIL
+import PIL.Image
+import cv2
+import numpy as np
+import torch
+from diffusers import DiffusionPipeline
+from lama_cleaner.model.base import InpaintModel
+from lama_cleaner.schema import Config
+
+
+class PaintByExample(InpaintModel):
+    pad_mod = 8
+    min_size = 512
+
+    def init_model(self, device: torch.device, **kwargs):
+        use_gpu = device == torch.device('cuda') and torch.cuda.is_available()
+        torch_dtype = torch.float16 if use_gpu else torch.float32
+        self.model = DiffusionPipeline.from_pretrained(
+            "Fantasy-Studio/Paint-by-Example",
+            torch_dtype=torch_dtype,
+        )
+        self.model.enable_attention_slicing()
+        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
+        """
+        seed = config.paint_by_example_seed
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        output = self.model(
+            image=PIL.Image.fromarray(image),
+            mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
+            example_image=config.paint_by_example_example_image,
+            num_inference_steps=config.paint_by_example_steps,
+            output_type='np.array',
+        ).images[0]
+
+        output = (output * 255).round().astype("uint8")
+        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
+        """
+        if config.use_croper:
+            crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config)
+            crop_image = self._pad_forward(crop_img, crop_mask, config)
+            inpaint_result = image[:, :, ::-1]
+            inpaint_result[t:b, l:r, :] = crop_image
+        else:
+            inpaint_result = self._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)
+
+        if config.paint_by_example_mask_blur != 0:
+            k = 2 * config.paint_by_example_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

lama_cleaner/model/sd.py

@@ -12,31 +12,6 @@ from lama_cleaner.model.base import InpaintModel
 from lama_cleaner.schema import Config, SDSampler
 
 
-#
-#
-# def preprocess_image(image):
-#     w, h = image.size
-#     w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
-#     image = image.resize((w, h), resample=PIL.Image.LANCZOS)
-#     image = np.array(image).astype(np.float32) / 255.0
-#     image = image[None].transpose(0, 3, 1, 2)
-#     image = torch.from_numpy(image)
-#     # [-1, 1]
-#     return 2.0 * image - 1.0
-#
-#
-# def preprocess_mask(mask):
-#     mask = mask.convert("L")
-#     w, h = mask.size
-#     w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
-#     mask = mask.resize((w // 8, h // 8), resample=PIL.Image.NEAREST)
-#     mask = np.array(mask).astype(np.float32) / 255.0
-#     mask = np.tile(mask, (4, 1, 1))
-#     mask = mask[None].transpose(0, 1, 2, 3)  # what does this step do?
-#     mask = 1 - mask  # repaint white, keep black
-#     mask = torch.from_numpy(mask)
-#     return mask
-
 class CPUTextEncoderWrapper:
     def __init__(self, text_encoder, torch_dtype):
         self.config = text_encoder.config
@@ -92,17 +67,6 @@ class SD(InpaintModel):
         return: BGR IMAGE
         """
 
-        # image = norm_img(image)  # [0, 1]
-        # image = image * 2 - 1  # [0, 1] -> [-1, 1]
-
-        # resize to latent feature map size
-        # h, w = mask.shape[:2]
-        # mask = cv2.resize(mask, (h // 8, w // 8), interpolation=cv2.INTER_AREA)
-        # mask = norm_img(mask)
-        #
-        # image = torch.from_numpy(image).unsqueeze(0).to(self.device)
-        # mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
-
         scheduler_config = self.model.scheduler.config
 
         if config.sd_sampler == SDSampler.ddim:
@@ -139,7 +103,6 @@ class SD(InpaintModel):
             prompt=config.prompt,
             negative_prompt=config.negative_prompt,
             mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
-            strength=config.sd_strength,
             num_inference_steps=config.sd_steps,
             guidance_scale=config.sd_guidance_scale,
             output_type="np.array",
@@ -159,30 +122,10 @@ class SD(InpaintModel):
         masks: [H, W]
         return: BGR IMAGE
         """
-        img_h, img_w = image.shape[:2]
-
         # boxes = boxes_from_mask(mask)
         if config.use_croper:
-            logger.info("use croper")
+            crop_img, crop_mask, (l, t, r, b) = self._apply_cropper(image, mask, config)
-            l, t, w, h = (
-                config.croper_x,
-                config.croper_y,
-                config.croper_width,
-                config.croper_height,
-            )
-            r = l + w
-            b = t + h
-
-            l = max(l, 0)
-            r = min(r, img_w)
-            t = max(t, 0)
-            b = min(b, img_h)
-
-            crop_img = image[t:b, l:r, :]
-            crop_mask = mask[t:b, l:r]
-
             crop_image = self._pad_forward(crop_img, crop_mask, config)
-
             inpaint_result = image[:, :, ::-1]
             inpaint_result[t:b, l:r, :] = crop_image
         else:

lama_cleaner/model_manager.py

@@ -5,13 +5,14 @@ from lama_cleaner.model.lama import LaMa
 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.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}
+          "sd2": SD2, "paint_by_example": PaintByExample}
 
 
 class ModelManager:

lama_cleaner/parse_args.py

@@ -10,7 +10,7 @@ def parse_args():
     parser.add_argument(
         "--model",
         default="lama",
-        choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.5", "cv2", "manga", "sd2"],
+        choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.5", "cv2", "manga", "sd2", "paint_by_example"],
     )
     parser.add_argument(
         "--hf_access_token",

lama_cleaner/schema.py

@@ -1,5 +1,6 @@
 from enum import Enum
 
+from PIL.Image import Image
 from pydantic import BaseModel
 
 
@@ -29,6 +30,9 @@ class SDSampler(str, Enum):
 
 
 class Config(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
     # Configs for ldm model
     ldm_steps: int
     ldm_sampler: str = LDMSampler.plms
@@ -73,3 +77,11 @@ class Config(BaseModel):
     # opencv document https://docs.opencv.org/4.6.0/d7/d8b/group__photo__inpaint.html#gga8002a65f5a3328fbf15df81b842d3c3ca05e763003a805e6c11c673a9f4ba7d07
     cv2_flag: str = 'INPAINT_NS'
     cv2_radius: int = 4
+
+    # Paint by Example
+    paint_by_example_steps: int = 50
+    paint_by_example_guidance_scale: float = 7.5
+    paint_by_example_mask_blur: int = 0
+    paint_by_example_seed: int = 42
+    paint_by_example_match_histograms: bool = False
+    paint_by_example_example_image: Image = None

lama_cleaner/server.py

@@ -10,6 +10,7 @@ import time
 import imghdr
 from pathlib import Path
 from typing import Union
+from PIL import Image
 
 import cv2
 import torch
@@ -97,8 +98,8 @@ def process():
     input = request.files
     # RGB
     origin_image_bytes = input["image"].read()
-
     image, alpha_channel = load_img(origin_image_bytes)
+
     mask, _ = load_img(input["mask"].read(), gray=True)
     mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
 
@@ -115,6 +116,12 @@ def process():
     else:
         size_limit = int(size_limit)
 
+    if "paintByExampleImage" in input:
+        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
+
     config = Config(
         ldm_steps=form["ldmSteps"],
         ldm_sampler=form["ldmSampler"],
@@ -138,11 +145,19 @@ 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,
     )
 
     if config.sd_seed == -1:
         config.sd_seed = random.randint(1, 999999999)
+    if config.paint_by_example_seed == -1:
+        config.paint_by_example_seed = random.randint(1, 999999999)
 
     logger.info(f"Origin image shape: {original_shape}")
     image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)

lama_cleaner/tests/test_paint_by_example.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+import cv2
+import pytest
+import torch
+from PIL import Image
+
+from lama_cleaner.model_manager import ModelManager
+from lama_cleaner.schema import HDStrategy
+from lama_cleaner.tests.test_model import get_config, get_data
+
+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'
+device = torch.device(device)
+
+
+def assert_equal(
+    model, config, gt_name,
+    fx: float = 1, fy: float = 1,
+    img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+    mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    example_p=current_dir / "rabbit.jpeg",
+):
+    img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
+
+    example_image = cv2.imread(str(example_p))
+    example_image = cv2.cvtColor(example_image, cv2.COLOR_BGRA2RGB)
+    example_image = cv2.resize(example_image, None, fx=fx, fy=fy, interpolation=cv2.INTER_AREA)
+
+    print(f"Input image shape: {img.shape}, example_image: {example_image.shape}")
+    config.paint_by_example_example_image = Image.fromarray(example_image)
+    res = model(img, mask, config)
+    cv2.imwrite(str(save_dir / gt_name), res)
+
+
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+def test_paint_by_example(strategy):
+    model = ModelManager(name="paint_by_example", device=device)
+    cfg = get_config(strategy, paint_by_example_steps=30)
+    assert_equal(
+        model,
+        cfg,
+        f"paint_by_example_{strategy.capitalize()}.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fy=0.9,
+        fx=1.3
+    )

requirements.txt

@@ -10,5 +10,5 @@ pytest
 yacs
 markupsafe==2.0.1
 scikit-image==0.19.3
-diffusers[torch]==0.9
+diffusers[torch]==0.10.2
-transformers==4.21.0
+transformers>=4.25.1