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

update

Browse Source
This commit is contained in:
Qing 2023-12-19 13:16:30 +08:00
parent f27fc51e34
commit 141936a937
18 changed files with 479 additions and 358 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lama_cleaner/const.py
View File

@ -84,8 +84,13 @@ SD2_CONTROLNET_CHOICES = [
DEFAULT_SDXL_CONTROLNET_METHOD = "diffusers/controlnet-canny-sdxl-1.0"
SDXL_CONTROLNET_CHOICES = [
"thibaud/controlnet-openpose-sdxl-1.0",
"destitech/controlnet-inpaint-dreamer-sdxl"
"diffusers/controlnet-canny-sdxl-1.0",
"diffusers/controlnet-canny-sdxl-1.0-mid",
"diffusers/controlnet-canny-sdxl-1.0-small"
"diffusers/controlnet-depth-sdxl-1.0",
"diffusers/controlnet-depth-sdxl-1.0-mid",
"diffusers/controlnet-depth-sdxl-1.0-small",
]
SD_LOCAL_MODEL_HELP = """

6
lama_cleaner/download.py
View File

@ -5,7 +5,7 @@ from typing import List
from loguru import logger
from pathlib import Path
from lama_cleaner.const import DIFFUSERS_MODEL_FP16_REVERSION
from lama_cleaner.const import DIFFUSERS_MODEL_FP16_REVERSION, DEFAULT_MODEL_DIR
from lama_cleaner.schema import (
ModelInfo,
ModelType,
@ -117,9 +117,7 @@ def scan_models() -> List[ModelInfo]:
available_models = []
available_models.extend(scan_inpaint_models())
available_models.extend(
scan_single_file_diffusion_models(os.environ["XDG_CACHE_HOME"])
)
available_models.extend(scan_single_file_diffusion_models(DEFAULT_MODEL_DIR))
cache_dir = Path(DIFFUSERS_CACHE)
diffusers_model_names = []

29
lama_cleaner/model/base.py
View File

@ -279,15 +279,12 @@ class DiffusionInpaintModel(InpaintModel):
"""
# boxes = boxes_from_mask(mask)
if config.use_croper:
if config.croper_is_outpainting:
inpaint_result = self._do_outpainting(image, config)
else:
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
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
elif config.use_extender:
inpaint_result = self._do_outpainting(image, config)
else:
inpaint_result = self._scaled_pad_forward(image, mask, config)
@ -297,10 +294,10 @@ class DiffusionInpaintModel(InpaintModel):
# cropper 和 image 在同一个坐标系下，croper_x/y 可能为负数
# 从 image 中 crop 出 outpainting 区域
image_h, image_w = image.shape[:2]
cropper_l = config.croper_x
cropper_t = config.croper_y
cropper_r = config.croper_x + config.croper_width
cropper_b = config.croper_y + config.croper_height
cropper_l = config.extender_x
cropper_t = config.extender_y
cropper_r = config.extender_x + config.extender_width
cropper_b = config.extender_y + config.extender_height
image_l = 0
image_t = 0
image_r = image_w
@ -356,8 +353,8 @@ class DiffusionInpaintModel(InpaintModel):
)[:, :, ::-1]
# 把 cropped_result_image 贴到 outpainting_image 上，这一步不需要 blend
paste_t = 0 if config.croper_y < 0 else config.croper_y
paste_l = 0 if config.croper_x < 0 else config.croper_x
paste_t = 0 if config.extender_y < 0 else config.extender_y
paste_l = 0 if config.extender_x < 0 else config.extender_x
outpainting_image[
paste_t : paste_t + expanded_cropped_result_image.shape[0],
@ -397,8 +394,6 @@ class DiffusionInpaintModel(InpaintModel):
def set_scheduler(self, config: Config):
scheduler_config = self.model.scheduler.config
sd_sampler = config.sd_sampler
if config.sd_lcm_lora:
sd_sampler = SDSampler.lcm
scheduler = get_scheduler(sd_sampler, scheduler_config)
self.model.scheduler = scheduler

18
lama_cleaner/model/controlnet.py
View File

@ -31,6 +31,20 @@ class ControlNet(DiffusionInpaintModel):
pad_mod = 8
min_size = 512
@property
def lcm_lora_id(self):
if self.model_info.model_type in [
ModelType.DIFFUSERS_SD,
ModelType.DIFFUSERS_SD_INPAINT,
]:
return "latent-consistency/lcm-lora-sdv1-5"
if self.model_info.model_type in [
ModelType.DIFFUSERS_SDXL,
ModelType.DIFFUSERS_SDXL_INPAINT,
]:
return "latent-consistency/lcm-lora-sdxl"
raise NotImplementedError(f"Unsupported controlnet lcm model {self.model_info}")
def init_model(self, device: torch.device, **kwargs):
fp16 = not kwargs.get("no_half", False)
model_info: ModelInfo = kwargs["model_info"]
@ -72,7 +86,7 @@ class ControlNet(DiffusionInpaintModel):
)
controlnet = ControlNetModel.from_pretrained(
sd_controlnet_method, torch_dtype=torch_dtype
sd_controlnet_method, torch_dtype=torch_dtype, resume_download=True
)
if model_info.is_single_file_diffusers:
if self.model_info.model_type == ModelType.DIFFUSERS_SD:
@ -81,7 +95,7 @@ class ControlNet(DiffusionInpaintModel):
model_kwargs["num_in_channels"] = 9
self.model = PipeClass.from_single_file(
model_info.path, controlnet=controlnet
model_info.path, controlnet=controlnet, **model_kwargs
).to(torch_dtype)
else:
self.model = PipeClass.from_pretrained(

2
lama_cleaner/model/sdxl.py
View File

@ -39,7 +39,7 @@ class SDXL(DiffusionInpaintModel):
)
else:
vae = AutoencoderKL.from_pretrained(
"madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
"madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch_dtype
)
self.model = StableDiffusionXLInpaintPipeline.from_pretrained(
self.model_id_or_path,

3
lama_cleaner/model/utils.py
View File

@ -16,6 +16,7 @@ from diffusers import (
EulerAncestralDiscreteScheduler,
DPMSolverMultistepScheduler,
UniPCMultistepScheduler,
LCMScheduler
)
from lama_cleaner.schema import SDSampler
@ -939,5 +940,7 @@ def get_scheduler(sd_sampler, scheduler_config):
return DPMSolverMultistepScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.uni_pc:
return UniPCMultistepScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.lcm:
return LCMScheduler.from_config(scheduler_config)
else:
raise ValueError(sd_sampler)

76
lama_cleaner/model_manager.py
View File

@ -1,9 +1,9 @@
import gc
from typing import List, Dict
import torch
from loguru import logger
from lama_cleaner.const import DEFAULT_SD_CONTROLNET_METHOD
from lama_cleaner.download import scan_models
from lama_cleaner.helper import switch_mps_device
from lama_cleaner.model import models, ControlNet, SD, SDXL
@ -18,6 +18,11 @@ class ModelManager:
self.kwargs = kwargs
self.available_models: Dict[str, ModelInfo] = {}
self.scan_models()
self.sd_controlnet = kwargs.get("sd_controlnet", False)
self.sd_controlnet_method = kwargs.get(
"sd_controlnet_method", DEFAULT_SD_CONTROLNET_METHOD
)
self.model = self.init_model(name, device, **kwargs)
def init_model(self, name: str, device, **kwargs):
@ -28,12 +33,17 @@ class ModelManager:
raise NotImplementedError(f"Unsupported model: {name}")
model_info = self.available_models[name]
kwargs = {**kwargs, "model_info": model_info}
sd_controlnet_enabled = kwargs.get("sd_controlnet", False)
kwargs = {
**kwargs,
"model_info": model_info,
"sd_controlnet": self.sd_controlnet,
"sd_controlnet_method": self.sd_controlnet_method,
}
if model_info.model_type in [ModelType.INPAINT, ModelType.DIFFUSERS_OTHER]:
return models[name](device, **kwargs)
if sd_controlnet_enabled:
if self.sd_controlnet:
return ControlNet(device, **kwargs)
else:
if model_info.model_type in [
@ -51,7 +61,7 @@ class ModelManager:
raise NotImplementedError(f"Unsupported model: {name}")
def __call__(self, image, mask, config: Config):
self.switch_controlnet_method(control_method=config.controlnet_method)
self.switch_controlnet_method(config)
self.enable_disable_freeu(config)
self.enable_disable_lcm_lora(config)
return self.model(image, mask, config)
@ -66,40 +76,56 @@ class ModelManager:
return
old_name = self.name
old_sd_controlnet_method = self.sd_controlnet_method
self.name = new_name
if (
self.available_models[new_name].support_controlnet
and self.sd_controlnet_method
not in self.available_models[new_name].controlnets
):
self.sd_controlnet_method = self.available_models[new_name].controlnets[0]
try:
if torch.cuda.memory_allocated() > 0:
# Clear current loaded model from memory
torch.cuda.empty_cache()
del self.model
gc.collect()
del self.model
torch_gc()
self.model = self.init_model(
new_name, switch_mps_device(new_name, self.device), **self.kwargs
)
except Exception as e:
self.name = old_name
self.sd_controlnet_method = old_sd_controlnet_method
logger.info(f"Switch model from {old_name} to {new_name} failed, rollback")
self.model = self.init_model(
old_name, switch_mps_device(old_name, self.device), **self.kwargs
)
raise e
def switch_controlnet_method(self, control_method: str):
if not self.kwargs.get("sd_controlnet"):
return
if self.kwargs["sd_controlnet_method"] == control_method:
return
def switch_controlnet_method(self, config):
if not self.available_models[self.name].support_controlnet:
return
del self.model
torch_gc()
if self.sd_controlnet != config.controlnet_enabled or (
self.sd_controlnet and self.sd_controlnet_method != config.controlnet_method
):
# 可能关闭/开启 controlnet
# 可能开启了 controlnet，切换 controlnet 的方法
old_sd_controlnet = self.sd_controlnet
old_sd_controlnet_method = self.sd_controlnet_method
self.sd_controlnet = config.controlnet_enabled
self.sd_controlnet_method = config.controlnet_method
old_method = self.kwargs["sd_controlnet_method"]
self.kwargs["sd_controlnet_method"] = control_method
self.model = self.init_model(
self.name, switch_mps_device(self.name, self.device), **self.kwargs
)
logger.info(f"Switch ControlNet method from {old_method} to {control_method}")
self.model = self.init_model(
self.name, switch_mps_device(self.name, self.device), **self.kwargs
)
if not config.controlnet_enabled:
logger.info(f"Disable controlnet")
elif old_sd_controlnet_method != config.controlnet_method:
logger.info(
f"Switch Controlnet method from {old_sd_controlnet_method} to {config.controlnet_method}"
)
else:
logger.info(f"Enable controlnet: {config.controlnet_method}")
def enable_disable_freeu(self, config: Config):
if str(self.model.device) == "mps":
@ -120,7 +146,7 @@ class ModelManager:
def enable_disable_lcm_lora(self, config: Config):
if self.available_models[self.name].support_lcm_lora:
if config.sd_lcm_lora:
if not self.model.model.pipe.get_list_adapters():
if not self.model.model.get_list_adapters():
self.model.model.load_lora_weights(self.model.lcm_lora_id)
else:
self.model.model.disable_lora()

12
lama_cleaner/parse_args.py
View File

@ -234,16 +234,6 @@ def parse_args():
"torch.cuda.is_available() is False, please use --device cpu or check your pytorch installation"
)
if args.sd_local_model_path and args.model == "sd1.5":
if not os.path.exists(args.sd_local_model_path):
parser.error(
f"invalid --sd-local-model-path: {args.sd_local_model_path} not exists"
)
if not os.path.isfile(args.sd_local_model_path):
parser.error(
f"invalid --sd-local-model-path: {args.sd_local_model_path} is a directory"
)
os.environ["U2NET_HOME"] = DEFAULT_MODEL_DIR
if args.model_dir and args.model_dir is not None:
if os.path.isfile(args.model_dir):
@ -264,7 +254,7 @@ def parse_args():
scanned_models = scan_models()
if args.model not in [it.name for it in scanned_models]:
parser.error(
f"invalid --model: {args.model} not exists. Available models: {AVAILABLE_MODELS} or {scanned_models}"
f"invalid --model: {args.model} not exists. Available models: {AVAILABLE_MODELS} or {[it.name for it in scanned_models]}"
)
if args.input and args.input is not None:

54
lama_cleaner/schema.py
View File

@ -65,6 +65,28 @@ class ModelInfo(BaseModel):
return SD_CONTROLNET_CHOICES
return []
@computed_field
@property
def support_strength(self) -> bool:
return self.model_type in [
ModelType.DIFFUSERS_SD,
ModelType.DIFFUSERS_SDXL,
ModelType.DIFFUSERS_SD_INPAINT,
ModelType.DIFFUSERS_SDXL_INPAINT,
]
@computed_field
@property
def support_outpainting(self) -> bool:
return self.model_type in [
ModelType.DIFFUSERS_SD,
ModelType.DIFFUSERS_SDXL,
ModelType.DIFFUSERS_SD_INPAINT,
ModelType.DIFFUSERS_SDXL_INPAINT,
] or self.name in [
"kandinsky-community/kandinsky-2-2-decoder-inpaint",
]
@computed_field
@property
def support_lcm_lora(self) -> bool:
@ -129,10 +151,10 @@ class SDSampler(str, Enum):
class FREEUConfig(BaseModel):
s1: float = 1.0
s2: float = 1.0
b1: float = 1.0
b2: float = 1.0
s1: float = 0.9
s2: float = 0.2
b1: float = 1.2
b2: float = 1.4
class Config(BaseModel):
@ -140,18 +162,18 @@ class Config(BaseModel):
arbitrary_types_allowed = True
# Configs for ldm model
ldm_steps: int
ldm_steps: int = 20
ldm_sampler: str = LDMSampler.plms
# Configs for zits model
zits_wireframe: bool = True
# Configs for High Resolution Strategy(different way to preprocess image)
hd_strategy: str # See HDStrategy Enum
hd_strategy_crop_margin: int
hd_strategy: str = HDStrategy.CROP # See HDStrategy Enum
hd_strategy_crop_margin: int = 128
# If the longer side of the image is larger than this value, use crop strategy
hd_strategy_crop_trigger_size: int
hd_strategy_resize_limit: int
hd_strategy_crop_trigger_size: int = 800
hd_strategy_resize_limit: int = 1280
# Configs for Stable Diffusion 1.5
prompt: str = ""
@ -159,11 +181,15 @@ class Config(BaseModel):
# Crop image to this size before doing sd inpainting
# The value is always on the original image scale
use_croper: bool = False
croper_is_outpainting: bool = False
croper_x: int = None
croper_y: int = None
croper_height: int = None
croper_width: int = None
use_extender: bool = False
extender_x: int = None
extender_y: int = None
extender_height: int = None
extender_width: int = None
# Resize the image before doing sd inpainting, the area outside the mask will not lose quality.
# Used by sd models and paint_by_example model
@ -207,18 +233,12 @@ class Config(BaseModel):
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: Optional[Image] = None
# InstructPix2Pix
p2p_steps: int = 50
p2p_image_guidance_scale: float = 1.5
p2p_guidance_scale: float = 7.5
# ControlNet
controlnet_enabled: bool = False
controlnet_conditioning_scale: float = 0.4
controlnet_method: str = "control_v11p_sd15_canny"

33
lama_cleaner/server.py
View File

@ -1,4 +1,5 @@
#!/usr/bin/env python3
import json
import os
import hashlib
import traceback
@ -103,7 +104,7 @@ logging.getLogger("werkzeug").addFilter(NoFlaskwebgui())
app = Flask(__name__, static_folder=os.path.join(BUILD_DIR, "static"))
app.config["JSON_AS_ASCII"] = False
CORS(app, expose_headers=["Content-Disposition"])
CORS(app, expose_headers=["Content-Disposition", "X-seed"])
sio_logger = logging.getLogger("sio-logger")
sio_logger.setLevel(logging.ERROR)
@ -115,8 +116,6 @@ output_dir: str = None
device = None
input_image_path: str = None
is_disable_model_switch: bool = False
is_controlnet: bool = False
controlnet_method: str = "control_v11p_sd15_canny"
enable_file_manager: bool = False
enable_auto_saving: bool = False
is_desktop: bool = False
@ -266,26 +265,21 @@ def process():
sd_guidance_scale=form["sdGuidanceScale"],
sd_sampler=form["sdSampler"],
sd_seed=form["sdSeed"],
sd_freeu=form["enableFreeu"],
sd_freeu_config=json.loads(form["freeuConfig"]),
sd_lcm_lora=form["enableLCMLora"],
sd_match_histograms=form["sdMatchHistograms"],
cv2_flag=form["cv2Flag"],
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"],
controlnet_enabled=form["controlnet_enabled"],
controlnet_conditioning_scale=form["controlnet_conditioning_scale"],
controlnet_method=form["controlnet_method"],
)
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)
config.sd_seed = random.randint(1, 99999999)
logger.info(f"Origin image shape: {original_shape}")
image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
@ -424,6 +418,8 @@ def get_server_config():
"plugins": list(plugins.keys()),
"enableFileManager": enable_file_manager,
"enableAutoSaving": enable_auto_saving,
"enableControlnet": model.sd_controlnet,
"controlnetMethod": model.sd_controlnet_method,
}, 200
@ -540,18 +536,12 @@ def main(args):
global is_desktop
global thumb
global output_dir
global is_controlnet
global controlnet_method
global image_quality
global enable_auto_saving
build_plugins(args)
image_quality = args.quality
if args.sd_controlnet and args.model in SD15_MODELS:
is_controlnet = True
controlnet_method = args.sd_controlnet_method
output_dir = args.output_dir
if output_dir:
output_dir = os.path.abspath(output_dir)
@ -609,9 +599,6 @@ def main(args):
hf_access_token=args.hf_access_token,
disable_nsfw=args.sd_disable_nsfw or args.disable_nsfw,
sd_cpu_textencoder=args.sd_cpu_textencoder,
sd_run_local=args.sd_run_local,
sd_local_model_path=args.sd_local_model_path,
local_files_only=args.local_files_only,
cpu_offload=args.cpu_offload,
enable_xformers=args.sd_enable_xformers or args.enable_xformers,
callback=diffuser_callback,

