Update main.py
Browse files
main.py
CHANGED
@@ -162,18 +162,18 @@ class Model:
|
|
162 |
if instyle is None or aligned_face is None:
|
163 |
return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
|
164 |
if self.style_name != style_type:
|
165 |
-
exstyle, _
|
166 |
if exstyle is None:
|
167 |
return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
|
168 |
with torch.no_grad():
|
169 |
s_w = instyle.clone()
|
170 |
-
s_w[
|
171 |
|
172 |
x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
|
173 |
-
x_p = F.interpolate(self.parsingpredictor(2*(F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0],
|
174 |
scale_factor=0.5, recompute_scale_factor=False).detach()
|
175 |
-
inputs = torch.cat((x, x_p/16.), dim=1)
|
176 |
-
y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s=style_degree)
|
177 |
y_tilde = torch.clamp(y_tilde, -1, 1)
|
178 |
print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
|
179 |
return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s' % (self.style_name)
|
|
|
162 |
if instyle is None or aligned_face is None:
|
163 |
return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
|
164 |
if self.style_name != style_type:
|
165 |
+
exstyle, _ = self.load_model(style_type)
|
166 |
if exstyle is None:
|
167 |
return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.'
|
168 |
with torch.no_grad():
|
169 |
s_w = instyle.clone()
|
170 |
+
s_w[:, :7] = exstyle[:, :7]
|
171 |
|
172 |
x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
|
173 |
+
x_p = F.interpolate(self.parsingpredictor(2 * (F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0],
|
174 |
scale_factor=0.5, recompute_scale_factor=False).detach()
|
175 |
+
inputs = torch.cat((x, x_p / 16.), dim=1)
|
176 |
+
y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s=style_degree)
|
177 |
y_tilde = torch.clamp(y_tilde, -1, 1)
|
178 |
print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
|
179 |
return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s' % (self.style_name)
|