find next error

app.py

@@ -96,11 +96,15 @@ def style_sofa(input_img: np.ndarray, style_img: np.ndarray):
     resized_style = resize_style(style_img)
     resized_style.save('resized_style.jpg')
     # generate mask for image
+    print('generating mask...')
     mask = get_mask(resized_img)
     mask.save('mask.jpg')
+    # Created a styled sofa
+    print('Styling sofa...')
     styled_sofa = create_styledSofa(resized_img,resized_style)
     styled_sofa.save('styled_sofa.jpg')
     # postprocess the final image
+    print('Replacing sofa...')
     new_sofa = replace_sofa(resized_img,mask,styled_sofa)
     new_sofa = new_sofa.crop(box)
     return new_sofa

styleTransfer.py

@@ -20,7 +20,6 @@ embedding_path  = 'embedding_iter_160000.pth'
 
 def style_transform(h,w):
     k = (h,w)
-    size = int(np.max(k))
     transform_list = []    
     transform_list.append(transforms.CenterCrop((h,w)))
     transform_list.append(transforms.ToTensor())
@@ -37,8 +36,6 @@ def content_transform():
 content_size=640
 style_size=640
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 vgg = StyTR.vgg
 vgg.load_state_dict(torch.load(vgg_path))
 vgg = nn.Sequential(*list(vgg.children())[:44])
@@ -70,14 +67,12 @@ content_tf = content_transform()
 style_tf   = style_transform(style_size,style_size)
 
 def StyleTransformer(content_img: Image, style_img: Image):
-   
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     network.to(device)
-      
     content = content_tf(content_img.convert("RGB"))   
     style = style_tf(style_img.convert("RGB"))
     style = style.to(device).unsqueeze(0)
     content = content.to(device).unsqueeze(0)
-
     with torch.no_grad():
         output= network(content,style)  
     output = output[0].cpu().squeeze()