diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx index d908c07..c3aefc2 100644 --- a/lama_cleaner/app/src/store/Atoms.tsx +++ b/lama_cleaner/app/src/store/Atoms.tsx @@ -231,6 +231,8 @@ export enum SDSampler { ddim = 'ddim', pndm = 'pndm', klms = 'k_lms', + kEuler = 'k_euler', + kEulerA = 'k_euler_a', } export enum SDMode { diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py index 08dcf63..eb32bb8 100644 --- a/lama_cleaner/model/sd.py +++ b/lama_cleaner/model/sd.py @@ -4,7 +4,8 @@ import PIL.Image import cv2 import numpy as np import torch -from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler +from diffusers import PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler, EulerDiscreteScheduler, \ + EulerAncestralDiscreteScheduler from loguru import logger from lama_cleaner.model.base import InpaintModel @@ -98,25 +99,27 @@ 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, + ) + if config.sd_sampler == SDSampler.ddim: scheduler = DDIMScheduler( - beta_start=0.00085, - beta_end=0.012, - beta_schedule="scaled_linear", + **scheduler_kwargs, clip_sample=False, set_alpha_to_one=False, ) elif config.sd_sampler == SDSampler.pndm: - PNDM_kwargs = { - "beta_schedule": "scaled_linear", - "beta_start": 0.00085, - "beta_end": 0.012, - "num_train_timesteps": 1000, - "skip_prk_steps": True, - } - scheduler = PNDMScheduler(**PNDM_kwargs) + scheduler = PNDMScheduler(**scheduler_kwargs, skip_prk_steps=True) elif config.sd_sampler == SDSampler.k_lms: - scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear") + scheduler = LMSDiscreteScheduler(**scheduler_kwargs) + elif config.sd_sampler == SDSampler.k_euler: + scheduler = EulerDiscreteScheduler(**scheduler_kwargs) + elif config.sd_sampler == SDSampler.k_euler_a: + scheduler = EulerAncestralDiscreteScheduler(**scheduler_kwargs) else: raise ValueError(config.sd_sampler) diff --git a/lama_cleaner/schema.py b/lama_cleaner/schema.py index a2569d1..4a374c5 100644 --- a/lama_cleaner/schema.py +++ b/lama_cleaner/schema.py @@ -18,6 +18,8 @@ class SDSampler(str, Enum): ddim = "ddim" pndm = "pndm" k_lms = "k_lms" + k_euler = 'k_euler' + k_euler_a = 'k_euler_a' class Config(BaseModel): diff --git a/lama_cleaner/tests/test_model.py b/lama_cleaner/tests/test_model.py index 2b24bfe..71cc7ef 100644 --- a/lama_cleaner/tests/test_model.py +++ b/lama_cleaner/tests/test_model.py @@ -161,79 +161,6 @@ def test_fcf(strategy): ) -@pytest.mark.parametrize("sd_device", ['cpu', 'cuda']) -@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) -@pytest.mark.parametrize("sampler", [SDSampler.ddim, SDSampler.pndm, SDSampler.k_lms]) -@pytest.mark.parametrize("cpu_textencoder", [True, False]) -@pytest.mark.parametrize("disable_nsfw", [True, False]) -def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_nsfw): - def callback(i, t, latents): - print(f"sd_step_{i}") - - if sd_device == 'cuda' and not torch.cuda.is_available(): - return - - sd_steps = 50 - model = ModelManager(name="sd1.5", - device=torch.device(sd_device), - hf_access_token="", - sd_run_local=True, - sd_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"{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']) -@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) -@pytest.mark.parametrize("sampler", [SDSampler.ddim]) -def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler): - def callback(i, t, latents): - pass - - if sd_device == 'cuda' and not torch.cuda.is_available(): - return - - sd_steps = 50 - model = ModelManager(name="sd1.5", - device=torch.device(sd_device), - hf_access_token="", - sd_run_local=True, - sd_disable_nsfw=True, - sd_cpu_textencoder=True, - callback=callback) - cfg = get_config( - strategy, - sd_steps=sd_steps, - prompt='Face of a fox, high resolution, sitting on a park bench', - negative_prompt='orange, yellow, small', - sd_sampler=sampler - ) - - name = f"{sampler}_negative_prompt" - - 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 - ) - - @pytest.mark.parametrize( "strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP] ) diff --git a/lama_cleaner/tests/test_sd_model.py b/lama_cleaner/tests/test_sd_model.py new file mode 100644 index 0000000..267a60d --- /dev/null +++ b/lama_cleaner/tests/test_sd_model.py @@ -0,0 +1,124 @@ +import os +from pathlib import Path + +import cv2 +import pytest +import torch + +from lama_cleaner.model_manager import ModelManager +from lama_cleaner.schema import Config, HDStrategy, LDMSampler, SDSampler +from lama_cleaner.tests.test_model import get_config, assert_equal + +current_dir = Path(__file__).parent.absolute().resolve() +save_dir = current_dir / 'result' +save_dir.mkdir(exist_ok=True, parents=True) +device = 'cuda' if torch.cuda.is_available() else 'cpu' +device = torch.device(device) + + +@pytest.mark.parametrize("sd_device", ['cpu', 'cuda']) +@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): + print(f"sd_step_{i}") + + 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, + sd_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']) +@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) +@pytest.mark.parametrize("sampler", [SDSampler.pndm, SDSampler.k_lms, SDSampler.k_euler, SDSampler.k_euler_a]) +@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}") + + 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, + sd_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']) +@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL]) +@pytest.mark.parametrize("sampler", [SDSampler.ddim]) +def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler): + def callback(i, t, latents): + pass + + if sd_device == 'cuda' and not torch.cuda.is_available(): + return + + sd_steps = 50 + model = ModelManager(name="sd1.5", + device=torch.device(sd_device), + hf_access_token="", + sd_run_local=True, + sd_disable_nsfw=True, + sd_cpu_textencoder=True, + callback=callback) + cfg = get_config( + strategy, + sd_steps=sd_steps, + prompt='Face of a fox, high resolution, sitting on a park bench', + negative_prompt='orange, yellow, small', + sd_sampler=sampler + ) + + name = f"{sampler}_negative_prompt" + + 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 + )