let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

get samplers from backend

Browse Source
This commit is contained in:
Qing
2024-01-02 14:34:36 +08:00
parent a2fd5bb3ea
commit f38be37f8c
14 changed files with 141 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
lama_cleaner/api.py
View File

@@ -37,6 +37,7 @@ from lama_cleaner.schema import (
SwitchModelRequest,
InpaintRequest,
RunPluginRequest,
SDSampler,
)
from lama_cleaner.file_manager import FileManager
@@ -129,6 +130,7 @@ class Api:
self.add_api_route("/api/v1/inputimage", self.api_input_image, methods=["GET"])
self.add_api_route("/api/v1/inpaint", self.api_inpaint, methods=["POST"])
self.add_api_route("/api/v1/run_plugin", self.api_run_plugin, methods=["POST"])
self.add_api_route("/api/v1/samplers", self.api_samplers, methods=["GET"])
self.app.mount("/", StaticFiles(directory=WEB_APP_DIR, html=True), name="assets")
# fmt: on
@@ -156,6 +158,7 @@ class Api:
controlnetMethod=self.model_manager.controlnet_method,
disableModelSwitch=self.config.disable_model_switch,
isDesktop=self.config.gui,
samplers=self.api_samplers(),
)
def api_input_image(self) -> FileResponse:
@@ -237,6 +240,9 @@ class Api:
media_type=f"image/{ext}",
)
def api_samplers(self) -> List[str]:
return [member.value for member in SDSampler.__members__.values()]
def launch(self):
self.app.include_router(self.router)
uvicorn.run(

5
lama_cleaner/download.py
View File

@@ -2,6 +2,7 @@ import json
import os
from typing import List
from huggingface_hub.constants import HF_HUB_CACHE
from loguru import logger
from pathlib import Path
@@ -101,13 +102,11 @@ def scan_inpaint_models(model_dir: Path) -> List[ModelInfo]:
def scan_models() -> List[ModelInfo]:
from diffusers.utils import DIFFUSERS_CACHE
model_dir = os.getenv("XDG_CACHE_HOME", DEFAULT_MODEL_DIR)
available_models = []
available_models.extend(scan_inpaint_models(model_dir))
available_models.extend(scan_single_file_diffusion_models(model_dir))
cache_dir = Path(DIFFUSERS_CACHE)
cache_dir = Path(HF_HUB_CACHE)
# logger.info(f"Scanning diffusers models in {cache_dir}")
diffusers_model_names = []
for it in cache_dir.glob("**/*/model_index.json"):

4
lama_cleaner/model/kandinsky.py
View File

@@ -35,9 +35,7 @@ class Kandinsky(DiffusionInpaintModel):
mask: [H, W, 1] 255 means area to repaint
return: BGR IMAGE
"""
scheduler_config = self.model.scheduler.config
scheduler = get_scheduler(config.sd_sampler, scheduler_config)
self.model.scheduler = scheduler
self.set_scheduler(config)
generator = torch.manual_seed(config.sd_seed)
mask = mask.astype(np.float32) / 255

59
lama_cleaner/model/utils.py
View File

@@ -1,3 +1,4 @@
import copy
import gc
import math
import random
@@ -18,10 +19,13 @@ from diffusers import (
DPMSolverMultistepScheduler,
UniPCMultistepScheduler,
LCMScheduler,
DPMSolverSinglestepScheduler,
KDPM2DiscreteScheduler,
KDPM2AncestralDiscreteScheduler,
HeunDiscreteScheduler,
)
from huggingface_hub.utils import RevisionNotFoundError
from diffusers.configuration_utils import FrozenDict
from loguru import logger
from requests import HTTPError
from lama_cleaner.schema import SDSampler
from torch import conv2d, conv_transpose2d
@@ -930,22 +934,41 @@ def set_seed(seed: int):
def get_scheduler(sd_sampler, scheduler_config):
if sd_sampler == SDSampler.ddim:
return DDIMScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.pndm:
return PNDMScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.k_lms:
return LMSDiscreteScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.k_euler:
return EulerDiscreteScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.k_euler_a:
return EulerAncestralDiscreteScheduler.from_config(scheduler_config)
elif sd_sampler == SDSampler.dpm_plus_plus:
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)
# https://github.com/huggingface/diffusers/issues/4167
keys_to_pop = ["use_karras_sigmas", "algorithm_type"]
scheduler_config = dict(scheduler_config)
for it in keys_to_pop:
scheduler_config.pop(it, None)
# fmt: off
samplers = {
SDSampler.dpm_plus_plus_2m: [DPMSolverMultistepScheduler],
SDSampler.dpm_plus_plus_2m_karras: [DPMSolverMultistepScheduler, dict(use_karras_sigmas=True)],
SDSampler.dpm_plus_plus_2m_sde: [DPMSolverMultistepScheduler, dict(algorithm_type="sde-dpmsolver++")],
SDSampler.dpm_plus_plus_2m_sde_karras: [DPMSolverMultistepScheduler, dict(algorithm_type="sde-dpmsolver++", use_karras_sigmas=True)],
SDSampler.dpm_plus_plus_sde: [DPMSolverSinglestepScheduler],
SDSampler.dpm_plus_plus_sde_karras: [DPMSolverSinglestepScheduler, dict(use_karras_sigmas=True)],
SDSampler.dpm2: [KDPM2DiscreteScheduler],
SDSampler.dpm2_karras: [KDPM2DiscreteScheduler, dict(use_karras_sigmas=True)],
SDSampler.dpm2_a: [KDPM2AncestralDiscreteScheduler],
SDSampler.dpm2_a_karras: [KDPM2AncestralDiscreteScheduler, dict(use_karras_sigmas=True)],
SDSampler.euler: [EulerDiscreteScheduler],
SDSampler.euler_a: [EulerAncestralDiscreteScheduler],
SDSampler.heun: [HeunDiscreteScheduler],
SDSampler.lms: [LMSDiscreteScheduler],
SDSampler.lms_karras: [LMSDiscreteScheduler, dict(use_karras_sigmas=True)],
SDSampler.ddim: [DDIMScheduler],
SDSampler.pndm: [PNDMScheduler],
SDSampler.uni_pc: [UniPCMultistepScheduler],
SDSampler.lcm: [LCMScheduler],
}
# fmt: on
if sd_sampler in samplers:
if len(samplers[sd_sampler]) == 2:
scheduler_cls, kwargs = samplers[sd_sampler]
else:
scheduler_cls, kwargs = samplers[sd_sampler][0], {}
return scheduler_cls.from_config(scheduler_config, **kwargs)
else:
raise ValueError(sd_sampler)

30
lama_cleaner/schema.py
View File

@@ -40,15 +40,26 @@ class LDMSampler(str, Enum):
class SDSampler(str, Enum):
ddim = "ddim"
pndm = "pndm"
k_lms = "k_lms"
k_euler = "k_euler"
k_euler_a = "k_euler_a"
dpm_plus_plus = "dpm++"
uni_pc = "uni_pc"
dpm_plus_plus_2m = "DPM++ 2M"
dpm_plus_plus_2m_karras = "DPM++ 2M Karras"
dpm_plus_plus_2m_sde = "DPM++ 2M SDE"
dpm_plus_plus_2m_sde_karras = "DPM++ 2M SDE Karras"
dpm_plus_plus_sde = "DPM++ SDE"
dpm_plus_plus_sde_karras = "DPM++ SDE Karras"
dpm2 = "DPM2"
dpm2_karras = "DPM2 Karras"
dpm2_a = "DPM2 a"
dpm2_a_karras = "DPM2 a Karras"
euler = "Euler"
euler_a = "Euler a"
heun = "Heun"
lms = "LMS"
lms_karras = "LMS Karras"
lcm = "lcm"
ddim = "DDIM"
pndm = "PNDM"
uni_pc = "UniPC"
lcm = "LCM"
class FREEUConfig(BaseModel):
@@ -143,7 +154,7 @@ class InpaintRequest(BaseModel):
le=1.0,
)
sd_mask_blur: int = Field(
33,
11,
description="Blur the edge of mask area. The higher the number the smoother blend with the original image",
)
sd_strength: float = Field(
@@ -268,6 +279,7 @@ class ServerConfigResponse(BaseModel):
controlnetMethod: Optional[str]
disableModelSwitch: bool
isDesktop: bool
samplers: List[str]
class SwitchModelRequest(BaseModel):

6
lama_cleaner/tests/test_outpainting.py
View File

@@ -49,7 +49,7 @@ def test_outpainting(name, device, rect):
extender_width=rect[2],
extender_height=rect[3],
sd_guidance_scale=8.0,
sd_sampler=SDSampler.dpm_plus_plus,
sd_sampler=SDSampler.dpm_plus_plus_2m,
)
assert_equal(
@@ -92,7 +92,7 @@ def test_kandinsky_outpainting(name, device, rect):
extender_width=rect[2],
extender_height=rect[3],
sd_guidance_scale=7,
sd_sampler=SDSampler.dpm_plus_plus,
sd_sampler=SDSampler.dpm_plus_plus_2m,
)
assert_equal(
@@ -136,7 +136,7 @@ def test_powerpaint_outpainting(name, device, rect):
extender_width=rect[2],
extender_height=rect[3],
sd_guidance_scale=8.0,
sd_sampler=SDSampler.dpm_plus_plus,
sd_sampler=SDSampler.dpm_plus_plus_2m,
powerpaint_task="outpainting",
)

63
lama_cleaner/tests/test_sd_model.py
View File

@@ -1,5 +1,7 @@
import os
from loguru import logger
from lama_cleaner.tests.utils import check_device, get_config, assert_equal
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
@@ -17,21 +19,7 @@ save_dir.mkdir(exist_ok=True, parents=True)
@pytest.mark.parametrize("device", ["cuda", "mps"])
@pytest.mark.parametrize(
"sampler",
[
SDSampler.ddim,
SDSampler.pndm,
SDSampler.k_lms,
SDSampler.k_euler,
SDSampler.k_euler_a,
SDSampler.lcm,
],
)
def test_runway_sd_1_5_all_samplers(
device,
sampler,
):
def test_runway_sd_1_5_all_samplers(device):
sd_steps = check_device(device)
model = ModelManager(
name="runwayml/stable-diffusion-inpainting",
@@ -39,22 +27,37 @@ def test_runway_sd_1_5_all_samplers(
disable_nsfw=True,
sd_cpu_textencoder=False,
)
cfg = get_config(
strategy=HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench",
sd_steps=sd_steps,
)
cfg.sd_sampler = sampler
name = f"device_{device}_{sampler}"
all_samplers = [member.value for member in SDSampler.__members__.values()]
print(all_samplers)
for sampler in all_samplers:
print(f"Testing sampler {sampler}")
if (
sampler
in [SDSampler.dpm2_karras, SDSampler.dpm2_a_karras, SDSampler.lms_karras]
and device == "mps"
):
# diffusers 0.25.0 still has bug on these sampler on mps, wait main branch released to fix it
logger.warning(
"skip dpm2_karras on mps, diffusers does not support it on mps. TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead."
)
continue
cfg = get_config(
strategy=HDStrategy.ORIGINAL,
prompt="a fox sitting on a bench",
sd_steps=sd_steps,
sd_sampler=sampler,
)
assert_equal(
model,
cfg,
f"runway_sd_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
name = f"device_{device}_{sampler}"
assert_equal(
model,
cfg,
f"runway_sd_{name}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("device", ["cuda", "mps", "cpu"])
@@ -171,7 +174,7 @@ def test_runway_norm_sd_model(device, strategy, sampler):
@pytest.mark.parametrize("device", ["cuda"])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
@pytest.mark.parametrize("sampler", [SDSampler.dpm_plus_plus_2m])
def test_runway_sd_1_5_cpu_offload(device, strategy, sampler):
sd_steps = check_device(device)
model = ModelManager(

8
lama_cleaner/tests/utils.py
View File

@@ -3,7 +3,9 @@ import cv2
import pytest
import torch
from lama_cleaner.helper import encode_pil_to_base64
from lama_cleaner.schema import LDMSampler, HDStrategy, InpaintRequest, SDSampler
from PIL import Image
current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / "result"
@@ -21,7 +23,7 @@ def check_device(device: str) -> int:
def assert_equal(
model,
config,
config: InpaintRequest,
gt_name,
fx: float = 1,
fy: float = 1,
@@ -29,6 +31,8 @@ def assert_equal(
mask_p=current_dir / "mask.png",
):
img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
config.image = encode_pil_to_base64(Image.fromarray(img), 95, {})[0]
config.mask = encode_pil_to_base64(Image.fromarray(mask), 95, {})[0]
print(f"Input image shape: {img.shape}")
res = model(img, mask, config)
ok = cv2.imwrite(
@@ -72,4 +76,4 @@ def get_config(**kwargs):
hd_strategy_resize_limit=200,
)
data.update(**kwargs)
return InpaintRequest(**data)
return InpaintRequest(image="", mask="", **data)