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

wip mat float16

Browse Source
This commit is contained in:
Qing
2023-03-25 22:46:28 +08:00
parent 7e028c3908
commit eb304ba696
3 changed files with 63 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
lama_cleaner/model/mat.py
View File

@@ -402,9 +402,6 @@ class MappingNet(torch.nn.Module):
def forward( def forward(
self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False
): ):
import ipdb
ipdb.set_trace()
# Embed, normalize, and concat inputs. # Embed, normalize, and concat inputs.
x = None x = None
if self.z_dim > 0: if self.z_dim > 0:
@@ -678,9 +675,9 @@ class Conv2dLayerPartial(nn.Module):
stride=self.stride, stride=self.stride,
padding=self.padding, padding=self.padding,
) )
mask_ratio = self.slide_winsize / (update_mask + 1e-8) mask_ratio = self.slide_winsize / (update_mask.to(torch.float32) + 1e-8)
update_mask = torch.clamp(update_mask, 0, 1) # 0 or 1 update_mask = torch.clamp(update_mask, 0, 1) # 0 or 1
mask_ratio = torch.mul(mask_ratio, update_mask) mask_ratio = torch.mul(mask_ratio, update_mask).to(x.dtype)
x = self.conv(x) x = self.conv(x)
x = torch.mul(x, mask_ratio) x = torch.mul(x, mask_ratio)
return x, update_mask return x, update_mask
@@ -734,7 +731,7 @@ class WindowAttention(nn.Module):
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
""" """
B_, N, C = x.shape B_, N, C = x.shape
norm_x = F.normalize(x, p=2.0, dim=-1) norm_x = F.normalize(x, p=2.0, dim=-1, eps=torch.finfo(x.dtype).eps)
q = ( q = (
self.q(norm_x) self.q(norm_x)
.reshape(B_, N, self.num_heads, C // self.num_heads) .reshape(B_, N, self.num_heads, C // self.num_heads)
@@ -771,7 +768,6 @@ class WindowAttention(nn.Module):
).repeat(1, N, 1) ).repeat(1, N, 1)
attn = self.softmax(attn) attn = self.softmax(attn)
x = (attn @ v).transpose(1, 2).reshape(B_, N, C) x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
x = self.proj(x) x = self.proj(x)
return x, mask_windows return x, mask_windows
@@ -935,7 +931,9 @@ class SwinTransformerBlock(nn.Module):
) # nW*B, window_size*window_size, C ) # nW*B, window_size*window_size, C
else: else:
attn_windows, mask_windows = self.attn( attn_windows, mask_windows = self.attn(
x_windows, mask_windows, mask=self.calculate_mask(x_size).to(x.device) x_windows,
mask_windows,
mask=self.calculate_mask(x_size).to(x.dtype).to(x.device),
) # nW*B, window_size*window_size, C ) # nW*B, window_size*window_size, C
# merge windows # merge windows
@@ -1874,7 +1872,7 @@ MAT_MODEL_MD5 = os.environ.get("MAT_MODEL_MD5", "8ca927835fa3f5e21d65ffcb165377e
class MAT(InpaintModel): class MAT(InpaintModel):
name = "mat" name = "mat"
min_size = 512 min_size = 1024
pad_mod = 512 pad_mod = 512
pad_to_square = True pad_to_square = True
@@ -1890,9 +1888,9 @@ class MAT(InpaintModel):
img_resolution=512, img_resolution=512,
img_channels=3, img_channels=3,
mapping_kwargs={"torch_dtype": self.torch_dtype}, mapping_kwargs={"torch_dtype": self.torch_dtype},
) ).to(self.torch_dtype)
# fmt: off # fmt: off
self.model = load_model(G, MAT_MODEL_URL, device, MAT_MODEL_MD5).to(self.torch_dtype) self.model = load_model(G, MAT_MODEL_URL, device, MAT_MODEL_MD5)
self.z = torch.from_numpy(np.random.randn(1, G.z_dim)).to(self.torch_dtype).to(device) self.z = torch.from_numpy(np.random.randn(1, G.z_dim)).to(self.torch_dtype).to(device)
self.label = torch.zeros([1, self.model.c_dim], device=device).to(self.torch_dtype) self.label = torch.zeros([1, self.model.c_dim], device=device).to(self.torch_dtype)
# fmt: on # fmt: on

