IOPaint/lama_cleaner/tests/test_sd_model.py

205 lines
6.4 KiB
Python
Raw Normal View History

2022-11-15 14:09:51 +01:00
from pathlib import Path
import pytest
import torch
from lama_cleaner.model_manager import ModelManager
2022-12-11 13:27:32 +01:00
from lama_cleaner.schema import HDStrategy, SDSampler
2022-11-15 14:09:51 +01:00
from lama_cleaner.tests.test_model import get_config, assert_equal
current_dir = Path(__file__).parent.absolute().resolve()
2023-02-07 14:00:19 +01:00
save_dir = current_dir / "result"
2022-11-15 14:09:51 +01:00
save_dir.mkdir(exist_ok=True, parents=True)
2023-02-07 14:00:19 +01:00
device = "cuda" if torch.cuda.is_available() else "cpu"
2022-11-15 14:09:51 +01:00
device = torch.device(device)
2023-02-07 14:00:19 +01:00
@pytest.mark.parametrize("sd_device", ["cuda"])
2022-11-15 14:09:51 +01:00
@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])
2023-02-07 14:00:19 +01:00
def test_runway_sd_1_5_ddim(
sd_device, strategy, sampler, cpu_textencoder, disable_nsfw
):
2022-11-15 14:09:51 +01:00
def callback(i, t, latents):
2023-02-07 14:00:19 +01:00
pass
2022-11-15 14:09:51 +01:00
2023-02-07 14:00:19 +01:00
if sd_device == "cuda" and not torch.cuda.is_available():
2022-11-15 14:09:51 +01:00
return
2023-02-07 14:00:19 +01:00
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)
2022-11-15 14:09:51 +01:00
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",
2023-02-07 14:00:19 +01:00
fx=1.3,
2022-11-15 14:09:51 +01:00
)
2023-02-07 14:00:19 +01:00
@pytest.mark.parametrize("sd_device", ["cuda"])
2022-11-15 14:09:51 +01:00
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
2023-02-07 14:00:19 +01:00
@pytest.mark.parametrize(
"sampler", [SDSampler.pndm, SDSampler.k_lms, SDSampler.k_euler, SDSampler.k_euler_a]
)
2022-11-15 14:09:51 +01:00
@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}")
2023-02-07 14:00:19 +01:00
if sd_device == "cuda" and not torch.cuda.is_available():
2022-11-15 14:09:51 +01:00
return
2023-02-07 14:00:19 +01:00
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)
2022-11-15 14:09:51 +01:00
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",
2023-02-07 14:00:19 +01:00
fx=1.3,
2022-11-15 14:09:51 +01:00
)
2023-02-07 14:00:19 +01:00
@pytest.mark.parametrize("sd_device", ["cuda"])
2022-11-15 14:09:51 +01:00
@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
2023-02-07 14:00:19 +01:00
if sd_device == "cuda" and not torch.cuda.is_available():
2022-11-15 14:09:51 +01:00
return
2023-02-07 14:00:19 +01:00
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=False,
sd_cpu_textencoder=False,
callback=callback,
)
2022-11-15 14:09:51 +01:00
cfg = get_config(
strategy,
sd_steps=sd_steps,
2023-02-07 14:00:19 +01:00
prompt="Face of a fox, high resolution, sitting on a park bench",
negative_prompt="orange, yellow, small",
2022-11-25 02:29:20 +01:00
sd_sampler=sampler,
2023-02-07 14:00:19 +01:00
sd_match_histograms=True,
2022-11-15 14:09:51 +01:00
)
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",
2023-02-07 14:00:19 +01:00
fx=1,
2022-11-15 14:09:51 +01:00
)
2023-01-05 15:07:39 +01:00
2023-02-07 14:00:19 +01:00
@pytest.mark.parametrize("sd_device", ["cuda"])
2023-01-05 15:07:39 +01:00
@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])
2023-02-07 14:00:19 +01:00
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():
2023-01-05 15:07:39 +01:00
return
2023-02-07 14:00:19 +01:00
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,
)
cfg = get_config(
strategy, prompt="a fox sitting on a bench", sd_steps=sd_steps, sd_scale=0.85
)
2023-01-05 15:07:39 +01:00
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",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
2023-02-07 14:00:19 +01:00
fx=1.3,
2023-01-05 15:07:39 +01:00
)
2023-02-07 14:00:19 +01:00
@pytest.mark.parametrize("sd_device", ["cuda"])
2023-01-05 15:07:39 +01:00
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
2023-02-07 14:00:19 +01:00
if sd_device == "cuda" and not torch.cuda.is_available():
2023-01-05 15:07:39 +01:00
return
2023-02-07 14:00:19 +01:00
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=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
)
2023-01-05 15:07:39 +01:00
cfg.sd_sampler = sampler
name = f"device_{sd_device}_{sampler}"
assert_equal(
model,
cfg,
f"runway_sd_{strategy.capitalize()}_{name}_cpu_offload.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
2023-01-07 01:52:11 +01:00