Upload folder using huggingface_hub

app.py

@@ -33,30 +33,33 @@ optimal_settings = {
     'Watercolor': (75, False),
 }
 
-cached_style_features = {}
-for style_name, style_img_path in style_options.items():
-    style_img_512 = preprocess_img_from_path(style_img_path, 512)[0].to(device)
-    style_img_1024 = preprocess_img_from_path(style_img_path, 1024)[0].to(device)
-    with torch.no_grad():
-        style_features = (model(style_img_512), model(style_img_1024))
-    cached_style_features[style_name] = style_features
-
 def gram_matrix(feature):
     batch_size, n_feature_maps, height, width = feature.size()
-    return torch.mm(
-        feature.view(batch_size * n_feature_maps, height * width), 
-        feature.view(batch_size * n_feature_maps, height * width).t()
-    )
+    new_feature = feature.view(batch_size * n_feature_maps, height * width)
+    return torch.mm(new_feature, new_feature.t())
 
-def compute_loss(generated_features, content_features, style_features, alpha, beta):
+cached_style_gram_matrices = {}
+for style_name, style_img_path in tqdm(style_options.items(), desc='Computing style gram matrices'):
+    style_img_512 = preprocess_img_from_path(style_img_path, 512)[0].to(device)
+    style_img_1024 = preprocess_img_from_path(style_img_path, 1024)[0].to(device)
+    with torch.no_grad():
+        style_features_512 = model(style_img_512)
+        style_features_1024 = model(style_img_1024)
+        # compute gram matrices
+        style_gram_matrices_512 = [gram_matrix(f) for f in style_features_512]
+        style_gram_matrices_1024 = [gram_matrix(f) for f in style_features_1024]
+    cached_style_gram_matrices[style_name] = (style_gram_matrices_512, style_gram_matrices_1024)
+print('Style caching complete.')
+
+def compute_loss(generated_features, content_features, style_gram_matrices, alpha, beta):
     content_loss = 0
     style_loss = 0
     
-    for generated_feature, content_feature, style_feature in zip(generated_features, content_features, style_features):
+    for generated_feature, content_feature, style_gram_matrix in zip(generated_features, content_features, style_gram_matrices):
         content_loss += torch.mean((generated_feature - content_feature) ** 2)
         
         G = gram_matrix(generated_feature)
-        A = gram_matrix(style_feature)
+        A = style_gram_matrix
         
         E_l = ((G - A) ** 2)
         w_l = 1 / 5
@@ -91,13 +94,13 @@ def inference(content_image, style_name, style_strength, output_quality, progres
     with torch.no_grad():
         content_features = model(content_img)
     
-    style_features = cached_style_features[style_name][0 if img_size == 512 else 1]
+    style_gram_matrices = cached_style_gram_matrices[style_name][0 if img_size == 512 else 1]
     
     for _ in tqdm(range(iters), desc='The magic is happening ✨'):
         optimizer.zero_grad()
 
         generated_features = model(generated_img)
-        total_loss = compute_loss(generated_features, content_features, style_features, alpha, beta)
+        total_loss = compute_loss(generated_features, content_features, style_gram_matrices, alpha, beta)
 
         total_loss.backward()
         optimizer.step()