fix error I think

styleTransfer.py

@@ -66,41 +66,42 @@ vgg.eval()
 
 new_state_dict = OrderedDict()
 state_dict = torch.load(decoder_path)
-for k, v in state_dict.items():
-    #namekey = k[7:] # remove `module.`
-    namekey = k
-    new_state_dict[namekey] = v
-decoder.load_state_dict(new_state_dict)
+# for k, v in state_dict.items():
+#     #namekey = k[7:] # remove `module.`
+#     namekey = k
+#     new_state_dict[namekey] = v
+decoder.load_state_dict(state_dict)
 
 new_state_dict = OrderedDict()
 state_dict = torch.load(Trans_path)
-for k, v in state_dict.items():
-    #namekey = k[7:] # remove `module.`
-    namekey = k
-    new_state_dict[namekey] = v
-Trans.load_state_dict(new_state_dict)
+# for k, v in state_dict.items():
+#     #namekey = k[7:] # remove `module.`
+#     namekey = k
+#     new_state_dict[namekey] = v
+Trans.load_state_dict(state_dict)
 
 new_state_dict = OrderedDict()
 state_dict = torch.load(embedding_path)
-for k, v in state_dict.items():
-    #namekey = k[7:] # remove `module.`
-    namekey = k
-    new_state_dict[namekey] = v
-embedding.load_state_dict(new_state_dict)
+# for k, v in state_dict.items():
+#     #namekey = k[7:] # remove `module.`
+#     namekey = k
+#     new_state_dict[namekey] = v
+embedding.load_state_dict(state_dict)
 
 network = StyTR.StyTrans(vgg,decoder,embedding,Trans)
 network.eval()
-network.to(device)
-
-content_tf = test_transform(content_size, crop)
-style_tf = test_transform(style_size, crop)
 
 def StyleTransformer(content_img: Image, style_img: Image):
+   
+    network.to(device)
+    # content_tf = test_transform(content_size, crop)
+    # style_tf = test_transform(style_size, crop)
+
     content_tf1 = content_transform()       
-    content = content_tf(content_img.convert("RGB"))
+    content = content_tf1(content_img.convert("RGB"))
     h,w,c=np.shape(content)    
     style_tf1 = style_transform(h,w)
-    style = style_tf(style_img.convert("RGB"))
+    style = style_tf1(style_img.convert("RGB"))
     style = style.to(device).unsqueeze(0)
     content = content.to(device).unsqueeze(0)
 
@@ -125,6 +126,6 @@ def StyleGAN(content_image, style_image):
 
 
 def create_styledSofa(sofa:Image, style:Image):
-    styled_sofa = StyleGAN(sofa,style)
-    #styled_sofa = StyleTransformer(sofa,style)
+    #styled_sofa = StyleGAN(sofa,style)
+    styled_sofa = StyleTransformer(sofa,style)
     return styled_sofa