12
lama_cleaner/model/utils.py
View File

@@ -134,8 +134,10 @@ def timestep_embedding(device, timesteps, dim, max_period=10000, repeat_only=Fal
###### MAT and FcF ####### ###### MAT and FcF #######
def normalize_2nd_moment(x, dim=1, eps=1e-8): def normalize_2nd_moment(x, dim=1):
return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt() return (
x * (x.square().mean(dim=dim, keepdim=True) + torch.finfo(x.dtype).eps).rsqrt()
)
class EasyDict(dict): class EasyDict(dict):
@@ -460,7 +462,7 @@ def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1):
if f is None: if f is None:
f = torch.ones([1, 1], dtype=torch.float32, device=x.device) f = torch.ones([1, 1], dtype=torch.float32, device=x.device)
assert isinstance(f, torch.Tensor) and f.ndim in [1, 2] assert isinstance(f, torch.Tensor) and f.ndim in [1, 2]
assert f.dtype == torch.float32 and not f.requires_grad assert not f.requires_grad
batch_size, num_channels, in_height, in_width = x.shape batch_size, num_channels, in_height, in_width = x.shape
# upx, upy = _parse_scaling(up) # upx, upy = _parse_scaling(up)
# downx, downy = _parse_scaling(down) # downx, downy = _parse_scaling(down)
@@ -733,9 +735,7 @@ def conv2d_resample(
# Validate arguments. # Validate arguments.
assert isinstance(x, torch.Tensor) and (x.ndim == 4) assert isinstance(x, torch.Tensor) and (x.ndim == 4)
assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype) assert isinstance(w, torch.Tensor) and (w.ndim == 4) and (w.dtype == x.dtype)
assert f is None or ( assert f is None or (isinstance(f, torch.Tensor) and f.ndim in [1, 2])
isinstance(f, torch.Tensor) and f.ndim in [1, 2] and f.dtype == torch.float32
)
assert isinstance(up, int) and (up >= 1) assert isinstance(up, int) and (up >= 1)
assert isinstance(down, int) and (down >= 1) assert isinstance(down, int) and (down >= 1)
# assert isinstance(groups, int) and (groups >= 1), f"!!!!!! groups: {groups} isinstance(groups, int) {isinstance(groups, int)} {type(groups)}" # assert isinstance(groups, int) and (groups >= 1), f"!!!!!! groups: {groups} isinstance(groups, int) {isinstance(groups, int)} {type(groups)}"

154
lama_cleaner/tests/test_model.py
View File

