update diffusers to 0.9; add SD2
This commit is contained in:
parent
15fe87e42d
commit
6a0ffdc96e
@ -191,6 +191,8 @@ function ModelSettingBlock() {
|
|||||||
return renderFCFModelDesc()
|
return renderFCFModelDesc()
|
||||||
case AIModel.SD15:
|
case AIModel.SD15:
|
||||||
return undefined
|
return undefined
|
||||||
|
case AIModel.SD2:
|
||||||
|
return undefined
|
||||||
case AIModel.Mange:
|
case AIModel.Mange:
|
||||||
return undefined
|
return undefined
|
||||||
case AIModel.CV2:
|
case AIModel.CV2:
|
||||||
@ -234,10 +236,16 @@ function ModelSettingBlock() {
|
|||||||
)
|
)
|
||||||
case AIModel.SD15:
|
case AIModel.SD15:
|
||||||
return renderModelDesc(
|
return renderModelDesc(
|
||||||
'Stable Diffusion',
|
'Stable Diffusion 1.5',
|
||||||
'https://ommer-lab.com/research/latent-diffusion-models/',
|
'https://ommer-lab.com/research/latent-diffusion-models/',
|
||||||
'https://github.com/CompVis/stable-diffusion'
|
'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:
|
case AIModel.Mange:
|
||||||
return renderModelDesc(
|
return renderModelDesc(
|
||||||
'Manga Inpainting',
|
'Manga Inpainting',
|
||||||
|
@ -10,6 +10,7 @@ export enum AIModel {
|
|||||||
MAT = 'mat',
|
MAT = 'mat',
|
||||||
FCF = 'fcf',
|
FCF = 'fcf',
|
||||||
SD15 = 'sd1.5',
|
SD15 = 'sd1.5',
|
||||||
|
SD2 = 'sd2',
|
||||||
CV2 = 'cv2',
|
CV2 = 'cv2',
|
||||||
Mange = 'manga',
|
Mange = 'manga',
|
||||||
}
|
}
|
||||||
@ -294,7 +295,14 @@ const defaultHDSettings: ModelsHDSettings = {
|
|||||||
hdStrategyResizeLimit: 768,
|
hdStrategyResizeLimit: 768,
|
||||||
hdStrategyCropTrigerSize: 512,
|
hdStrategyCropTrigerSize: 512,
|
||||||
hdStrategyCropMargin: 128,
|
hdStrategyCropMargin: 128,
|
||||||
enabled: true,
|
enabled: false,
|
||||||
|
},
|
||||||
|
[AIModel.SD2]: {
|
||||||
|
hdStrategy: HDStrategy.ORIGINAL,
|
||||||
|
hdStrategyResizeLimit: 768,
|
||||||
|
hdStrategyCropTrigerSize: 512,
|
||||||
|
hdStrategyCropMargin: 128,
|
||||||
|
enabled: false,
|
||||||
},
|
},
|
||||||
[AIModel.Mange]: {
|
[AIModel.Mange]: {
|
||||||
hdStrategy: HDStrategy.CROP,
|
hdStrategy: HDStrategy.CROP,
|
||||||
@ -318,6 +326,7 @@ export enum SDSampler {
|
|||||||
klms = 'k_lms',
|
klms = 'k_lms',
|
||||||
kEuler = 'k_euler',
|
kEuler = 'k_euler',
|
||||||
kEulerA = 'k_euler_a',
|
kEulerA = 'k_euler_a',
|
||||||
|
dpmPlusPlus = 'dpm++',
|
||||||
}
|
}
|
||||||
|
|
||||||
export enum SDMode {
|
export enum SDMode {
|
||||||
@ -422,7 +431,7 @@ export const isSDState = selector({
|
|||||||
key: 'isSD',
|
key: 'isSD',
|
||||||
get: ({ get }) => {
|
get: ({ get }) => {
|
||||||
const settings = get(settingState)
|
const settings = get(settingState)
|
||||||
return settings.model === AIModel.SD15
|
return settings.model === AIModel.SD15 || settings.model === AIModel.SD2
|
||||||
},
|
},
|
||||||
})
|
})
|
||||||
|
|
||||||
|
@ -5,7 +5,7 @@ import cv2
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import torch
|
import torch
|
||||||
from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, \
|
from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, \
|
||||||
EulerAncestralDiscreteScheduler
|
EulerAncestralDiscreteScheduler, DPMSolverMultistepScheduler
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
from lama_cleaner.model.base import InpaintModel
|
from lama_cleaner.model.base import InpaintModel
|
||||||
@ -102,27 +102,20 @@ class SD(InpaintModel):
|
|||||||
# image = torch.from_numpy(image).unsqueeze(0).to(self.device)
|
# image = torch.from_numpy(image).unsqueeze(0).to(self.device)
|
||||||
# mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
|
# mask = torch.from_numpy(mask).unsqueeze(0).to(self.device)
|
||||||
|
|
||||||
scheduler_kwargs = dict(
|
scheduler_config = self.model.scheduler.config
|
||||||
beta_schedule="scaled_linear",
|
|
||||||
beta_start=0.00085,
|
|
||||||
beta_end=0.012,
|
|
||||||
num_train_timesteps=1000,
|
|
||||||
)
|
|
||||||
|
|
||||||
if config.sd_sampler == SDSampler.ddim:
|
if config.sd_sampler == SDSampler.ddim:
|
||||||
scheduler = DDIMScheduler(
|
scheduler = DDIMScheduler.from_config(scheduler_config)
|
||||||
**scheduler_kwargs,
|
|
||||||
clip_sample=False,
|
|
||||||
set_alpha_to_one=False,
|
|
||||||
)
|
|
||||||
elif config.sd_sampler == SDSampler.pndm:
|
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:
|
elif config.sd_sampler == SDSampler.k_lms:
|
||||||
scheduler = LMSDiscreteScheduler(**scheduler_kwargs)
|
scheduler = LMSDiscreteScheduler.from_config(scheduler_config)
|
||||||
elif config.sd_sampler == SDSampler.k_euler:
|
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:
|
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:
|
else:
|
||||||
raise ValueError(config.sd_sampler)
|
raise ValueError(config.sd_sampler)
|
||||||
|
|
||||||
@ -138,13 +131,10 @@ class SD(InpaintModel):
|
|||||||
k = 2 * config.sd_mask_blur + 1
|
k = 2 * config.sd_mask_blur + 1
|
||||||
mask = cv2.GaussianBlur(mask, (k, k), 0)[:, :, np.newaxis]
|
mask = cv2.GaussianBlur(mask, (k, k), 0)[:, :, np.newaxis]
|
||||||
|
|
||||||
_kwargs = {
|
|
||||||
self.image_key: PIL.Image.fromarray(image),
|
|
||||||
}
|
|
||||||
|
|
||||||
img_h, img_w = image.shape[:2]
|
img_h, img_w = image.shape[:2]
|
||||||
|
|
||||||
output = self.model(
|
output = self.model(
|
||||||
|
image=PIL.Image.fromarray(image),
|
||||||
prompt=config.prompt,
|
prompt=config.prompt,
|
||||||
negative_prompt=config.negative_prompt,
|
negative_prompt=config.negative_prompt,
|
||||||
mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
|
mask_image=PIL.Image.fromarray(mask[:, :, -1], mode="L"),
|
||||||
@ -155,7 +145,6 @@ class SD(InpaintModel):
|
|||||||
callback=self.callback,
|
callback=self.callback,
|
||||||
height=img_h,
|
height=img_h,
|
||||||
width=img_w,
|
width=img_w,
|
||||||
**_kwargs
|
|
||||||
).images[0]
|
).images[0]
|
||||||
|
|
||||||
output = (output * 255).round().astype("uint8")
|
output = (output * 255).round().astype("uint8")
|
||||||
@ -217,4 +206,7 @@ class SD(InpaintModel):
|
|||||||
|
|
||||||
class SD15(SD):
|
class SD15(SD):
|
||||||
model_id_or_path = "runwayml/stable-diffusion-inpainting"
|
model_id_or_path = "runwayml/stable-diffusion-inpainting"
|
||||||
image_key = "image"
|
|
||||||
|
|
||||||
|
class SD2(SD):
|
||||||
|
model_id_or_path = "stabilityai/stable-diffusion-2-inpainting"
|
||||||
|
@ -5,12 +5,13 @@ from lama_cleaner.model.lama import LaMa
|
|||||||
from lama_cleaner.model.ldm import LDM
|
from lama_cleaner.model.ldm import LDM
|
||||||
from lama_cleaner.model.manga import Manga
|
from lama_cleaner.model.manga import Manga
|
||||||
from lama_cleaner.model.mat import MAT
|
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.zits import ZITS
|
||||||
from lama_cleaner.model.opencv2 import OpenCV2
|
from lama_cleaner.model.opencv2 import OpenCV2
|
||||||
from lama_cleaner.schema import Config
|
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:
|
class ModelManager:
|
||||||
|
@ -10,7 +10,7 @@ def parse_args():
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
default="lama",
|
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(
|
parser.add_argument(
|
||||||
"--hf_access_token",
|
"--hf_access_token",
|
||||||
@ -59,7 +59,7 @@ def parse_args():
|
|||||||
if imghdr.what(args.input) is None:
|
if imghdr.what(args.input) is None:
|
||||||
parser.error(f"invalid --input: {args.input} is not a valid image file")
|
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_"):
|
if not args.hf_access_token.startswith("hf_"):
|
||||||
parser.error(
|
parser.error(
|
||||||
f"sd(stable-diffusion) model requires huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens"
|
f"sd(stable-diffusion) model requires huggingface access token. Check how to get token from: https://huggingface.co/docs/hub/security-tokens"
|
||||||
|
@ -25,6 +25,7 @@ class SDSampler(str, Enum):
|
|||||||
k_lms = "k_lms"
|
k_lms = "k_lms"
|
||||||
k_euler = 'k_euler'
|
k_euler = 'k_euler'
|
||||||
k_euler_a = 'k_euler_a'
|
k_euler_a = 'k_euler_a'
|
||||||
|
dpm_plus_plus = 'dpm++'
|
||||||
|
|
||||||
|
|
||||||
class Config(BaseModel):
|
class Config(BaseModel):
|
||||||
|
@ -10,5 +10,5 @@ pytest
|
|||||||
yacs
|
yacs
|
||||||
markupsafe==2.0.1
|
markupsafe==2.0.1
|
||||||
scikit-image==0.19.3
|
scikit-image==0.19.3
|
||||||
diffusers[torch]==0.7.2
|
diffusers[torch]==0.9
|
||||||
transformers==4.21.0
|
transformers==4.21.0
|
||||||
|
Loading…
Reference in New Issue
Block a user