5
lama_cleaner/tests/test_controlnet.py
View File

@ -39,7 +39,6 @@ def test_runway_sd_1_5(
sd_controlnet=True,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=False,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=cpu_textencoder,
sd_controlnet_method=sd_controlnet_method,
@ -88,11 +87,9 @@ def test_local_file_path(sd_device, sampler):
sd_controlnet=True,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=False,
disable_nsfw=True,
sd_cpu_textencoder=False,
cpu_offload=True,
sd_local_model_path="/Users/cwq/data/models/sd-v1-5-inpainting.ckpt",
sd_controlnet_method="control_v11p_sd15_canny",
)
cfg = get_config(
@ -128,7 +125,6 @@ def test_local_file_path_controlnet_native_inpainting(sd_device, sampler):
sd_controlnet=True,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=False,
disable_nsfw=True,
sd_cpu_textencoder=False,
cpu_offload=True,
@ -170,7 +166,6 @@ def test_controlnet_switch(sd_device, sampler):
sd_controlnet=True,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=False,
disable_nsfw=True,
sd_cpu_textencoder=False,
cpu_offload=True,

54
lama_cleaner/tests/test_instruct_pix2pix.py
View File

@ -8,23 +8,30 @@ 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 = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
device = 'cuda' if torch.cuda.is_available() else 'mps'
device = "cuda" if torch.cuda.is_available() else "mps"
model_name = "timbrooks/instruct-pix2pix"
@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 20
model = ModelManager(name="instruct_pix2pix",
device=torch.device(device),
hf_access_token="",
sd_run_local=False,
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)
sd_steps = 50 if device == "cuda" else 20
model = ModelManager(
name=model_name,
device=torch.device(device),
hf_access_token="",
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}"
@ -34,22 +41,27 @@ def test_instruct_pix2pix(disable_nsfw, cpu_offload):
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
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 20
model = ModelManager(name="instruct_pix2pix",
device=torch.device(device),
hf_access_token="",
sd_run_local=False,
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_steps = 50 if device == "cuda" else 20
model = ModelManager(
name=model_name,
device=torch.device(device),
hf_access_token="",
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"

2
lama_cleaner/tests/test_model_md5.py
View File

@ -20,8 +20,6 @@ def test_load_model():
hf_access_token="",
disable_nsfw=False,
sd_cpu_textencoder=True,
sd_run_local=True,
local_files_only=True,
cpu_offload=True,
enable_xformers=False,
)

80
lama_cleaner/tests/test_model_switch.py
View File

@ -0,0 +1,80 @@
import logging
import os
from lama_cleaner.schema import Config
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
import torch
from lama_cleaner.model_manager import ModelManager
def test_model_switch():
model = ModelManager(
name="runwayml/stable-diffusion-inpainting",
sd_controlnet=True,
sd_controlnet_method="lllyasviel/control_v11p_sd15_canny",
device=torch.device("mps"),
hf_access_token="",
disable_nsfw=True,
sd_cpu_textencoder=True,
cpu_offload=False,
enable_xformers=False,
callback=None,
)
model.switch("lama")
def test_controlnet_switch_onoff(caplog):
name = "runwayml/stable-diffusion-inpainting"
model = ModelManager(
name=name,
sd_controlnet=True,
sd_controlnet_method="lllyasviel/control_v11p_sd15_canny",
device=torch.device("mps"),
hf_access_token="",
disable_nsfw=True,
sd_cpu_textencoder=True,
cpu_offload=False,
enable_xformers=False,
callback=None,
)
model.switch_controlnet_method(
Config(
name=name,
controlnet_enabled=False,
)
)
assert "Disable controlnet" in caplog.text
def test_controlnet_switch_method(caplog):
name = "runwayml/stable-diffusion-inpainting"
old_method = "lllyasviel/control_v11p_sd15_canny"
new_method = "lllyasviel/control_v11p_sd15_openpose"
model = ModelManager(
name=name,
sd_controlnet=True,
sd_controlnet_method=old_method,
device=torch.device("mps"),
hf_access_token="",
disable_nsfw=True,
sd_cpu_textencoder=True,
cpu_offload=False,
enable_xformers=False,
callback=None,
)
model.switch_controlnet_method(
Config(
name=name,
controlnet_enabled=True,
controlnet_method=new_method,
)
)
assert f"Switch Controlnet method from {old_method} to {new_method}" in caplog.text

34
lama_cleaner/tests/test_outpainting.py
View File

@ -17,7 +17,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device)
@pytest.mark.parametrize("name", ["sd1.5"])
@pytest.mark.parametrize("name", ["runwayml/stable-diffusion-inpainting"])
@pytest.mark.parametrize("sd_device", ["mps"])
@pytest.mark.parametrize(
"rect",
@ -42,7 +42,6 @@ def test_outpainting(name, sd_device, rect):
name=name,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
@ -51,12 +50,11 @@ def test_outpainting(name, sd_device, rect):
HDStrategy.ORIGINAL,
prompt="a dog sitting on a bench in the park",
sd_steps=50,
use_croper=True,
croper_is_outpainting=True,
croper_x=rect[0],
croper_y=rect[1],
croper_width=rect[2],
croper_height=rect[3],
use_extender=True,
extender_x=rect[0],
extender_y=rect[1],
extender_width=rect[2],
extender_height=rect[3],
sd_guidance_scale=8.0,
sd_sampler=SDSampler.dpm_plus_plus,
)
@ -64,13 +62,13 @@ def test_outpainting(name, sd_device, rect):
assert_equal(
model,
cfg,
f"{name.replace('.', '_')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
f"{name.replace('/', '--')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("name", ["kandinsky2.2"])
@pytest.mark.parametrize("name", ["kandinsky-community/kandinsky-2-2-decoder-inpaint"])
@pytest.mark.parametrize("sd_device", ["mps"])
@pytest.mark.parametrize(
"rect",
@ -86,10 +84,9 @@ def test_kandinsky_outpainting(name, sd_device, rect):
return
model = ModelManager(
name="sd1.5",
name=name,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
@ -99,12 +96,11 @@ def test_kandinsky_outpainting(name, sd_device, rect):
prompt="a cat",
negative_prompt="lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature",
sd_steps=50,
use_croper=True,
croper_is_outpainting=True,
croper_x=rect[0],
croper_y=rect[1],
croper_width=rect[2],
croper_height=rect[3],
use_extender=True,
extender_x=rect[0],
extender_y=rect[1],
extender_width=rect[2],
extender_height=rect[3],
sd_guidance_scale=7,
sd_sampler=SDSampler.dpm_plus_plus,
)
@ -112,7 +108,7 @@ def test_kandinsky_outpainting(name, sd_device, rect):
assert_equal(
model,
cfg,
f"{name.replace('.', '_')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
f"{name.replace('/', '--')}_outpainting_dpm++_{'_'.join(map(str, rect))}.png",
img_p=current_dir / "cat.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1,

77
lama_cleaner/tests/test_paint_by_example.py
View File

@ -10,15 +10,19 @@ 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 = 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 "mps"
device = torch.device(device)
model_name = "Fantasy-Studio/Paint-by-Example"
def assert_equal(
model, config, gt_name,
fx: float = 1, fy: float = 1,
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 / "bunny.jpeg",
@ -27,7 +31,9 @@ def assert_equal(
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)
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)
@ -35,14 +41,13 @@ def assert_equal(
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, disable_nsfw=True)
cfg = get_config(strategy, paint_by_example_steps=30)
def test_paint_by_example():
model = ModelManager(name=model_name, device=device, disable_nsfw=True)
cfg = get_config(HDStrategy.ORIGINAL, sd_steps=30)
assert_equal(
model,
cfg,
f"paint_by_example_{strategy.capitalize()}.png",
f"paint_by_example.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fy=0.9,
@ -50,57 +55,31 @@ def test_paint_by_example(strategy):
)
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_paint_by_example_disable_nsfw(strategy):
model = ModelManager(name="paint_by_example", device=device, disable_nsfw=False)
cfg = get_config(strategy, paint_by_example_steps=30)
def test_paint_by_example_cpu_offload():
model = ModelManager(
name=model_name, device=device, cpu_offload=True, disable_nsfw=False
)
cfg = get_config(HDStrategy.ORIGINAL, sd_steps=30)
assert_equal(
model,
cfg,
f"paint_by_example_{strategy.capitalize()}_disable_nsfw.png",
f"paint_by_example_cpu_offload.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_paint_by_example_sd_scale(strategy):
model = ModelManager(name="paint_by_example", device=device, disable_nsfw=True)
cfg = get_config(strategy, paint_by_example_steps=30, sd_scale=0.85)
def test_paint_by_example_cpu_offload_cpu_device():
model = ModelManager(
name=model_name, device=torch.device("cpu"), cpu_offload=True, disable_nsfw=True
)
cfg = get_config(HDStrategy.ORIGINAL, sd_steps=1)
assert_equal(
model,
cfg,
f"paint_by_example_{strategy.capitalize()}_sdscale.png",
f"paint_by_example_cpu_offload_cpu_device.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fy=0.9,
fx=1.3
)
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_paint_by_example_cpu_offload(strategy):
model = ModelManager(name="paint_by_example", device=device, cpu_offload=True, disable_nsfw=False)
cfg = get_config(strategy, paint_by_example_steps=30, sd_scale=0.85)
assert_equal(
model,
cfg,
f"paint_by_example_{strategy.capitalize()}_cpu_offload.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_paint_by_example_cpu_offload_cpu_device(strategy):
model = ModelManager(name="paint_by_example", device=torch.device('cpu'), cpu_offload=True, disable_nsfw=True)
cfg = get_config(strategy, paint_by_example_steps=1, sd_scale=0.85)
assert_equal(
model,
cfg,
f"paint_by_example_{strategy.capitalize()}_cpu_offload_cpu_device.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fy=0.9,
fx=1.3
fx=1.3,
)

246
lama_cleaner/tests/test_sd_model.py
View File

@ -7,7 +7,7 @@ import pytest
import torch
from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy, SDSampler
from lama_cleaner.schema import HDStrategy, SDSampler, FREEUConfig
from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve()
@ -16,178 +16,127 @@ save_dir.mkdir(exist_ok=True, parents=True)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@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 callback(i, t, latents):
pass
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",
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)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
assert_equal(
model,
cfg,
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,
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize(
"sampler", [SDSampler.pndm, SDSampler.k_lms, SDSampler.k_euler, SDSampler.k_euler_a]
"sampler",
[
SDSampler.ddim,
SDSampler.pndm,
SDSampler.k_lms,
SDSampler.k_euler,
SDSampler.k_euler_a,
SDSampler.lcm,
],
)
@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}")
def test_runway_sd_1_5_all_samplers(
sd_device,
sampler,
):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == "cuda" else 1
sd_steps = 30
model = ModelManager(
name="sd1.5",
name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=cpu_textencoder,
callback=callback,
disable_nsfw=True,
sd_cpu_textencoder=False,
)
cfg = get_config(
HDStrategy.ORIGINAL, prompt="a fox sitting on a bench", sd_steps=sd_steps
)
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}"
name = f"device_{sd_device}_{sampler}"
assert_equal(
model,
cfg,
f"runway_sd_{strategy.capitalize()}_{name}.png",
f"runway_sd_{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("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
@pytest.mark.parametrize("sd_prevent_unmasked_area", [False, True])
def test_runway_sd_1_5_negative_prompt(
sd_device, strategy, sampler, sd_prevent_unmasked_area
):
def callback(i, t, latents):
pass
@pytest.mark.parametrize("sampler", [SDSampler.lcm])
def test_runway_sd_lcm_lora(sd_device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == "cuda" else 20
sd_steps = 5
model = ModelManager(
name="sd1.5",
name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=False,
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
)
cfg = get_config(
strategy,
prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps,
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_prevent_unmasked_area=sd_prevent_unmasked_area,
)
name = f"{sampler}_negative_prompt"
assert_equal(
model,
cfg,
f"runway_sd_{strategy.capitalize()}_{name}_prevent_unmasked_area_{sd_prevent_unmasked_area}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1,
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@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():
return
sd_steps = 50 if sd_device == "cuda" else 20
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_guidance_scale=2,
sd_lcm_lora=True,
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
assert_equal(
model,
cfg,
f"runway_sd_{strategy.capitalize()}_{name}_sdscale.png",
f"runway_sd_1_5_lcm_lora.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=1.3,
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_sd_freeu(sd_device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 30
model = ModelManager(
name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device),
hf_access_token="",
disable_nsfw=True,
sd_cpu_textencoder=False,
)
cfg = get_config(
strategy,
prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps,
sd_guidance_scale=7.5,
sd_freeu=True,
sd_freeu_config=FREEUConfig(),
)
cfg.sd_sampler = sampler
assert_equal(
model,
cfg,
f"runway_sd_1_5_freeu.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_sd_sd_strength(sd_device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == "cuda" else 20
sd_steps = 30
model = ModelManager(
name="sd1.5",
name="runwayml/stable-diffusion-inpainting",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=True,
sd_cpu_textencoder=False,
)
@ -205,6 +154,33 @@ def test_runway_sd_sd_strength(sd_device, strategy, sampler):
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_runway_norm_sd_model(sd_device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 30
model = ModelManager(
name="runwayml/stable-diffusion-v1-5",
device=torch.device(sd_device),
hf_access_token="",
disable_nsfw=True,
sd_cpu_textencoder=False,
)
cfg = get_config(strategy, prompt="face of a fox, sitting on a bench", sd_steps=sd_steps)
cfg.sd_sampler = sampler
assert_equal(
model,
cfg,
f"runway_{sd_device}_norm_sd_model.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("sd_device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
@ -212,19 +188,16 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 50 if sd_device == "cuda" else 20
sd_steps = 30
model = ModelManager(
name="sd1.5",
name="runwayml/stable-diffusion-inpainting",
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
)
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}"
@ -239,28 +212,27 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("sampler", [SDSampler.uni_pc])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
@pytest.mark.parametrize(
"local_model_path",
"name",
[
"/Users/cwq/data/models/sd-v1-5-inpainting.ckpt",
"/Users/cwq/data/models/sd-v1-5-inpainting.safetensors",
"sd-v1-5-inpainting.ckpt",
"sd-v1-5-inpainting.safetensors",
"v1-5-pruned-emaonly.safetensors",
],
)
def test_local_file_path(sd_device, sampler, local_model_path):
def test_local_file_path(sd_device, sampler, name):
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 1 if sd_device == "cpu" else 30
sd_steps = 30
model = ModelManager(
name="sd1.5",
name=name,
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=True,
sd_cpu_textencoder=False,
cpu_offload=True,
sd_local_model_path=local_model_path,
cpu_offload=False,
)
cfg = get_config(
HDStrategy.ORIGINAL,
@ -269,7 +241,7 @@ def test_local_file_path(sd_device, sampler, local_model_path):
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_{Path(local_model_path).stem}"
name = f"device_{sd_device}_{sampler}_{name}"
assert_equal(
model,

101
lama_cleaner/tests/test_sdxl.py
View File

@ -7,7 +7,7 @@ import pytest
import torch
from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import HDStrategy, SDSampler
from lama_cleaner.schema import HDStrategy, SDSampler, FREEUConfig
from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve()
@ -15,12 +15,10 @@ save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True)
@pytest.mark.parametrize("sd_device", ["mps"])
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
@pytest.mark.parametrize("cpu_textencoder", [False])
@pytest.mark.parametrize("disable_nsfw", [True])
def test_sdxl(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
def test_sdxl(sd_device, strategy, sampler):
def callback(i, t, latents):
pass
@ -29,24 +27,23 @@ def test_sdxl(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
sd_steps = 20
model = ModelManager(
name="sdxl",
name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=False,
disable_nsfw=disable_nsfw,
sd_cpu_textencoder=cpu_textencoder,
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
)
cfg = get_config(
strategy,
prompt="a fox sitting on a bench",
prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps,
sd_strength=0.99,
sd_strength=1.0,
sd_guidance_scale=7.0,
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}_cpu_textencoder_{cpu_textencoder}_disnsfw_{disable_nsfw}"
name = f"device_{sd_device}_{sampler}"
assert_equal(
model,
@ -59,6 +56,67 @@ def test_sdxl(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
)
@pytest.mark.parametrize("sd_device", ["cuda", "mps"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
def test_sdxl_lcm_lora_and_freeu(sd_device, strategy, sampler):
def callback(i, t, latents):
pass
if sd_device == "cuda" and not torch.cuda.is_available():
return
sd_steps = 5
model = ModelManager(
name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
device=torch.device(sd_device),
hf_access_token="",
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
)
cfg = get_config(
strategy,
prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps,
sd_strength=1.0,
sd_guidance_scale=2.0,
sd_lcm_lora=True,
)
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
assert_equal(
model,
cfg,
f"sdxl_{name}_lcm_lora.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=2,
fy=2,
)
cfg = get_config(
strategy,
prompt="face of a fox, sitting on a bench",
sd_steps=sd_steps,
sd_guidance_scale=7.5,
sd_freeu=True,
sd_freeu_config=FREEUConfig(),
)
assert_equal(
model,
cfg,
f"sdxl_{name}_freeu.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
fx=2,
fy=2,
)
@pytest.mark.parametrize("sd_device", ["mps"])
@pytest.mark.parametrize(
"rect",
@ -67,33 +125,26 @@ def test_sdxl(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw):
],
)
def test_sdxl_outpainting(sd_device, rect):
def callback(i, t, latents):
pass
if sd_device == "cuda" and not torch.cuda.is_available():
return
model = ModelManager(
name="sdxl",
name="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
disable_nsfw=True,
sd_cpu_textencoder=False,
callback=callback,
)
cfg = get_config(
HDStrategy.ORIGINAL,
prompt="a dog sitting on a bench in the park",
negative_prompt="lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature",
sd_steps=20,
use_croper=True,
croper_is_outpainting=True,
croper_x=rect[0],
croper_y=rect[1],
croper_width=rect[2],
croper_height=rect[3],
use_extender=True,
extender_x=rect[0],
extender_y=rect[1],
extender_width=rect[2],
extender_height=rect[3],
sd_strength=1.0,
sd_guidance_scale=8.0,
sd_sampler=SDSampler.ddim,