@@ -9,13 +9,18 @@ from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import Config, HDStrategy, LDMSampler, SDSampler from lama_cleaner.schema import Config, HDStrategy, LDMSampler, SDSampler
current_dir = Path(__file__).parent.absolute().resolve() current_dir = Path(__file__).parent.absolute().resolve()
save_dir = current_dir / 'result' save_dir = current_dir / "result"
save_dir.mkdir(exist_ok=True, parents=True) save_dir.mkdir(exist_ok=True, parents=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu' device = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(device) device = torch.device(device)
def get_data(fx: float = 1, fy: float = 1.0, img_p=current_dir / "image.png", mask_p=current_dir / "mask.png"): def get_data(
fx: float = 1,
fy: float = 1.0,
img_p=current_dir / "image.png",
mask_p=current_dir / "mask.png",
):
img = cv2.imread(str(img_p)) img = cv2.imread(str(img_p))
img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB) img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
mask = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE) mask = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE)
@@ -37,10 +42,15 @@ def get_config(strategy, **kwargs):
return Config(**data) return Config(**data)
def assert_equal(model, config, gt_name, def assert_equal(
fx: float = 1, fy: float = 1, model,
config,
gt_name,
fx: float = 1,
fy: float = 1,
img_p=current_dir / "image.png", img_p=current_dir / "image.png",
mask_p=current_dir / "mask.png"): mask_p=current_dir / "mask.png",
):
img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p) img, mask = get_data(fx=fx, fy=fy, img_p=img_p, mask_p=mask_p)
print(f"Input image shape: {img.shape}") print(f"Input image shape: {img.shape}")
res = model(img, mask, config) res = model(img, mask, config)
@@ -59,139 +69,13 @@ def assert_equal(model, config, gt_name,
# assert np.array_equal(res, gt) # assert np.array_equal(res, gt)
@pytest.mark.parametrize( @pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
def test_lama(strategy):
model = ModelManager(name="lama", device=device)
assert_equal(
model,
get_config(strategy),
f"lama_{strategy[0].upper() + strategy[1:]}_result.png",
)
fx = 1.3
assert_equal(
model,
get_config(strategy),
f"lama_{strategy[0].upper() + strategy[1:]}_fx_{fx}_result.png",
fx=1.3,
)
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
@pytest.mark.parametrize("ldm_sampler", [LDMSampler.ddim, LDMSampler.plms])
def test_ldm(strategy, ldm_sampler):
model = ModelManager(name="ldm", device=device)
cfg = get_config(strategy, ldm_sampler=ldm_sampler)
assert_equal(
model, cfg, f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_result.png"
)
fx = 1.3
assert_equal(
model,
cfg,
f"ldm_{strategy[0].upper() + strategy[1:]}_{ldm_sampler}_fx_{fx}_result.png",
fx=fx,
)
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
@pytest.mark.parametrize("zits_wireframe", [False, True])
def test_zits(strategy, zits_wireframe):
model = ModelManager(name="zits", device=device)
cfg = get_config(strategy, zits_wireframe=zits_wireframe)
# os.environ['ZITS_DEBUG_LINE_PATH'] = str(current_dir / 'zits_debug_line.jpg')
# os.environ['ZITS_DEBUG_EDGE_PATH'] = str(current_dir / 'zits_debug_edge.jpg')
assert_equal(
model,
cfg,
f"zits_{strategy[0].upper() + strategy[1:]}_wireframe_{zits_wireframe}_result.png",
)
fx = 1.3
assert_equal(
model,
cfg,
f"zits_{strategy.capitalize()}_wireframe_{zits_wireframe}_fx_{fx}_result.png",
fx=fx,
)
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL]
)
def test_mat(strategy): def test_mat(strategy):
model = ModelManager(name="mat", device=device) model = ModelManager(name="mat", device=device)
cfg = get_config(strategy) cfg = get_config(strategy)
for _ in range(10):
assert_equal( assert_equal(
model, model, cfg, f"mat_{strategy.capitalize()}_result.png",
cfg,
f"mat_{strategy.capitalize()}_result.png",
) )
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL]
)
def test_fcf(strategy):
model = ModelManager(name="fcf", device=device)
cfg = get_config(strategy)
assert_equal(
model,
cfg,
f"fcf_{strategy.capitalize()}_result.png",
fx=2,
fy=2
)
assert_equal(
model,
cfg,
f"fcf_{strategy.capitalize()}_result.png",
fx=3.8,
fy=2
)
@pytest.mark.parametrize(
"strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP]
)
@pytest.mark.parametrize("cv2_flag", ['INPAINT_NS', 'INPAINT_TELEA'])
@pytest.mark.parametrize("cv2_radius", [3, 15])
def test_cv2(strategy, cv2_flag, cv2_radius):
model = ModelManager(
name="cv2",
device=torch.device(device),
)
cfg = get_config(strategy, cv2_flag=cv2_flag, cv2_radius=cv2_radius)
assert_equal(
model,
cfg,
f"sd_{strategy.capitalize()}_{cv2_flag}_{cv2_radius}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)
@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL, HDStrategy.RESIZE, HDStrategy.CROP])
def test_manga(strategy):
model = ModelManager(
name="manga",
device=torch.device(device),
)
cfg = get_config(strategy)
assert_equal(
model,
cfg,
f"sd_{strategy.capitalize()}.png",
img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
)