add missing files

lineart_models/__init__.py

@@ -0,0 +1,3 @@
+from .lineart import LineartDetector
+from .lineart_anime import LineartAnimeDetector
+from .mangaline_preprocessor import MangaLineExtraction

lineart_models/lineart.py

@@ -0,0 +1,146 @@
+# From https://github.com/carolineec/informative-drawings
+# MIT License
+'''
+MIT License
+
+Copyright (c) 2022 Caroline Chan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+'''
+
+import os
+import cv2
+import torch
+import numpy as np
+
+import torch.nn as nn
+from einops import rearrange
+from .utils import load_file_from_url
+
+
+norm_layer = nn.InstanceNorm2d
+
+
+class ResidualBlock(nn.Module):
+ def __init__(self, in_features):
+ super(ResidualBlock, self).__init__()
+
+ conv_block = [ nn.ReflectionPad2d(1),
+ nn.Conv2d(in_features, in_features, 3),
+ norm_layer(in_features),
+ nn.ReLU(inplace=True),
+ nn.ReflectionPad2d(1),
+ nn.Conv2d(in_features, in_features, 3),
+ norm_layer(in_features)
+ ]
+
+ self.conv_block = nn.Sequential(*conv_block)
+
+ def forward(self, x):
+ return x + self.conv_block(x)
+
+
+class Generator(nn.Module):
+ def __init__(self, input_nc, output_nc, n_residual_blocks=9, sigmoid=True):
+ super(Generator, self).__init__()
+
+ # Initial convolution block
+ model0 = [ nn.ReflectionPad2d(3),
+ nn.Conv2d(input_nc, 64, 7),
+ norm_layer(64),
+ nn.ReLU(inplace=True) ]
+ self.model0 = nn.Sequential(*model0)
+
+ # Downsampling
+ model1 = []
+ in_features = 64
+ out_features = in_features*2
+ for _ in range(2):
+ model1 += [ nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+ norm_layer(out_features),
+ nn.ReLU(inplace=True) ]
+ in_features = out_features
+ out_features = in_features*2
+ self.model1 = nn.Sequential(*model1)
+
+ model2 = []
+ # Residual blocks
+ for _ in range(n_residual_blocks):
+ model2 += [ResidualBlock(in_features)]
+ self.model2 = nn.Sequential(*model2)
+
+ # Upsampling
+ model3 = []
+ out_features = in_features//2
+ for _ in range(2):
+ model3 += [ nn.ConvTranspose2d(in_features, out_features, 3, stride=2, padding=1, output_padding=1),
+ norm_layer(out_features),
+ nn.ReLU(inplace=True) ]
+ in_features = out_features
+ out_features = in_features//2
+ self.model3 = nn.Sequential(*model3)
+
+ # Output layer
+ model4 = [ nn.ReflectionPad2d(3),
+ nn.Conv2d(64, output_nc, 7)]
+ if sigmoid:
+ model4 += [nn.Sigmoid()]
+
+ self.model4 = nn.Sequential(*model4)
+
+ def forward(self, x, cond=None):
+ out = self.model0(x)
+ out = self.model1(out)
+ out = self.model2(out)
+ out = self.model3(out)
+ out = self.model4(out)
+
+ return out
+
+
+class LineartDetector:
+ def __init__(self, model_path="hf_download"):
+ self.model = self.load_model('sk_model.pth', model_path)
+ self.model_coarse = self.load_model('sk_model2.pth', model_path)
+
+ def load_model(self, name, model_path="hf_download"):
+ remote_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/" + name
+ modelpath = os.path.join(model_path, name)
+ if not os.path.exists(modelpath):
+ load_file_from_url(remote_model_path, model_dir=model_path)
+ model = Generator(3, 1, 3)
+ model.load_state_dict(torch.load(modelpath, map_location=torch.device('cpu')))
+ model.eval()
+ model = model.cuda()
+ return model
+
+ def __call__(self, input_image, coarse=False):
+ model = self.model_coarse if coarse else self.model
+ assert input_image.ndim == 3
+ image = input_image
+ with torch.no_grad():
+ image = torch.from_numpy(image).float().cuda()
+ image = image / 255.0
+ image = rearrange(image, 'h w c -> 1 c h w')
+ line = model(image)[0][0]
+
+ line = line.cpu().numpy()
+ line = (line * 255.0).clip(0, 255).astype(np.uint8)
+
+ return line

