File size: 1,787 Bytes
cbbdd92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import os
from iopaint.schema import InpaintRequest
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
import torch
from iopaint.model_manager import ModelManager
def test_model_switch():
model = ModelManager(
name="runwayml/stable-diffusion-inpainting",
enable_controlnet=True,
controlnet_method="lllyasviel/control_v11p_sd15_canny",
device=torch.device("mps"),
disable_nsfw=True,
sd_cpu_textencoder=True,
cpu_offload=False,
)
model.switch("lama")
def test_controlnet_switch_onoff(caplog):
name = "runwayml/stable-diffusion-inpainting"
model = ModelManager(
name=name,
enable_controlnet=True,
controlnet_method="lllyasviel/control_v11p_sd15_canny",
device=torch.device("mps"),
disable_nsfw=True,
sd_cpu_textencoder=True,
cpu_offload=False,
)
model.switch_controlnet_method(
InpaintRequest(
name=name,
enable_controlnet=False,
)
)
assert "Disable controlnet" in caplog.text
def test_switch_controlnet_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,
enable_controlnet=True,
controlnet_method=old_method,
device=torch.device("mps"),
disable_nsfw=True,
sd_cpu_textencoder=True,
cpu_offload=False,
)
model.switch_controlnet_method(
InpaintRequest(
name=name,
enable_controlnet=True,
controlnet_method=new_method,
)
)
assert f"Switch Controlnet method from {old_method} to {new_method}" in caplog.text
|