John Guerrerio commited on
Commit
b947961
1 Parent(s): 5792adf

first take at deployment

Browse files
Files changed (3) hide show
  1. .gitignore +2 -0
  2. app.py +113 -0
  3. requirements.txt +4 -0
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ generated_*
2
+ venv
app.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torch.optim as optim
4
+ import torchvision.transforms as transforms
5
+ import torchvision.models as models
6
+ from torchvision.utils import save_image
7
+ import numpy as np
8
+ import gradio as gr
9
+
10
+ IMAGE_SIZE = 244 # VGG image input size - we use VGG 19 as our pretrained CNN
11
+
12
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
13
+
14
+ cnn = models.vgg19(weights=None)
15
+ state_dict = torch.load("vgg19-dcbb9e9d.pth")
16
+ cnn.load_state_dict(state_dict)
17
+
18
+ class VGG(nn.Module):
19
+ def __init__(self):
20
+ super(VGG, self).__init__()
21
+ self.layers = ['0', '5', '10', '19', '28'] # layers we use as representations
22
+ self.model = cnn.features[:29] # we don't care about later layers
23
+
24
+ def forward(self, x):
25
+ features = []
26
+
27
+ for layer_num, layer in enumerate(self.model):
28
+ x = layer(x)
29
+
30
+ # we don't care about the model output - we care about the output of individual layers
31
+ if str(layer_num) in self.layers:
32
+ features.append(x)
33
+
34
+ return features
35
+
36
+
37
+ gradio_transforms = transforms.Compose([
38
+ transforms.ToTensor(),
39
+ transforms.Resize([IMAGE_SIZE, IMAGE_SIZE])
40
+ ])
41
+
42
+
43
+ def sanitize_inputs(epochs, lr, cl, sl):
44
+ if epochs < 1:
45
+ return ["Epochs must be positive", None]
46
+ if not isinstance(epochs, int):
47
+ return ["Epochs must be an integer", None]
48
+ if lr < 0:
49
+ return ["Learning rate must be positive", None]
50
+ if lr > 1:
51
+ return ["Learning rate must be less than one", None]
52
+ if cl < 0 or cl > 1:
53
+ return ["Content loss weight must be between 0 and 1", None]
54
+ if sl < 0 or sl > 1:
55
+ return ["Style loss weight must be between 0 and 1", None]
56
+
57
+ return None
58
+
59
+ def train(Epochs, Learning_Rate, Content_Loss, Style_Loss, Content_Image, Style_Image):
60
+ errors = sanitize_inputs(Epochs, Learning_Rate, Content_Loss, Style_Loss)
61
+ if errors is not None:
62
+ return errors
63
+
64
+ test = Content_Image
65
+
66
+ content = gradio_transforms(Content_Image).unsqueeze(0).to(device)
67
+ style = gradio_transforms(Style_Image).unsqueeze(0).to(device)
68
+ generated = content.clone().requires_grad_(True).to(device)
69
+
70
+ model = VGG().to(device).eval()
71
+ optimizer = optim.Adam([generated], lr=Learning_Rate)
72
+
73
+ for epoch in range(Epochs):
74
+ generatedFeatures = model(generated)
75
+ contentFeatures = model(content)
76
+ styleFeatures = model(style)
77
+
78
+ styleLoss = 0
79
+ contentLoss = 0
80
+
81
+ for genFeat, contFeat, styleFeat in zip(generatedFeatures, contentFeatures, styleFeatures):
82
+ batch_size, channel, height, width = genFeat.shape
83
+
84
+ contentLoss += torch.mean((genFeat - contFeat) ** 2)
85
+
86
+ G = genFeat.view(channel, height * width).mm(genFeat.view(channel, height * width).t())
87
+ A = styleFeat.view(channel, height * width).mm(styleFeat.view(channel, height * width).t())
88
+
89
+ styleLoss += torch.mean((G - A) ** 2)
90
+
91
+ total_loss = Content_Loss * contentLoss + Style_Loss * styleLoss
92
+ optimizer.zero_grad()
93
+ total_loss.backward()
94
+ optimizer.step()
95
+
96
+
97
+ save_image(generated, "generated_gradio.png")
98
+
99
+ return ["No errors! Enjoy your new image!", "generated_gradio.png"]
100
+
101
+ demo = gr.Interface(
102
+ fn=train,
103
+ inputs=["number", "number", "number", "number", "image", "image"],
104
+ outputs=[
105
+ gr.Label(label="Error Messages"),
106
+ gr.Image(label="Generated Image"),
107
+ ],
108
+ title="Neural Style Transfer",
109
+ description="Perform neural style transfer on images of your choice! Provide a content image that contains the content you want to transform and a style image that contains the style you want to emulate.\n\nNote: Huggingface requires users to pay to gain access to GPUs, so this model is hosted on a cpu. Training for many epochs will take a VERY long time. Using a larger learning rate (e.g., 0.01) can help reduce the number of epochs you need.",
110
+ theme=gr.themes.Soft()
111
+ )
112
+
113
+ demo.launch(debug=True, share=True)
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ torch==2.2.1
2
+ torchvision==0.17.1
3
+ matplotlib==3.7.1
4
+ numpy==1.25.2