Update app.py

app.py

@@ -1,66 +1,66 @@
 import gradio as gr
-# import torch
-# import torch.optim as optim
-# import torchvision.models as models
-# import torchvision.transforms as transforms
-# from PIL import Image
-# import numpy as np
-# import requests
-# from io import BytesIO
+import torch
+import torch.optim as optim
+import torchvision.models as models
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+import requests
+from io import BytesIO
 
-# Load VGG19 model
-# vgg = models.vgg19(pretrained=True).features
-# for param in vgg.parameters():
-#     param.requires_grad_(False)
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# vgg.to(device)
+Load VGG19 model
+vgg = models.vgg19(pretrained=True).features
+for param in vgg.parameters():
+    param.requires_grad_(False)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+vgg.to(device)
 
-# # Helper functions (load_image, im_convert, get_features, gram_matrix)
-# # ... (Include the helper functions you provided earlier here)
+# Helper functions (load_image, im_convert, get_features, gram_matrix)
+# ... (Include the helper functions you provided earlier here)
 
-# def style_transfer(content_image, style_image, alpha, beta, conv1_1, conv2_1, conv3_1, conv4_1, conv5_1, steps):
-#     content = load_image(content_image).to(device)
-#     style = load_image(style_image, shape=content.shape[-2:]).to(device)
+def style_transfer(content_image, style_image, alpha, beta, conv1_1, conv2_1, conv3_1, conv4_1, conv5_1, steps):
+    content = load_image(content_image).to(device)
+    style = load_image(style_image, shape=content.shape[-2:]).to(device)
     
-#     content_features = get_features(content, vgg)
-#     style_features = get_features(style, vgg)
-#     style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}
+    content_features = get_features(content, vgg)
+    style_features = get_features(style, vgg)
+    style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}
     
-#     target = content.clone().requires_grad_(True).to(device)
+    target = content.clone().requires_grad_(True).to(device)
     
-#     style_weights = {
-#         'conv1_1': conv1_1,
-#         'conv2_1': conv2_1,
-#         'conv3_1': conv3_1,
-#         'conv4_1': conv4_1,
-#         'conv5_1': conv5_1
-#     }
+    style_weights = {
+        'conv1_1': conv1_1,
+        'conv2_1': conv2_1,
+        'conv3_1': conv3_1,
+        'conv4_1': conv4_1,
+        'conv5_1': conv5_1
+    }
     
-#     content_weight = alpha
-#     style_weight = beta * 1e6
+    content_weight = alpha
+    style_weight = beta * 1e6
     
-#     optimizer = optim.Adam([target], lr=0.003)
+    optimizer = optim.Adam([target], lr=0.003)
     
-#     for ii in range(1, steps+1):
-#         target_features = get_features(target, vgg)
-#         content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)
+    for ii in range(1, steps+1):
+        target_features = get_features(target, vgg)
+        content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)
         
-#         style_loss = 0
-#         for layer in style_weights:
-#             target_feature = target_features[layer]
-#             target_gram = gram_matrix(target_feature)
-#             _, d, h, w = target_feature.shape
-#             style_gram = style_grams[layer]
-#             layer_style_loss = style_weights[layer] * torch.mean((target_gram - style_gram)**2)
-#             style_loss += layer_style_loss / (d * h * w)
+        style_loss = 0
+        for layer in style_weights:
+            target_feature = target_features[layer]
+            target_gram = gram_matrix(target_feature)
+            _, d, h, w = target_feature.shape
+            style_gram = style_grams[layer]
+            layer_style_loss = style_weights[layer] * torch.mean((target_gram - style_gram)**2)
+            style_loss += layer_style_loss / (d * h * w)
         
-#         total_loss = content_weight * content_loss + style_weight * style_loss
+        total_loss = content_weight * content_loss + style_weight * style_loss
         
-#         optimizer.zero_grad()
-#         total_loss.backward()
-#         optimizer.step()
+        optimizer.zero_grad()
+        total_loss.backward()
+        optimizer.step()
     
-#     return im_convert(target)
+    return im_convert(target)
 
 # Example images
 examples = [