optimize sd/paint_by_example modle VRAM usage

This commit is contained in:
Qing
2023-01-18 18:34:10 +08:00
parent 384f16dcd0
commit 148e97e8da
7 changed files with 61 additions and 24 deletions

View File

@@ -37,7 +37,7 @@ def assert_equal(
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
def test_paint_by_example(strategy):
model = ModelManager(name="paint_by_example", device=device)
model = ModelManager(name="paint_by_example", device=device, disable_nsfw=True)
cfg = get_config(strategy, paint_by_example_steps=30)
assert_equal(
model,
@@ -50,9 +50,22 @@ 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)
assert_equal(
model,
cfg,
f"paint_by_example_{strategy.capitalize()}_disable_nsfw.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)
model = ModelManager(name="paint_by_example", device=device, disable_nsfw=True)
cfg = get_config(strategy, paint_by_example_steps=30, sd_scale=0.85)
assert_equal(
model,
@@ -67,7 +80,7 @@ def test_paint_by_example_sd_scale(strategy):
@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)
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,
@@ -75,14 +88,12 @@ def test_paint_by_example_cpu_offload(strategy):
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",
fy=0.9,
fx=1.3
)
@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)
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,

View File

@@ -14,7 +14,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = torch.device(device)
@pytest.mark.parametrize("sd_device", ['cpu', 'cuda'])
@pytest.mark.parametrize("sd_device", ['cuda'])
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
@pytest.mark.parametrize("sampler", [SDSampler.ddim])
@pytest.mark.parametrize("cpu_textencoder", [True, False])
@@ -31,7 +31,7 @@ def test_runway_sd_1_5_ddim(sd_device, strategy, sampler, cpu_textencoder, disab
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
sd_disable_nsfw=disable_nsfw,
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)
@@ -66,7 +66,7 @@ def test_runway_sd_1_5(sd_device, strategy, sampler, cpu_textencoder, disable_ns
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
sd_disable_nsfw=disable_nsfw,
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)
@@ -99,7 +99,7 @@ def test_runway_sd_1_5_negative_prompt(sd_device, strategy, sampler):
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
sd_disable_nsfw=False,
disable_nsfw=False,
sd_cpu_textencoder=False,
callback=callback)
cfg = get_config(
@@ -137,7 +137,7 @@ def test_runway_sd_1_5_sd_scale(sd_device, strategy, sampler, cpu_textencoder, d
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
sd_disable_nsfw=disable_nsfw,
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)
cfg.sd_sampler = sampler
@@ -166,7 +166,7 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
sd_disable_nsfw=False,
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)
@@ -191,7 +191,7 @@ def test_runway_sd_1_5_cpu_offload_cpu_device(sd_device, strategy, sampler):
device=torch.device(sd_device),
hf_access_token="",
sd_run_local=True,
sd_disable_nsfw=False,
disable_nsfw=False,
sd_cpu_textencoder=False,
cpu_offload=True)
cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=1, sd_scale=0.85)