add pix2pix

lama_cleaner/app/src/adapters/inpainting.ts

@@ -82,6 +82,11 @@ export default async function inpaint(
     fd.append('paintByExampleImage', paintByExampleImage)
   }
 
+  // InstructPix2Pix
+  fd.append('p2pSteps', settings.p2pSteps.toString())
+  fd.append('p2pImageGuidanceScale', settings.p2pImageGuidanceScale.toString())
+  fd.append('p2pGuidanceScale', settings.p2pGuidanceScale.toString())
+
   if (sizeLimit === undefined) {
     fd.append('sizeLimit', '1080')
   } else {

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

@@ -41,10 +41,10 @@ import {
   croperState,
   enableFileManagerState,
   fileState,
-  gifImageState,
   imageHeightState,
   imageWidthState,
   interactiveSegClicksState,
+  isDiffusionModelsState,
   isInpaintingState,
   isInteractiveSegRunningState,
   isInteractiveSegState,
@@ -118,8 +118,7 @@ export default function Editor() {
   const setToastState = useSetRecoilState(toastState)
   const [isInpainting, setIsInpainting] = useRecoilState(isInpaintingState)
   const runMannually = useRecoilValue(runManuallyState)
-  const isSD = useRecoilValue(isSDState)
+  const isDiffusionModels = useRecoilValue(isDiffusionModelsState)
-  const isPaintByExample = useRecoilValue(isPaintByExampleState)
   const [isInteractiveSeg, setIsInteractiveSeg] = useRecoilState(
     isInteractiveSegState
   )
@@ -842,7 +841,7 @@ export default function Editor() {
     }
 
     if (
-      (isSD || isPaintByExample) &&
+      isDiffusionModels &&
       settings.showCroper &&
       isOutsideCroper(mouseXY(ev))
     ) {
@@ -1388,7 +1387,7 @@ export default function Editor() {
             minHeight={Math.min(256, original.naturalHeight)}
             minWidth={Math.min(256, original.naturalWidth)}
             scale={scale}
-            show={(isSD || isPaintByExample) && settings.showCroper}
+            show={isDiffusionModels && settings.showCroper}
           />
 
           {isInteractiveSeg ? <InteractiveSeg /> : <></>}
@@ -1442,7 +1441,7 @@ export default function Editor() {
       )}
 
       <div className="editor-toolkit-panel">
-        {isSD || isPaintByExample || file === undefined ? (
+        {isDiffusionModels || file === undefined ? (
           <></>
         ) : (
           <SizeSelector
@@ -1545,7 +1544,7 @@ export default function Editor() {
             onClick={download}
           />
 
-          {settings.runInpaintingManually && !isSD && !isPaintByExample && (
+          {settings.runInpaintingManually && !isDiffusionModels && (
             <Button
               toolTip="Run Inpainting"
               icon={

lama_cleaner/app/src/components/Header/Header.tsx

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

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

@@ -179,28 +179,16 @@ function ModelSettingBlock() {
 
   const renderOptionDesc = (): ReactNode => {
     switch (setting.model) {
-      case AIModel.LAMA:
-        return undefined
       case AIModel.LDM:
         return renderLDMModelDesc()
       case AIModel.ZITS:
         return renderZITSModelDesc()
-      case AIModel.MAT:
-        return undefined
       case AIModel.FCF:
         return renderFCFModelDesc()
-      case AIModel.SD15:
-        return undefined
-      case AIModel.SD2:
-        return undefined
-      case AIModel.PAINT_BY_EXAMPLE:
-        return undefined
-      case AIModel.Mange:
-        return undefined
       case AIModel.CV2:
         return renderOpenCV2Desc()
       default:
-        return <></>
+        return undefined
     }
   }
 
@@ -266,6 +254,12 @@ function ModelSettingBlock() {
           'https://arxiv.org/abs/2211.13227',
           'https://github.com/Fantasy-Studio/Paint-by-Example'
         )
+      case AIModel.PIX2PIX:
+        return renderModelDesc(
+          'InstructPix2Pix',
+          'https://arxiv.org/abs/2211.09800',
+          'https://github.com/timothybrooks/instruct-pix2pix'
+        )
       default:
         return <></>
     }

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

@@ -1,6 +1,7 @@
 import React from 'react'
 import { useRecoilState, useRecoilValue } from 'recoil'
 import {
+  isDiffusionModelsState,
   isPaintByExampleState,
   isSDState,
   settingState,
@@ -28,7 +29,7 @@ export default function SettingModal(props: SettingModalProps) {
   const { onClose } = props
   const [setting, setSettingState] = useRecoilState(settingState)
   const isSD = useRecoilValue(isSDState)
-  const isPaintByExample = useRecoilValue(isPaintByExampleState)
+  const isDiffusionModels = useRecoilValue(isDiffusionModelsState)
 
   const handleOnClose = () => {
     setSettingState(old => {
@@ -56,9 +57,9 @@ export default function SettingModal(props: SettingModalProps) {
       show={setting.show}
     >
       <DownloadMaskSettingBlock />
-      {isSD || isPaintByExample ? <></> : <ManualRunInpaintingSettingBlock />}
+      {isDiffusionModels ? <></> : <ManualRunInpaintingSettingBlock />}
       <ModelSettingBlock />
-      {isSD ? <></> : <HDSettingBlock />}
+      {isDiffusionModels ? <></> : <HDSettingBlock />}
     </Modal>
   )
 }

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

@@ -0,0 +1,224 @@
+import React, { FormEvent } from 'react'
+import { useRecoilState, useRecoilValue } from 'recoil'
+import * as PopoverPrimitive from '@radix-ui/react-popover'
+import { useToggle } from 'react-use'
+import {
+  isInpaintingState,
+  negativePropmtState,
+  propmtState,
+  settingState,
+} from '../../store/Atoms'
+import NumberInputSetting from '../Settings/NumberInputSetting'
+import SettingBlock from '../Settings/SettingBlock'
+import { Switch, SwitchThumb } from '../shared/Switch'
+import TextAreaInput from '../shared/Textarea'
+import emitter, { EVENT_PROMPT } from '../../event'
+import ImageResizeScale from './ImageResizeScale'
+
+const INPUT_WIDTH = 30
+
+const P2PSidePanel = () => {
+  const [open, toggleOpen] = useToggle(true)
+  const [setting, setSettingState] = useRecoilState(settingState)
+  const [negativePrompt, setNegativePrompt] =
+    useRecoilState(negativePropmtState)
+  const isInpainting = useRecoilValue(isInpaintingState)
+  const prompt = useRecoilValue(propmtState)
+
+  const handleOnInput = (evt: FormEvent<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

lama_cleaner/app/src/components/Workspace.tsx

@@ -8,6 +8,7 @@ import {
   AIModel,
   fileState,
   isPaintByExampleState,
+  isPix2PixState,
   isSDState,
   settingState,
   showFileManagerState,
@@ -22,6 +23,7 @@ import {
 import SidePanel from './SidePanel/SidePanel'
 import PESidePanel from './SidePanel/PESidePanel'
 import FileManager from './FileManager/FileManager'
+import P2PSidePanel from './SidePanel/P2PSidePanel'
 
 const Workspace = () => {
   const setFile = useSetRecoilState(fileState)
@@ -29,6 +31,7 @@ const Workspace = () => {
   const [toastVal, setToastState] = useRecoilState(toastState)
   const isSD = useRecoilValue(isSDState)
   const isPaintByExample = useRecoilValue(isPaintByExampleState)
+  const isPix2Pix = useRecoilValue(isPix2PixState)
 
   const [showFileManager, setShowFileManager] =
     useRecoilState(showFileManagerState)
@@ -98,6 +101,7 @@ const Workspace = () => {
     <>
       {isSD ? <SidePanel /> : <></>}
       {isPaintByExample ? <PESidePanel /> : <></>}
+      {isPix2Pix ? <P2PSidePanel /> : <></>}
       <FileManager
         photoWidth={256}
         show={showFileManager}

lama_cleaner/app/src/store/Atoms.tsx

@@ -14,6 +14,7 @@ export enum AIModel {
   CV2 = 'cv2',
   Mange = 'manga',
   PAINT_BY_EXAMPLE = 'paint_by_example',
+  PIX2PIX = 'pix2pix',
 }
 
 export const maskState = atom<File | undefined>({
@@ -343,6 +344,11 @@ export interface Settings {
   paintByExampleSeedFixed: boolean
   paintByExampleMaskBlur: number
   paintByExampleMatchHistograms: boolean
+
+  // InstructPix2Pix
+  p2pSteps: number
+  p2pImageGuidanceScale: number
+  p2pGuidanceScale: number
 }
 
 const defaultHDSettings: ModelsHDSettings = {
@@ -402,6 +408,13 @@ const defaultHDSettings: ModelsHDSettings = {
     hdStrategyCropMargin: 128,
     enabled: false,
   },
+  [AIModel.PIX2PIX]: {
+    hdStrategy: HDStrategy.ORIGINAL,
+    hdStrategyResizeLimit: 768,
+    hdStrategyCropTrigerSize: 512,
+    hdStrategyCropMargin: 128,
+    enabled: false,
+  },
   [AIModel.Mange]: {
     hdStrategy: HDStrategy.CROP,
     hdStrategyResizeLimit: 1280,
@@ -471,6 +484,11 @@ export const settingStateDefault: Settings = {
   paintByExampleMaskBlur: 5,
   paintByExampleSeedFixed: false,
   paintByExampleMatchHistograms: false,
+
+  // InstructPix2Pix
+  p2pSteps: 50,
+  p2pImageGuidanceScale: 1.5,
+  p2pGuidanceScale: 7.5,
 }
 
 const localStorageEffect =
@@ -567,12 +585,33 @@ export const isPaintByExampleState = selector({
   },
 })
 
+export const isPix2PixState = selector({
+  key: 'isPix2PixState',
+  get: ({ get }) => {
+    const settings = get(settingState)
+    return settings.model === AIModel.PIX2PIX
+  },
+})
+
 export const runManuallyState = selector({
   key: 'runManuallyState',
   get: ({ get }) => {
     const settings = get(settingState)
     const isSD = get(isSDState)
     const isPaintByExample = get(isPaintByExampleState)
-    return settings.runInpaintingManually || isSD || isPaintByExample
+    const isPix2Pix = get(isPix2PixState)
+    return (
+      settings.runInpaintingManually || isSD || isPaintByExample || isPix2Pix
+    )
+  },
+})
+
+export const isDiffusionModelsState = selector({
+  key: 'isDiffusionModelsState',
+  get: ({ get }) => {
+    const isSD = get(isSDState)
+    const isPaintByExample = get(isPaintByExampleState)
+    const isPix2Pix = get(isPix2PixState)
+    return isSD || isPaintByExample || isPix2Pix
   },
 })

lama_cleaner/const.py

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

lama_cleaner/model/pix2pix.py

@@ -0,0 +1,83 @@
+import PIL.Image
+import cv2
+import torch
+from loguru import logger
+
+from lama_cleaner.model.base import DiffusionInpaintModel
+from lama_cleaner.model.utils import set_seed
+from lama_cleaner.schema import Config
+
+
+class Pix2Pix(DiffusionInpaintModel):
+    pad_mod = 8
+    min_size = 512
+
+    def init_model(self, device: torch.device, **kwargs):
+        from diffusers import StableDiffusionInstructPix2PixPipeline
+        fp16 = not kwargs.get('no_half', False)
+
+        model_kwargs = {"local_files_only": kwargs.get('local_files_only', kwargs['sd_run_local'])}
+        if kwargs['disable_nsfw'] or kwargs.get('cpu_offload', False):
+            logger.info("Disable Stable Diffusion Model NSFW checker")
+            model_kwargs.update(dict(
+                safety_checker=None,
+                feature_extractor=None,
+                requires_safety_checker=False
+            ))
+
+        use_gpu = device == torch.device('cuda') and torch.cuda.is_available()
+        torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
+        self.model = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+            "timbrooks/instruct-pix2pix",
+            revision="fp16" if use_gpu and fp16 else "main",
+            torch_dtype=torch_dtype,
+            **model_kwargs
+        )
+
+        self.model.enable_attention_slicing()
+        if kwargs.get('enable_xformers', False):
+            self.model.enable_xformers_memory_efficient_attention()
+
+        if kwargs.get('cpu_offload', False) and use_gpu:
+            logger.info("Enable sequential cpu offload")
+            self.model.enable_sequential_cpu_offload(gpu_id=0)
+        else:
+            self.model = self.model.to(device)
+
+    def forward(self, image, mask, config: Config):
+        """Input image and output image have same size
+        image: [H, W, C] RGB
+        mask: [H, W, 1] 255 means area to repaint
+        return: BGR IMAGE
+        edit = pipe(prompt, image=image, num_inference_steps=20, image_guidance_scale=1.5, guidance_scale=7).images[0]
+        """
+        set_seed(config.sd_seed)
+
+        output = self.model(
+            image=PIL.Image.fromarray(image),
+            prompt=config.prompt,
+            negative_prompt=config.negative_prompt,
+            num_inference_steps=config.p2p_steps,
+            image_guidance_scale=config.p2p_image_guidance_scale,
+            guidance_scale=config.p2p_guidance_scale,
+            output_type="np.array",
+        ).images[0]
+
+        output = (output * 255).round().astype("uint8")
+        output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
+        return output
+
+    #
+    # def forward_post_process(self, result, image, mask, config):
+    #     if config.sd_match_histograms:
+    #         result = self._match_histograms(result, image[:, :, ::-1], mask)
+    #
+    #     if config.sd_mask_blur != 0:
+    #         k = 2 * config.sd_mask_blur + 1
+    #         mask = cv2.GaussianBlur(mask, (k, k), 0)
+    #     return result, image, mask
+
+    @staticmethod
+    def is_downloaded() -> bool:
+        # model will be downloaded when app start, and can't switch in frontend settings
+        return True

lama_cleaner/model_manager.py

@@ -7,13 +7,14 @@ from lama_cleaner.model.ldm import LDM
 from lama_cleaner.model.manga import Manga
 from lama_cleaner.model.mat import MAT
 from lama_cleaner.model.paint_by_example import PaintByExample
+from lama_cleaner.model.pix2pix import Pix2Pix
 from lama_cleaner.model.sd import SD15, SD2
 from lama_cleaner.model.zits import ZITS
 from lama_cleaner.model.opencv2 import OpenCV2
 from lama_cleaner.schema import Config
 
 models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.5": SD15, "cv2": OpenCV2, "manga": Manga,
-          "sd2": SD2, "paint_by_example": PaintByExample}
+          "sd2": SD2, "paint_by_example": PaintByExample, "pix2pix": Pix2Pix}
 
 
 class ModelManager:

lama_cleaner/schema.py

@@ -88,3 +88,8 @@ class Config(BaseModel):
     paint_by_example_seed: int = 42
     paint_by_example_match_histograms: bool = False
     paint_by_example_example_image: Image = None
+
+    # InstructPix2Pix
+    p2p_steps: int = 50
+    p2p_image_guidance_scale: float = 1.5
+    p2p_guidance_scale: float = 7.5

lama_cleaner/server.py

@@ -228,6 +228,9 @@ def process():
         paint_by_example_seed=form["paintByExampleSeed"],
         paint_by_example_match_histograms=form["paintByExampleMatchHistograms"],
         paint_by_example_example_image=paint_by_example_example_image,
+        p2p_steps=form["p2pSteps"],
+        p2p_image_guidance_scale=form["p2pImageGuidanceScale"],
+        p2p_guidance_scale=form["p2pGuidanceScale"],
     )
 
     if config.sd_seed == -1:

requirements.txt

@@ -11,7 +11,7 @@ pytest
 yacs
 markupsafe==2.0.1
 scikit-image==0.19.3
-diffusers[torch]==0.10.2
+diffusers[torch]==0.12.1
 transformers>=4.25.1
 watchdog==2.2.1
 gradio