lineart_models/lineart_anime.py

@@ -0,0 +1,226 @@
+# Anime2sketch
+# https://github.com/Mukosame/Anime2Sketch
+'''
+MIT License
+
+Copyright (c) 2022 Caroline Chan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+'''
+
+import numpy as np
+import torch
+import torch.nn as nn
+import functools
+
+import os
+import cv2
+from einops import rearrange
+
+
+class UnetGenerator(nn.Module):
+ """Create a Unet-based generator"""
+
+ def __init__(
+ self,
+ input_nc,
+ output_nc,
+ num_downs,
+ ngf=64,
+ norm_layer=nn.BatchNorm2d,
+ use_dropout=False,
+ ):
+ """Construct a Unet generator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ output_nc (int) -- the number of channels in output images
+ num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+ image of size 128x128 will become of size 1x1 # at the bottleneck
+ ngf (int) -- the number of filters in the last conv layer
+ norm_layer -- normalization layer
+ We construct the U-Net from the innermost layer to the outermost layer.
+ It is a recursive process.
+ """
+ super(UnetGenerator, self).__init__()
+ # construct unet structure
+ unet_block = UnetSkipConnectionBlock(
+ ngf * 8,
+ ngf * 8,
+ input_nc=None,
+ submodule=None,
+ norm_layer=norm_layer,
+ innermost=True,
+ ) # add the innermost layer
+ for _ in range(num_downs - 5): # add intermediate layers with ngf * 8 filters
+ unet_block = UnetSkipConnectionBlock(
+ ngf * 8,
+ ngf * 8,
+ input_nc=None,
+ submodule=unet_block,
+ norm_layer=norm_layer,
+ use_dropout=use_dropout,
+ )
+ # gradually reduce the number of filters from ngf * 8 to ngf
+ unet_block = UnetSkipConnectionBlock(
+ ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+ )
+ unet_block = UnetSkipConnectionBlock(
+ ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+ )
+ unet_block = UnetSkipConnectionBlock(
+ ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer
+ )
+ self.model = UnetSkipConnectionBlock(
+ output_nc,
+ ngf,
+ input_nc=input_nc,
+ submodule=unet_block,
+ outermost=True,
+ norm_layer=norm_layer,
+ ) # add the outermost layer
+
+ def forward(self, input):
+ """Standard forward"""
+ return self.model(input)
+
+
+class UnetSkipConnectionBlock(nn.Module):
+ """Defines the Unet submodule with skip connection.
+ X -------------------identity----------------------
+ |-- downsampling -- |submodule| -- upsampling --|
+ """
+
+ def __init__(
+ self,
+ outer_nc,
+ inner_nc,
+ input_nc=None,
+ submodule=None,
+ outermost=False,
+ innermost=False,
+ norm_layer=nn.BatchNorm2d,
+ use_dropout=False,
+ ):
+ """Construct a Unet submodule with skip connections.
+ Parameters:
+ outer_nc (int) -- the number of filters in the outer conv layer
+ inner_nc (int) -- the number of filters in the inner conv layer
+ input_nc (int) -- the number of channels in input images/features
+ submodule (UnetSkipConnectionBlock) -- previously defined submodules
+ outermost (bool) -- if this module is the outermost module
+ innermost (bool) -- if this module is the innermost module
+ norm_layer -- normalization layer
+ use_dropout (bool) -- if use dropout layers.
+ """
+ super(UnetSkipConnectionBlock, self).__init__()
+ self.outermost = outermost
+ if type(norm_layer) == functools.partial:
+ use_bias = norm_layer.func == nn.InstanceNorm2d
+ else:
+ use_bias = norm_layer == nn.InstanceNorm2d
+ if input_nc is None:
+ input_nc = outer_nc
+ downconv = nn.Conv2d(
+ input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias
+ )
+ downrelu = nn.LeakyReLU(0.2, True)
+ downnorm = norm_layer(inner_nc)
+ uprelu = nn.ReLU(True)
+ upnorm = norm_layer(outer_nc)
+
+ if outermost:
+ upconv = nn.ConvTranspose2d(
+ inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1
+ )
+ down = [downconv]
+ up = [uprelu, upconv, nn.Tanh()]
+ model = down + [submodule] + up
+ elif innermost:
+ upconv = nn.ConvTranspose2d(
+ inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias
+ )
+ down = [downrelu, downconv]
+ up = [uprelu, upconv, upnorm]
+ model = down + up
+ else:
+ upconv = nn.ConvTranspose2d(
+ inner_nc * 2,
+ outer_nc,
+ kernel_size=4,
+ stride=2,
+ padding=1,
+ bias=use_bias,
+ )
+ down = [downrelu, downconv, downnorm]
+ up = [uprelu, upconv, upnorm]
+
+ if use_dropout:
+ model = down + [submodule] + up + [nn.Dropout(0.5)]
+ else:
+ model = down + [submodule] + up
+
+ self.model = nn.Sequential(*model)
+
+ def forward(self, x):
+ if self.outermost:
+ return self.model(x)
+ else: # add skip connections
+ return torch.cat([x, self.model(x)], 1)
+
+
+class LineartAnimeDetector:
+ def __init__(self, model_path="hf_download"):
+ remote_model_path = (
+ "https://huggingface.co/lllyasviel/Annotators/resolve/main/netG.pth"
+ )
+ modelpath = os.path.join(model_path, "netG.pth")
+ if not os.path.exists(modelpath):
+ from .utils import load_file_from_url
+
+ load_file_from_url(remote_model_path, model_dir=model_path)
+ norm_layer = functools.partial(
+ nn.InstanceNorm2d, affine=False, track_running_stats=False
+ )
+ net = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)
+ ckpt = torch.load(modelpath)
+ for key in list(ckpt.keys()):
+ if "module." in key:
+ ckpt[key.replace("module.", "")] = ckpt[key]
+ del ckpt[key]
+ net.load_state_dict(ckpt)
+ net = net.cuda()
+ net.eval()
+ self.model = net
+
+ def __call__(self, input_image):
+ H, W, C = input_image.shape
+ Hn = 256 * int(np.ceil(float(H) / 256.0))
+ Wn = 256 * int(np.ceil(float(W) / 256.0))
+ img = cv2.resize(input_image, (Wn, Hn), interpolation=cv2.INTER_CUBIC)
+ with torch.no_grad():
+ image_feed = torch.from_numpy(img).float().cuda()
+ image_feed = image_feed / 127.5 - 1.0
+ image_feed = rearrange(image_feed, "h w c -> 1 c h w")
+
+ line = self.model(image_feed)[0, 0] * 127.5 + 127.5
+ line = line.cpu().numpy()
+
+ line = cv2.resize(line, (W, H), interpolation=cv2.INTER_CUBIC)
+ line = line.clip(0, 255).astype(np.uint8)
+ return line

