Upload folder using huggingface_hub

app.py

@@ -49,13 +49,13 @@ def compute_loss(generated_features, content_features, style_features, alpha, be
     for generated_feature, content_feature, style_feature in zip(generated_features, content_features, style_features):
         batch_size, n_feature_maps, height, width = generated_feature.size()
         
-        content_loss += (torch.mean((generated_feature - content_feature) ** 2))
+        content_loss += torch.mean((generated_feature - content_feature) ** 2)
         
         G = torch.mm((generated_feature.view(batch_size * n_feature_maps, height * width)), (generated_feature.view(batch_size * n_feature_maps, height * width)).t())
         A = torch.mm((style_feature.view(batch_size * n_feature_maps, height * width)), (style_feature.view(batch_size * n_feature_maps, height * width)).t())
         
         E_l = ((G - A) ** 2)
-        w_l = 1/5
+        w_l = 1 / 5
         style_loss += torch.mean(w_l * E_l)
 
     return alpha * content_loss + beta * style_loss
@@ -86,8 +86,8 @@ def inference(content_image, style_name, style_strength, output_quality, progres
 
     with torch.no_grad():
         content_features = model(content_img)
-    if img_size == 512: style_features = cached_style_features[style_name][0]
-    else: style_features = cached_style_features[style_name][1]
+    
+    style_features = cached_style_features[style_name][0 if img_size == 512 else 1]
     
     for _ in tqdm(range(iters), desc='The magic is happening ✨'):
         optimizer.zero_grad()

utils.py

@@ -3,18 +3,11 @@ from PIL import Image
 import torch
 import torchvision.transforms as transforms
 
-def preprocess_img(img: Image, img_size):
-    original_size = img.size
-    
-    transform = transforms.Compose([
-        transforms.Resize((img_size, img_size)),
-        transforms.ToTensor()
-    ])
-    img = transform(img).unsqueeze(0)
-    return img, original_size
-
 def preprocess_img_from_path(path_to_image, img_size):
     img = Image.open(path_to_image)
+    return preprocess_img(img, img_size)
+
+def preprocess_img(img: Image, img_size):
     original_size = img.size
     
     transform = transforms.Compose([
@@ -25,8 +18,7 @@ def preprocess_img_from_path(path_to_image, img_size):
     return img, original_size
 
 def postprocess_img(img, original_size):
-    img = img.cpu().clone()
-    img = img.squeeze(0)
+    img = img.detach().cpu().squeeze(0)
     
     # address tensor value scaling and quantization
     img = torch.clamp(img, 0, 1)