From 6a0ffdc96e131a1778a55a30a7a202f617607407 Mon Sep 17 00:00:00 2001 From: Qing Date: Sun, 4 Dec 2022 13:41:48 +0800 Subject: [PATCH] update diffusers to 0.9; add SD2 --- .../components/Settings/ModelSettingBlock.tsx | 10 +++++- lama_cleaner/app/src/store/Atoms.tsx | 13 +++++-- lama_cleaner/model/sd.py | 36 ++++++++----------- lama_cleaner/model_manager.py | 5 +-- lama_cleaner/parse_args.py | 4 +-- lama_cleaner/schema.py | 1 + requirements.txt | 2 +- 7 files changed, 41 insertions(+), 30 deletions(-) diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx index e6291a7..d004f74 100644 --- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx +++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx @@ -191,6 +191,8 @@ function ModelSettingBlock() { return renderFCFModelDesc() case AIModel.SD15: return undefined + case AIModel.SD2: + return undefined case AIModel.Mange: return undefined case AIModel.CV2: @@ -234,10 +236,16 @@ function ModelSettingBlock() { ) case AIModel.SD15: return renderModelDesc( - 'Stable Diffusion', + 'Stable Diffusion 1.5', 'https://ommer-lab.com/research/latent-diffusion-models/', 'https://github.com/CompVis/stable-diffusion' ) + case AIModel.SD2: + return renderModelDesc( + 'Stable Diffusion 2', + 'https://ommer-lab.com/research/latent-diffusion-models/', + 'https://github.com/Stability-AI/stablediffusion' + ) case AIModel.Mange: return renderModelDesc( 'Manga Inpainting', diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx index d2d2130..353c7d6 100644 --- a/lama_cleaner/app/src/store/Atoms.tsx +++ b/lama_cleaner/app/src/store/Atoms.tsx @@ -10,6 +10,7 @@ export enum AIModel { MAT = 'mat', FCF = 'fcf', SD15 = 'sd1.5', + SD2 = 'sd2', CV2 = 'cv2', Mange = 'manga', } @@ -294,7 +295,14 @@ const defaultHDSettings: ModelsHDSettings = { hdStrategyResizeLimit: 768, hdStrategyCropTrigerSize: 512, hdStrategyCropMargin: 128, - enabled: true, + enabled: false, + }, + [AIModel.SD2]: { + hdStrategy: HDStrategy.ORIGINAL, + hdStrategyResizeLimit: 768, + hdStrategyCropTrigerSize: 512, + hdStrategyCropMargin: 128, + enabled: false, }, [AIModel.Mange]: { hdStrategy: HDStrategy.CROP, @@ -318,6 +326,7 @@ export enum SDSampler { klms = 'k_lms', kEuler = 'k_euler', kEulerA = 'k_euler_a', + dpmPlusPlus = 'dpm++', } export enum SDMode { @@ -422,7 +431,7 @@ export const isSDState = selector({ key: 'isSD', get: ({ get }) => { const settings = get(settingState) - return settings.model === AIModel.SD15 + return settings.model === AIModel.SD15 || settings.model === AIModel.SD2 }, }) diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py index 1b0abba..d983f29 100644 --- a/lama_cleaner/model/sd.py +++ b/lama_cleaner/model/sd.py @@ -5,7 +5,7 @@ import cv2 import numpy as np import torch from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, \ - EulerAncestralDiscreteScheduler + EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler from loguru import logger from lama_cleaner.model.base import InpaintModel @@ -102,27 +102,20 @@ class SD(InpaintModel): # image = torch.from_numpy(image).unsqueeze(0).to(self.device) # mask = torch.from_numpy(mask).unsqueeze(0).to(self.device) - scheduler_kwargs = dict( - beta_schedule="scaled_linear", - beta_start=0.00085, - beta_end=0.012, - num_train_timesteps=1000, - ) + scheduler_config = self.model.scheduler.config if config.sd_sampler == SDSampler.ddim: - scheduler = DDIMScheduler( - **scheduler_kwargs, - clip_sample=False, - set_alpha_to_one=False, - ) + scheduler = DDIMScheduler.from_config(scheduler_config) elif config.sd_sampler == SDSampler.pndm: - scheduler = PNDMScheduler(**scheduler_kwargs, skip_prk_steps=True) + scheduler = PNDMScheduler.from_config(scheduler_config) elif config.sd_sampler == SDSampler.k_lms: - scheduler = LMSDiscreteScheduler(**scheduler_kwargs) + scheduler = LMSDiscreteScheduler.from_config(scheduler_config) elif config.sd_sampler == SDSampler.k_euler: - scheduler = EulerDiscreteScheduler(**scheduler_kwargs) + scheduler = EulerDiscreteScheduler.from_config(scheduler_config) elif config.sd_sampler == SDSampler.k_euler_a: - scheduler = EulerAncestralDiscreteScheduler(**scheduler_kwargs) + scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler_config) + elif config.sd_sampler == SDSampler.dpm_plus_plus: + scheduler = DPMSolverMultistepScheduler.from_config(scheduler_config) else: raise ValueError(config.sd_sampler) @@ -138,13 +131,10 @@ class SD(InpaintModel): k = 2 * config.sd_mask_blur + 1 mask = cv2.GaussianBlur(mask, (k, k), 0)[:, :, np.newaxis] - _kwargs = { - self.image_key: PIL.Image.fromarray(image), - } - img_h, img_w = image.shape[:2] output = self.model( + image=PIL.Image.fromarray(image), prompt=config.prompt, negative_prompt=config.negative_prompt, mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"), @@ -155,7 +145,6 @@ class SD(InpaintModel): callback=self.callback, height=img_h, width=img_w, - **_kwargs ).images[0] output = (output * 255).round().astype("uint8") @@ -217,4 +206,7 @@ class SD(InpaintModel): class SD15(SD): model_id_or_path = "runwayml/stable-diffusion-inpainting" - image_key = "image" + + +class SD2(SD): + model_id_or_path = "stabilityai/stable-diffusion-2-inpainting" diff --git a/lama_cleaner/model_manager.py b/lama_cleaner/model_manager.py index 70411e1..c9f2b9b 100644 --- a/lama_cleaner/model_manager.py +++ b/lama_cleaner/model_manager.py @@ -5,12 +5,13 @@ 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.sd import SD15 +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} +models = {"lama": LaMa, "ldm": LDM, "zits": ZITS, "mat": MAT, "fcf": FcF, "sd1.5": SD15, "cv2": OpenCV2, "manga": Manga, + "sd2": SD2} class ModelManager: diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py index b6de29d..6b51e09 100644 --- a/lama_cleaner/parse_args.py +++ b/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"], + choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.5", "cv2", "manga", "sd2"], ) parser.add_argument( "--hf_access_token", @@ -59,7 +59,7 @@ def parse_args(): if imghdr.what(args.input) is None: parser.error(f"invalid --input: {args.input} is not a valid image file") - if args.model.startswith("sd") and not args.sd_run_local: + if args.model == 'sd1.5' and not args.sd_run_local: if not args.hf_access_token.startswith("hf_"): parser.error( f"sd(stable-diffusion) model requires huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens" diff --git a/lama_cleaner/schema.py b/lama_cleaner/schema.py index 41f80bc..47baaa7 100644 --- a/lama_cleaner/schema.py +++ b/lama_cleaner/schema.py @@ -25,6 +25,7 @@ class SDSampler(str, Enum): k_lms = "k_lms" k_euler = 'k_euler' k_euler_a = 'k_euler_a' + dpm_plus_plus = 'dpm++' class Config(BaseModel): diff --git a/requirements.txt b/requirements.txt index 6b32a72..b9d1fa2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,5 +10,5 @@ pytest yacs markupsafe==2.0.1 scikit-image==0.19.3 -diffusers[torch]==0.7.2 +diffusers[torch]==0.9 transformers==4.21.0