lineart_models/mangaline_preprocessor.py

@@ -0,0 +1,356 @@
+'''
+MIT License
+
+Copyright (c) 2021 Miaomiao Li
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+'''
+import os
+
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from .utils import load_file_from_url
+
+
+class _bn_relu_conv(nn.Module):
+ def __init__(self, in_filters, nb_filters, fw, fh, subsample=1):
+ super(_bn_relu_conv, self).__init__()
+ self.model = nn.Sequential(
+ nn.BatchNorm2d(in_filters, eps=1e-3),
+ nn.LeakyReLU(0.2),
+ nn.Conv2d(
+ in_filters,
+ nb_filters,
+ (fw, fh),
+ stride=subsample,
+ padding=(fw // 2, fh // 2),
+ padding_mode="zeros",
+ ),
+ )
+
+ def forward(self, x):
+ return self.model(x)
+
+ # the following are for debugs
+ print(
+ "****",
+ np.max(x.cpu().numpy()),
+ np.min(x.cpu().numpy()),
+ np.mean(x.cpu().numpy()),
+ np.std(x.cpu().numpy()),
+ x.shape,
+ )
+ for i, layer in enumerate(self.model):
+ if i != 2:
+ x = layer(x)
+ else:
+ x = layer(x)
+ # x = nn.functional.pad(x, (1, 1, 1, 1), mode='constant', value=0)
+ print(
+ "____",
+ np.max(x.cpu().numpy()),
+ np.min(x.cpu().numpy()),
+ np.mean(x.cpu().numpy()),
+ np.std(x.cpu().numpy()),
+ x.shape,
+ )
+ print(x[0])
+ return x
+
+
+class _u_bn_relu_conv(nn.Module):
+ def __init__(self, in_filters, nb_filters, fw, fh, subsample=1):
+ super(_u_bn_relu_conv, self).__init__()
+ self.model = nn.Sequential(
+ nn.BatchNorm2d(in_filters, eps=1e-3),
+ nn.LeakyReLU(0.2),
+ nn.Conv2d(
+ in_filters,
+ nb_filters,
+ (fw, fh),
+ stride=subsample,
+ padding=(fw // 2, fh // 2),
+ ),
+ nn.Upsample(scale_factor=2, mode="nearest"),
+ )
+
+ def forward(self, x):
+ return self.model(x)
+
+
+class _shortcut(nn.Module):
+ def __init__(self, in_filters, nb_filters, subsample=1):
+ super(_shortcut, self).__init__()
+ self.process = False
+ self.model = None
+ if in_filters != nb_filters or subsample != 1:
+ self.process = True
+ self.model = nn.Sequential(
+ nn.Conv2d(in_filters, nb_filters, (1, 1), stride=subsample)
+ )
+
+ def forward(self, x, y):
+ # print(x.size(), y.size(), self.process)
+ if self.process:
+ y0 = self.model(x)
+ # print("merge+", torch.max(y0+y), torch.min(y0+y),torch.mean(y0+y), torch.std(y0+y), y0.shape)
+ return y0 + y
+ else:
+ # print("merge", torch.max(x+y), torch.min(x+y),torch.mean(x+y), torch.std(x+y), y.shape)
+ return x + y
+
+
+class _u_shortcut(nn.Module):
+ def __init__(self, in_filters, nb_filters, subsample):
+ super(_u_shortcut, self).__init__()
+ self.process = False
+ self.model = None
+ if in_filters != nb_filters:
+ self.process = True
+ self.model = nn.Sequential(
+ nn.Conv2d(
+ in_filters,
+ nb_filters,
+ (1, 1),
+ stride=subsample,
+ padding_mode="zeros",
+ ),
+ nn.Upsample(scale_factor=2, mode="nearest"),
+ )
+
+ def forward(self, x, y):
+ if self.process:
+ return self.model(x) + y
+ else:
+ return x + y
+
+
+class basic_block(nn.Module):
+ def __init__(self, in_filters, nb_filters, init_subsample=1):
+ super(basic_block, self).__init__()
+ self.conv1 = _bn_relu_conv(
+ in_filters, nb_filters, 3, 3, subsample=init_subsample
+ )
+ self.residual = _bn_relu_conv(nb_filters, nb_filters, 3, 3)
+ self.shortcut = _shortcut(in_filters, nb_filters, subsample=init_subsample)
+
+ def forward(self, x):
+ x1 = self.conv1(x)
+ x2 = self.residual(x1)
+ return self.shortcut(x, x2)
+
+
+class _u_basic_block(nn.Module):
+ def __init__(self, in_filters, nb_filters, init_subsample=1):
+ super(_u_basic_block, self).__init__()
+ self.conv1 = _u_bn_relu_conv(
+ in_filters, nb_filters, 3, 3, subsample=init_subsample
+ )
+ self.residual = _bn_relu_conv(nb_filters, nb_filters, 3, 3)
+ self.shortcut = _u_shortcut(in_filters, nb_filters, subsample=init_subsample)
+
+ def forward(self, x):
+ y = self.residual(self.conv1(x))
+ return self.shortcut(x, y)
+
+
+class _residual_block(nn.Module):
+ def __init__(self, in_filters, nb_filters, repetitions, is_first_layer=False):
+ super(_residual_block, self).__init__()
+ layers = []
+ for i in range(repetitions):
+ init_subsample = 1
+ if i == repetitions - 1 and not is_first_layer:
+ init_subsample = 2
+ if i == 0:
+ l = basic_block(
+ in_filters=in_filters,
+ nb_filters=nb_filters,
+ init_subsample=init_subsample,
+ )
+ else:
+ l = basic_block(
+ in_filters=nb_filters,
+ nb_filters=nb_filters,
+ init_subsample=init_subsample,
+ )
+ layers.append(l)
+
+ self.model = nn.Sequential(*layers)
+
+ def forward(self, x):
+ return self.model(x)
+
+
+class _upsampling_residual_block(nn.Module):
+ def __init__(self, in_filters, nb_filters, repetitions):
+ super(_upsampling_residual_block, self).__init__()
+ layers = []
+ for i in range(repetitions):
+ l = None
+ if i == 0:
+ l = _u_basic_block(
+ in_filters=in_filters, nb_filters=nb_filters
+ ) # (input)
+ else:
+ l = basic_block(in_filters=nb_filters, nb_filters=nb_filters) # (input)
+ layers.append(l)
+
+ self.model = nn.Sequential(*layers)
+
+ def forward(self, x):
+ return self.model(x)
+
+
+class res_skip(nn.Module):
+ def __init__(self):
+ super(res_skip, self).__init__()
+ self.block0 = _residual_block(
+ in_filters=1, nb_filters=24, repetitions=2, is_first_layer=True
+ ) # (input)
+ self.block1 = _residual_block(
+ in_filters=24, nb_filters=48, repetitions=3
+ ) # (block0)
+ self.block2 = _residual_block(
+ in_filters=48, nb_filters=96, repetitions=5
+ ) # (block1)
+ self.block3 = _residual_block(
+ in_filters=96, nb_filters=192, repetitions=7
+ ) # (block2)
+ self.block4 = _residual_block(
+ in_filters=192, nb_filters=384, repetitions=12
+ ) # (block3)
+
+ self.block5 = _upsampling_residual_block(
+ in_filters=384, nb_filters=192, repetitions=7
+ ) # (block4)
+ self.res1 = _shortcut(
+ in_filters=192, nb_filters=192
+ ) # (block3, block5, subsample=(1,1))
+
+ self.block6 = _upsampling_residual_block(
+ in_filters=192, nb_filters=96, repetitions=5
+ ) # (res1)
+ self.res2 = _shortcut(
+ in_filters=96, nb_filters=96
+ ) # (block2, block6, subsample=(1,1))
+
+ self.block7 = _upsampling_residual_block(
+ in_filters=96, nb_filters=48, repetitions=3
+ ) # (res2)
+ self.res3 = _shortcut(
+ in_filters=48, nb_filters=48
+ ) # (block1, block7, subsample=(1,1))
+
+ self.block8 = _upsampling_residual_block(
+ in_filters=48, nb_filters=24, repetitions=2
+ ) # (res3)
+ self.res4 = _shortcut(
+ in_filters=24, nb_filters=24
+ ) # (block0,block8, subsample=(1,1))
+
+ self.block9 = _residual_block(
+ in_filters=24, nb_filters=16, repetitions=2, is_first_layer=True
+ ) # (res4)
+ self.conv15 = _bn_relu_conv(
+ in_filters=16, nb_filters=1, fh=1, fw=1, subsample=1
+ ) # (block7)
+
+ def forward(self, x):
+ x0 = self.block0(x)
+ x1 = self.block1(x0)
+ x2 = self.block2(x1)
+ x3 = self.block3(x2)
+ x4 = self.block4(x3)
+
+ x5 = self.block5(x4)
+ res1 = self.res1(x3, x5)
+
+ x6 = self.block6(res1)
+ res2 = self.res2(x2, x6)
+
+ x7 = self.block7(res2)
+ res3 = self.res3(x1, x7)
+
+ x8 = self.block8(res3)
+ res4 = self.res4(x0, x8)
+
+ x9 = self.block9(res4)
+ y = self.conv15(x9)
+
+ return y
+
+
+class MangaLineExtraction:
+ def __init__(self, device=None, model_dir=None):
+ self.model = None
+ self.device = device
+ MangaLineExtraction.model_dir = model_dir
+
+ def load_model(self):
+ remote_model_path = (
+ "https://huggingface.co/lllyasviel/Annotators/resolve/main/erika.pth"
+ )
+ modelpath = os.path.join(self.model_dir, "erika.pth")
+ if not os.path.exists(modelpath):
+ load_file_from_url(remote_model_path, model_dir=self.model_dir)
+ # norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
+ net = res_skip()
+ ckpt = torch.load(modelpath)
+ for key in list(ckpt.keys()):
+ if "module." in key:
+ ckpt[key.replace("module.", "")] = ckpt[key]
+ del ckpt[key]
+ net.load_state_dict(ckpt)
+ net.eval()
+ self.model = net.to(self.device)
+
+ def unload_model(self):
+ if self.model is not None:
+ self.model.cpu()
+
+ def __call__(self, input_image):
+ if self.model is None:
+ self.load_model()
+ self.model.to(self.device)
+ # if width or height is not divisible by 16, pad the image
+ h, w = input_image.shape[:2]
+ # get adjusted pixel amount to max 1280x1280
+ total_pixels = h * w
+ if total_pixels > 1280 * 1280:
+ ratio = (1280 * 1280) / total_pixels
+ ratio = ratio**0.5
+ h = int(h * ratio)
+ w = int(w * ratio)
+ divisible = 16
+ h = h + (divisible - h % divisible) % divisible
+ w = w + (divisible - w % divisible) % divisible
+ input_image = cv2.resize(input_image, (w, h))
+ img = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY)
+ img = np.ascontiguousarray(img.copy()).copy()
+ with torch.no_grad():
+ image_feed = torch.from_numpy(img).float().to(self.device)
+ image_feed = rearrange(image_feed, "h w -> 1 1 h w")
+ line = self.model(image_feed).cpu().numpy()[0, 0]
+ # line = 255 - line
+ return line.clip(0, 255).astype(np.uint8)

lineart_models/utils.py

@@ -0,0 +1,39 @@
+import os
+
+import cv2
+import numpy as np
+from urllib.parse import urlparse
+
+
+def load_file_from_url(
+ url: str,
+ *,
+ model_dir: str,
+ progress: bool = True,
+ file_name: str | None = None,
+) -> str:
+ """Download a file from `url` into `model_dir`, using the file present if possible.
+
+ Returns the path to the downloaded file.
+ """
+ os.makedirs(model_dir, exist_ok=True)
+ if not file_name:
+ parts = urlparse(url)
+ file_name = os.path.basename(parts.path)
+ cached_file = os.path.abspath(os.path.join(model_dir, file_name))
+ if not os.path.exists(cached_file):
+ print(f'Downloading: "{url}" to {cached_file}\n')
+ from torch.hub import download_url_to_file
+
+ download_url_to_file(url, cached_file, progress=progress)
+ return cached_file
+
+
+def combine_linearts(lineart1: np.ndarray, lineart2: np.ndarray, erode=[False, False]) -> np.ndarray:
+ if erode[0]:
+ lineart1 = cv2.erode(lineart1, np.ones((3, 3), np.uint8))
+ if erode[1]:
+ lineart2 = cv2.erode(lineart2, np.ones((3, 3), np.uint8))
+ # unify the dark part of lineart1 and lineart2
+ union = np.where(lineart1 < lineart2, lineart1, lineart2)
+ return union