Upload 4 files

CaffeLoader.py

@@ -0,0 +1,254 @@
+import torch
+import torch.nn as nn
+
+
+class VGG(nn.Module):
+    def __init__(self, features, num_classes=1000):
+        super(VGG, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, num_classes),
+        )
+
+
+class VGG_SOD(nn.Module):
+    def __init__(self, features, num_classes=100):
+        super(VGG_SOD, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(),
+            nn.Linear(4096, 100),
+        )
+
+
+class VGG_FCN32S(nn.Module):
+    def __init__(self, features, num_classes=1000):
+        super(VGG_FCN32S, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Conv2d(512,4096,(7, 7)),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+            nn.Conv2d(4096,4096,(1, 1)),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+        )
+
+
+class VGG_PRUNED(nn.Module):
+    def __init__(self, features, num_classes=1000):
+        super(VGG_PRUNED, self).__init__()
+        self.features = features
+        self.classifier = nn.Sequential(
+            nn.Linear(512 * 7 * 7, 4096),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+            nn.Linear(4096, 4096),
+            nn.ReLU(True),
+            nn.Dropout(0.5),
+        )
+
+
+class NIN(nn.Module):
+    def __init__(self, pooling):
+        super(NIN, self).__init__()
+        if pooling == 'max':
+            pool2d = nn.MaxPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
+        elif pooling == 'avg':
+            pool2d = nn.AvgPool2d((3, 3),(2, 2),(0, 0),ceil_mode=True)
+
+        self.features = nn.Sequential(
+            nn.Conv2d(3,96,(11, 11),(4, 4)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(96,96,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(96,96,(1, 1)),
+            nn.ReLU(inplace=True),
+            pool2d,
+            nn.Conv2d(96,256,(5, 5),(1, 1),(2, 2)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256,256,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256,256,(1, 1)),
+            nn.ReLU(inplace=True),
+            pool2d,
+            nn.Conv2d(256,384,(3, 3),(1, 1),(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384,384,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384,384,(1, 1)),
+            nn.ReLU(inplace=True),
+            pool2d,
+            nn.Dropout(0.5),
+            nn.Conv2d(384,1024,(3, 3),(1, 1),(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024,1024,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024,1000,(1, 1)),
+            nn.ReLU(inplace=True),
+            nn.AvgPool2d((6, 6),(1, 1),(0, 0),ceil_mode=True),
+            nn.Softmax(),
+        )
+
+
+
+class ModelParallel(nn.Module):
+    def __init__(self, net, device_ids, device_splits):
+        super(ModelParallel, self).__init__()
+        self.device_list = self.name_devices(device_ids.split(','))
+        self.chunks = self.chunks_to_devices(self.split_net(net, device_splits.split(',')))
+
+    def name_devices(self, input_list):
+        device_list = []
+        for i, device in enumerate(input_list):
+            if str(device).lower() != 'c':
+                device_list.append("cuda:" + str(device))
+            else:
+                device_list.append("cpu")
+        return device_list
+
+    def split_net(self, net, device_splits):
+        chunks, cur_chunk = [], nn.Sequential()
+        for i, l in enumerate(net):
+            cur_chunk.add_module(str(i), net[i])
+            if str(i) in device_splits and device_splits != '':
+                del device_splits[0]
+                chunks.append(cur_chunk)
+                cur_chunk = nn.Sequential()
+        chunks.append(cur_chunk)
+        return chunks
+
+    def chunks_to_devices(self, chunks):
+        for i, chunk in enumerate(chunks):
+            chunk.to(self.device_list[i])
+        return chunks
+
+    def c(self, input, i):
+        if input.type() == 'torch.FloatTensor' and 'cuda' in self.device_list[i]:
+            input = input.type('torch.cuda.FloatTensor')
+        elif input.type() == 'torch.cuda.FloatTensor' and 'cpu' in self.device_list[i]:
+            input = input.type('torch.FloatTensor')
+        return input
+
+    def forward(self, input):
+        for i, chunk in enumerate(self.chunks):
+            if i < len(self.chunks) -1:
+                input = self.c(chunk(self.c(input, i).to(self.device_list[i])), i+1).to(self.device_list[i+1])
+            else:
+                input = chunk(input)
+        return input
+
+
+
+def buildSequential(channel_list, pooling):
+    layers = []
+    in_channels = 3
+    if pooling == 'max':
+        pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
+    elif pooling == 'avg':
+        pool2d = nn.AvgPool2d(kernel_size=2, stride=2)
+    else:
+        raise ValueError("Unrecognized pooling parameter")
+    for c in channel_list:
+        if c == 'P':
+            layers += [pool2d]
+        else:
+            conv2d = nn.Conv2d(in_channels, c, kernel_size=3, padding=1)
+            layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = c
+    return nn.Sequential(*layers)
+
+
+channel_list = {
+'VGG-16p': [24, 22, 'P', 41, 51, 'P', 108, 89, 111, 'P', 184, 276, 228, 'P', 512, 512, 512, 'P'],
+'VGG-16': [64, 64, 'P', 128, 128, 'P', 256, 256, 256, 'P', 512, 512, 512, 'P', 512, 512, 512, 'P'],
+'VGG-19': [64, 64, 'P', 128, 128, 'P', 256, 256, 256, 256, 'P', 512, 512, 512, 512, 'P', 512, 512, 512, 512, 'P'],
+}
+
+nin_dict = {
+'C': ['conv1', 'cccp1', 'cccp2', 'conv2', 'cccp3', 'cccp4', 'conv3', 'cccp5', 'cccp6', 'conv4-1024', 'cccp7-1024', 'cccp8-1024'],
+'R': ['relu0', 'relu1', 'relu2', 'relu3', 'relu5', 'relu6', 'relu7', 'relu8', 'relu9', 'relu10', 'relu11', 'relu12'],
+'P': ['pool1', 'pool2', 'pool3', 'pool4'],
+'D': ['drop'],
+}
+vgg16_dict = {
+'C': ['conv1_1', 'conv1_2', 'conv2_1', 'conv2_2', 'conv3_1', 'conv3_2', 'conv3_3', 'conv4_1', 'conv4_2', 'conv4_3', 'conv5_1', 'conv5_2', 'conv5_3'],
+'R': ['relu1_1', 'relu1_2', 'relu2_1', 'relu2_2', 'relu3_1', 'relu3_2', 'relu3_3', 'relu4_1', 'relu4_2', 'relu4_3', 'relu5_1', 'relu5_2', 'relu5_3'],
+'P': ['pool1', 'pool2', 'pool3', 'pool4', 'pool5'],
+}
+vgg19_dict = {
+'C': ['conv1_1', 'conv1_2', 'conv2_1', 'conv2_2', 'conv3_1', 'conv3_2', 'conv3_3', 'conv3_4', 'conv4_1', 'conv4_2', 'conv4_3', 'conv4_4', 'conv5_1', 'conv5_2', 'conv5_3', 'conv5_4'],
+'R': ['relu1_1', 'relu1_2', 'relu2_1', 'relu2_2', 'relu3_1', 'relu3_2', 'relu3_3', 'relu3_4', 'relu4_1', 'relu4_2', 'relu4_3', 'relu4_4', 'relu5_1', 'relu5_2', 'relu5_3', 'relu5_4'],
+'P': ['pool1', 'pool2', 'pool3', 'pool4', 'pool5'],
+}
+
+
+def modelSelector(model_file, pooling):
+    vgg_list = ["fcn32s", "pruning", "sod", "vgg"]
+    if any(name in model_file for name in vgg_list):
+        if "pruning" in model_file:
+            print("VGG-16 Architecture Detected")
+            print("Using The Channel Pruning Model")
+            cnn, layerList = VGG_PRUNED(buildSequential(channel_list['VGG-16p'], pooling)), vgg16_dict
+        elif "fcn32s" in model_file:
+            print("VGG-16 Architecture Detected")
+            print("Using the fcn32s-heavy-pascal Model")
+            cnn, layerList = VGG_FCN32S(buildSequential(channel_list['VGG-16'], pooling)), vgg16_dict
+        elif "sod" in model_file:
+            print("VGG-16 Architecture Detected")
+            print("Using The SOD Fintune Model")
+            cnn, layerList = VGG_SOD(buildSequential(channel_list['VGG-16'], pooling)), vgg16_dict
+        elif "19" in model_file:
+            print("VGG-19 Architecture Detected")
+            cnn, layerList = VGG(buildSequential(channel_list['VGG-19'], pooling)), vgg19_dict
+        elif "16" in model_file:
+            print("VGG-16 Architecture Detected")
+            cnn, layerList = VGG(buildSequential(channel_list['VGG-16'], pooling)), vgg16_dict
+        else:
+            raise ValueError("VGG architecture not recognized.")
+    elif "nin" in model_file:
+        print("NIN Architecture Detected")
+        cnn, layerList = NIN(pooling), nin_dict
+    else:
+        raise ValueError("Model architecture not recognized.")
+    return cnn, layerList
+
+
+# Print like Torch7/loadcaffe
+def print_loadcaffe(cnn, layerList):
+    c = 0
+    for l in list(cnn):
+         if "Conv2d" in str(l):
+             in_c, out_c, ks  = str(l.in_channels), str(l.out_channels), str(l.kernel_size)
+             print(layerList['C'][c] +": " +  (out_c + " " + in_c + " " + ks).replace(")",'').replace("(",'').replace(",",'') )
+             c+=1
+         if c == len(layerList['C']):
+             break
+
+
+# Load the model, and configure pooling layer type
+def loadCaffemodel(model_file, pooling, use_gpu, disable_check):
+    cnn, layerList = modelSelector(str(model_file).lower(), pooling)
+
+    cnn.load_state_dict(torch.load(model_file), strict=(not disable_check))
+    print("Successfully loaded " + str(model_file))
+
+    # Maybe convert the model to cuda now, to avoid later issues
+    if "c" not in str(use_gpu).lower() or "c" not in str(use_gpu[0]).lower():
+        cnn = cnn.cuda()
+    cnn = cnn.features
+
+    print_loadcaffe(cnn, layerList)
+
+    return cnn, layerList

README.md

@@ -1,10 +1,308 @@
----
-title: Style Transfer
-emoji: 🏃
-colorFrom: gray
-colorTo: blue
-sdk: static
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# neural-style-pt
+
+[![DOI](https://zenodo.org/badge/142345353.svg)](https://zenodo.org/badge/latestdoi/142345353)
+
+This is a PyTorch implementation of the paper [A Neural Algorithm of Artistic Style](http://arxiv.org/abs/1508.06576)
+by Leon A. Gatys, Alexander S. Ecker, and Matthias Bethge. The code is based on Justin Johnson's [Neural-Style](https://github.com/jcjohnson/neural-style).
+
+The paper presents an algorithm for combining the content of one image with the style of another image using
+convolutional neural networks. Here's an example that maps the artistic style of
+[The Starry Night](https://en.wikipedia.org/wiki/The_Starry_Night)
+onto a night-time photograph of the Stanford campus:
+
+<div align="center">
+ <img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/starry_night_google.jpg" height="223px">
+ <img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/hoovertowernight.jpg" height="223px">
+ <img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/starry_stanford_bigger.png" width="710px">
+</div>
+
+Applying the style of different images to the same content image gives interesting results.
+Here we reproduce Figure 2 from the paper, which renders a photograph of the Tubingen in Germany in a
+variety of styles:
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/tubingen.jpg" height="250px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_shipwreck.png" height="250px">
+
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_starry.png" height="250px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_scream.png" height="250px">
+
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_seated_nude.png" height="250px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_composition_vii.png" height="250px">
+</div>
+
+Here are the results of applying the style of various pieces of artwork to this photograph of the
+golden gate bridge:
+
+
+<div align="center"
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/golden_gate.jpg" height="200px">
+
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/frida_kahlo.jpg" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_kahlo.png" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/escher_sphere.jpg" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_escher.png" height="160px">
+</div>
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/woman-with-hat-matisse.jpg" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_matisse.png" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/the_scream.jpg" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_scream.png" height="160px">
+</div>
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/starry_night_crop.png" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry.png" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/seated-nude.jpg" height="160px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_seated.png" height="160px">
+</div>
+
+### Content / Style Tradeoff
+
+The algorithm allows the user to trade-off the relative weight of the style and content reconstruction terms,
+as shown in this example where we port the style of [Picasso's 1907 self-portrait](http://www.wikiart.org/en/pablo-picasso/self-portrait-1907) onto Brad Pitt:
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/picasso_selfport1907.jpg" height="220px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/inputs/brad_pitt.jpg" height="220px">
+</div>
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/pitt_picasso_content_5_style_100.png" height="220px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/pitt_picasso_content_1_style_100.png" height="220px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/pitt_picasso_content_01_style_100.png" height="220px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/pitt_picasso_content_0025_style_100.png" height="220px">
+</div>
+
+### Style Scale
+
+By resizing the style image before extracting style features, we can control the types of artistic
+features that are transfered from the style image; you can control this behavior with the `-style_scale` flag.
+Below we see three examples of rendering the Golden Gate Bridge in the style of The Starry Night.
+From left to right, `-style_scale` is 2.0, 1.0, and 0.5.
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry_scale2.png" height=175px>
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry_scale1.png" height=175px>
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry_scale05.png" height=175px>
+</div>
+
+### Multiple Style Images
+You can use more than one style image to blend multiple artistic styles.
+
+Clockwise from upper left: "The Starry Night" + "The Scream", "The Scream" + "Composition VII",
+"Seated Nude" + "Composition VII", and "Seated Nude" + "The Starry Night"
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_starry_scream.png" height="250px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_scream_composition_vii.png" height="250px">
+
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_starry_seated.png" height="250px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_seated_nude_composition_vii.png" height="250px">
+</div>
+
+
+### Style Interpolation
+When using multiple style images, you can control the degree to which they are blended:
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry_scream_3_7.png" height="175px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry_scream_5_5.png" height="175px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/golden_gate_starry_scream_7_3.png" height="175px">
+</div>
+
+### Transfer style but not color
+If you add the flag `-original_colors 1` then the output image will retain the colors of the original image.
+
+<div align="center">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_starry.png" height="185px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_scream.png" height="185px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/tubingen_composition_vii.png" height="185px">
+
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/original_color/tubingen_starry.png" height="185px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/original_color/tubingen_scream.png" height="185px">
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/original_color/tubingen_composition_vii.png" height="185px">
+</div>
+
+## Setup:
+
+Dependencies:
+* [PyTorch](http://pytorch.org/)
+
+
+Optional dependencies:
+* For CUDA backend:
+  * CUDA 7.5 or above
+* For cuDNN backend:
+  * cuDNN v6 or above
+* For ROCm backend:
+  * ROCm 2.1 or above
+* For MKL backend:
+  * MKL 2019 or above
+* For OpenMP backend:
+  * OpenMP 5.0 or above
+
+After installing the dependencies, you'll need to run the following script to download the VGG model:
+```
+python models/download_models.py
+```
+This will download the original [VGG-19 model](https://gist.github.com/ksimonyan/3785162f95cd2d5fee77#file-readme-md).
+The original [VGG-16 model](https://gist.github.com/ksimonyan/211839e770f7b538e2d8#file-readme-md) will also be downloaded. By default the original VGG-19 model is used.
+
+If you have a smaller memory GPU then using NIN Imagenet model will be better and gives slightly worse yet comparable results. You can get the details on the model from [BVLC Caffe ModelZoo](https://github.com/BVLC/caffe/wiki/Model-Zoo). The NIN model is downloaded when you run the `download_models.py` script.
+
+You can find detailed installation instructions for Ubuntu and Windows in the [installation guide](INSTALL.md).
+
+## Usage
+Basic usage:
+```
+python neural_style.py -style_image <image.jpg> -content_image <image.jpg>
+```
+
+cuDNN usage with NIN Model:
+```
+python neural_style.py -style_image examples/inputs/picasso_selfport1907.jpg -content_image examples/inputs/brad_pitt.jpg -output_image profile.png -model_file models/nin_imagenet.pth -gpu 0 -backend cudnn -num_iterations 1000 -seed 123 -content_layers relu0,relu3,relu7,relu12 -style_layers relu0,relu3,relu7,relu12 -content_weight 10 -style_weight 500 -image_size 512 -optimizer adam
+```
+
+![cuDNN NIN Model Picasso Brad Pitt](https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/pitt_picasso_nin_cudnn.png)
+
+
+To use multiple style images, pass a comma-separated list like this:
+
+`-style_image starry_night.jpg,the_scream.jpg`.
+
+Note that paths to images should not contain the `~` character to represent your home directory; you should instead use a relative
+path or a full absolute path.
+
+**Options**:
+* `-image_size`: Maximum side length (in pixels) of the generated image. Default is 512.
+* `-style_blend_weights`: The weight for blending the style of multiple style images, as a
+  comma-separated list, such as `-style_blend_weights 3,7`. By default all style images
+  are equally weighted.
+* `-gpu`: Zero-indexed ID of the GPU to use; for CPU mode set `-gpu` to `c`.
+
+**Optimization options**:
+* `-content_weight`: How much to weight the content reconstruction term. Default is 5e0.
+* `-style_weight`: How much to weight the style reconstruction term. Default is 1e2.
+* `-tv_weight`: Weight of total-variation (TV) regularization; this helps to smooth the image.
+  Default is 1e-3. Set to 0 to disable TV regularization.
+* `-num_iterations`: Default is 1000.
+* `-init`: Method for generating the generated image; one of `random` or `image`.
+  Default is `random` which uses a noise initialization as in the paper; `image`
+  initializes with the content image.
+* `-init_image`: Replaces the initialization image with a user specified image.
+* `-optimizer`: The optimization algorithm to use; either `lbfgs` or `adam`; default is `lbfgs`.
+  L-BFGS tends to give better results, but uses more memory. Switching to ADAM will reduce memory usage;
+  when using ADAM you will probably need to play with other parameters to get good results, especially
+  the style weight, content weight, and learning rate.
+* `-learning_rate`: Learning rate to use with the ADAM optimizer. Default is 1e1.
+* `-normalize_gradients`: If this flag is present, style and content gradients from each layer will be L1 normalized.
+
+**Output options**:
+* `-output_image`: Name of the output image. Default is `out.png`.
+* `-print_iter`: Print progress every `print_iter` iterations. Set to 0 to disable printing.
+* `-save_iter`: Save the image every `save_iter` iterations. Set to 0 to disable saving intermediate results.
+
+**Layer options**:
+* `-content_layers`: Comma-separated list of layer names to use for content reconstruction.
+  Default is `relu4_2`.
+* `-style_layers`: Comma-separated list of layer names to use for style reconstruction.
+  Default is `relu1_1,relu2_1,relu3_1,relu4_1,relu5_1`.
+
+**Other options**:
+* `-style_scale`: Scale at which to extract features from the style image. Default is 1.0.
+* `-original_colors`: If you set this to 1, then the output image will keep the colors of the content image.
+* `-model_file`: Path to the `.pth` file for the VGG Caffe model. Default is the original VGG-19 model; you can also try the original VGG-16 model.
+* `-pooling`: The type of pooling layers to use; one of `max` or `avg`. Default is `max`.
+  The VGG-19 models uses max pooling layers, but the paper mentions that replacing these layers with average
+  pooling layers can improve the results. I haven't been able to get good results using average pooling, but
+  the option is here.
+* `-seed`: An integer value that you can specify for repeatable results. By default this value is random for each run.
+* `-multidevice_strategy`: A comma-separated list of layer indices at which to split the network when using multiple devices. See [Multi-GPU scaling](https://github.com/ProGamerGov/neural-style-pt#multi-gpu-scaling) for more details.
+* `-backend`: `nn`, `cudnn`, `openmp`, or `mkl`. Default is `nn`. `mkl` requires Intel's MKL backend.
+* `-cudnn_autotune`: When using the cuDNN backend, pass this flag to use the built-in cuDNN autotuner to select
+  the best convolution algorithms for your architecture. This will make the first iteration a bit slower and can
+  take a bit more memory, but may significantly speed up the cuDNN backend.
+
+## Frequently Asked Questions
+
+**Problem:** The program runs out of memory and dies
+
+**Solution:** Try reducing the image size: `-image_size 256` (or lower). Note that different image sizes will likely
+require non-default values for `-style_weight` and `-content_weight` for optimal results.
+If you are running on a GPU, you can also try running with `-backend cudnn` to reduce memory usage.
+
+**Problem:** `-backend cudnn` is slower than default NN backend
+
+**Solution:** Add the flag `-cudnn_autotune`; this will use the built-in cuDNN autotuner to select the best convolution algorithms.
+
+**Problem:** Get the following error message:
+
+`Missing key(s) in state_dict: "classifier.0.bias", "classifier.0.weight", "classifier.3.bias", "classifier.3.weight".
+        Unexpected key(s) in state_dict: "classifier.1.weight", "classifier.1.bias", "classifier.4.weight", "classifier.4.bias".`
+
+**Solution:** Due to a mix up with layer locations, older models require a fix to be compatible with newer versions of PyTorch. The included [`donwload_models.py`](https://github.com/ProGamerGov/neural-style-pt/blob/master/models/download_models.py) script will automatically perform these fixes after downloading the models.
+
+
+
+## Memory Usage
+By default, `neural-style-pt` uses the `nn` backend for convolutions and L-BFGS for optimization. These give good results, but can both use a lot of memory. You can reduce memory usage with the following:
+
+* **Use cuDNN**: Add the flag `-backend cudnn` to use the cuDNN backend. This will only work in GPU mode.
+* **Use ADAM**: Add the flag `-optimizer adam` to use ADAM instead of L-BFGS. This should significantly
+  reduce memory usage, but may require tuning of other parameters for good results; in particular you should
+  play with the learning rate, content weight, and style weight.
+  This should work in both CPU and GPU modes.
+* **Reduce image size**: If the above tricks are not enough, you can reduce the size of the generated image;
+  pass the flag `-image_size 256` to generate an image at half the default size.
+
+With the default settings, neural-style-pt uses about 3.7 GB of GPU memory on my system; switching to ADAM and cuDNN reduces the GPU memory footprint to about 1GB.
+
+## Speed
+Speed can vary a lot depending on the backend and the optimizer.
+Here are some times for running 500 iterations with `-image_size=512` on a Tesla K80 with different settings:
+* `-backend nn -optimizer lbfgs`: 117 seconds
+* `-backend nn -optimizer adam`: 100 seconds
+* `-backend cudnn -optimizer lbfgs`: 124 seconds
+* `-backend cudnn -optimizer adam`: 107 seconds
+* `-backend cudnn -cudnn_autotune -optimizer lbfgs`: 109 seconds
+* `-backend cudnn -cudnn_autotune -optimizer adam`: 91 seconds
+
+Here are the same benchmarks on a GTX 1080:
+* `-backend nn -optimizer lbfgs`: 56 seconds
+* `-backend nn -optimizer adam`: 38 seconds
+* `-backend cudnn -optimizer lbfgs`: 40 seconds
+* `-backend cudnn -optimizer adam`: 40 seconds
+* `-backend cudnn -cudnn_autotune -optimizer lbfgs`: 23 seconds
+* `-backend cudnn -cudnn_autotune -optimizer adam`: 24 seconds
+
+## Multi-GPU scaling
+You can use multiple CPU and GPU devices to process images at higher resolutions; different layers of the network will be
+computed on different devices. You can control which GPU and CPU devices are used with the `-gpu` flag, and you can control
+how to split layers across devices using the `-multidevice_strategy` flag.
+
+For example in a server with four GPUs, you can give the flag `-gpu 0,1,2,3` to process on GPUs 0, 1, 2, and 3 in that order; by also giving the flag `-multidevice_strategy 3,6,12` you indicate that the first two layers should be computed on GPU 0, layers 3 to 5 should be computed on GPU 1, layers 6 to 11 should be computed on GPU 2, and the remaining layers should be computed on GPU 3. You will need to tune the `-multidevice_strategy` for your setup in order to achieve maximal resolution.
+
+We can achieve very high quality results at high resolution by combining multi-GPU processing with multiscale
+generation as described in the paper
+<a href="https://arxiv.org/abs/1611.07865">**Controlling Perceptual Factors in Neural Style Transfer**</a> by Leon A. Gatys,
+Alexander S. Ecker, Matthias Bethge, Aaron Hertzmann and Eli Shechtman.
+
+
+Here is a 4016 x 2213 image generated on a server with eight Tesla K80 GPUs:
+
+<img src="https://raw.githubusercontent.com/ProGamerGov/neural-style-pt/master/examples/outputs/starry_stanford_bigger.png" height="400px">
+
+The script used to generate this image <a href='examples/scripts/starry_stanford_bigger.sh'>can be found here</a>.
+
+## Implementation details
+Images are initialized with white noise and optimized using L-BFGS.
+
+We perform style reconstructions using the `conv1_1`, `conv2_1`, `conv3_1`, `conv4_1`, and `conv5_1` layers
+and content reconstructions using the `conv4_2` layer. As in the paper, the five style reconstruction losses have
+equal weights.
+
+## Citation
+
+If you find this code useful for your research, please cite it using the provided citation.

app.py

@@ -0,0 +1,37 @@
+from flask import Flask, render_template, request, redirect, url_for
+import os
+from werkzeug.utils import secure_filename
+from neural_style import main
+
+app = Flask(__name__)
+
+
+# 首页，用于上传图片和显示结果
+@app.route('/', methods=['GET', 'POST'])
+def upload_and_process():
+    if request.method == 'POST':
+        # 获取上传的图片列表
+        images = [request.files['image1'], request.files['image2']]
+        filenames = []
+
+        for image in images:
+            if image:
+                # 使用secure_filename获取安全的文件名
+                filename = secure_filename(image.filename)
+                print(filename)
+                # 保存上传的图片到本地
+                image.save(os.path.join('static', filename))
+                filenames.append(filename)
+
+        # 调用AI模型对图片进行处理（在这里，您需要编写AI模型的代码）
+        main(filenames[0], filenames[1])
+
+        # 返回结果页面并展示处理后的图片
+        return render_template('index.html', image_path="out.png", \
+                               filename1=filenames[0], filename2=filenames[1])
+
+    return render_template('index.html')
+
+
+if __name__ == '__main__':
+    app.run(debug=True)

neural_style.py

@@ -0,0 +1,514 @@
+import os
+import copy
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+
+from PIL import Image
+from CaffeLoader import loadCaffemodel, ModelParallel
+
+import argparse
+
+parser = argparse.ArgumentParser()
+# Basic options
+parser.add_argument("-style_image", help="Style target image", default='examples/inputs/seated-nude.jpg')
+parser.add_argument("-style_blend_weights", default=None)
+parser.add_argument("-content_image", help="Content target image", default='examples/inputs/tubingen.jpg')
+parser.add_argument("-image_size", help="Maximum height / width of generated image", type=int, default=512)
+parser.add_argument("-gpu", help="Zero-indexed ID of the GPU to use; for CPU mode set -gpu = c", default='c')
+
+# Optimization options
+parser.add_argument("-content_weight", type=float, default=5e0)
+parser.add_argument("-style_weight", type=float, default=1e2)
+parser.add_argument("-normalize_weights", action='store_true')
+parser.add_argument("-normalize_gradients", action='store_true')
+parser.add_argument("-tv_weight", type=float, default=1e-3)
+parser.add_argument("-num_iterations", type=int, default=200)
+parser.add_argument("-init", choices=['random', 'image'], default='random')
+parser.add_argument("-init_image", default=None)
+parser.add_argument("-optimizer", choices=['lbfgs', 'adam'], default='lbfgs')
+parser.add_argument("-learning_rate", type=float, default=1e0)
+parser.add_argument("-lbfgs_num_correction", type=int, default=100)
+
+# Output options
+parser.add_argument("-print_iter", type=int, default=50)
+parser.add_argument("-save_iter", type=int, default=100)
+parser.add_argument("-output_image", default='out.png')
+
+# Other options
+parser.add_argument("-style_scale", type=float, default=1.0)
+parser.add_argument("-original_colors", type=int, choices=[0, 1], default=0)
+parser.add_argument("-pooling", choices=['avg', 'max'], default='max')
+parser.add_argument("-model_file", type=str, default='models/vgg19-d01eb7cb.pth')
+parser.add_argument("-disable_check", action='store_true')
+parser.add_argument("-backend", choices=['nn', 'cudnn', 'mkl', 'mkldnn', 'openmp', 'mkl,cudnn', 'cudnn,mkl'],
+                    default='nn')
+parser.add_argument("-cudnn_autotune", action='store_true')
+parser.add_argument("-seed", type=int, default=-1)
+
+parser.add_argument("-content_layers", help="layers for content", default='relu4_2')
+parser.add_argument("-style_layers", help="layers for style", default='relu1_1,relu2_1,relu3_1,relu4_1,relu5_1')
+
+parser.add_argument("-multidevice_strategy", default='4,7,29')
+params = parser.parse_args()
+
+Image.MAX_IMAGE_PIXELS = 1000000000  # Support gigapixel images
+
+
+def main(filename1, filename2):
+    params.content_image = "C:/Users/86136/Desktop/web/ai2/neural-style-pt-master/static/" + filename1
+    params.style_image = "C:/Users/86136/Desktop/web/ai2/neural-style-pt-master/static/" + filename2
+    dtype, multidevice, backward_device = setup_gpu()
+
+    cnn, layerList = loadCaffemodel(params.model_file, params.pooling, params.gpu, params.disable_check)
+
+    content_image = preprocess(params.content_image, params.image_size).type(dtype)
+    style_image_input = params.style_image.split(',')
+    style_image_list, ext = [], [".jpg", ".jpeg", ".png", ".tiff"]
+    for image in style_image_input:
+        if os.path.isdir(image):
+            images = (image + "/" + file for file in os.listdir(image)
+                      if os.path.splitext(file)[1].lower() in ext)
+            style_image_list.extend(images)
+        else:
+            style_image_list.append(image)
+    style_images_caffe = []
+    for image in style_image_list:
+        style_size = int(params.image_size * params.style_scale)
+        img_caffe = preprocess(image, style_size).type(dtype)
+        style_images_caffe.append(img_caffe)
+
+    if params.init_image != None:
+        image_size = (content_image.size(2), content_image.size(3))
+        init_image = preprocess(params.init_image, image_size).type(dtype)
+
+    # Handle style blending weights for multiple style inputs
+    style_blend_weights = []
+    if params.style_blend_weights == None:
+        # Style blending not specified, so use equal weighting
+        for i in style_image_list:
+            style_blend_weights.append(1.0)
+        for i, blend_weights in enumerate(style_blend_weights):
+            style_blend_weights[i] = int(style_blend_weights[i])
+    else:
+        style_blend_weights = params.style_blend_weights.split(',')
+        assert len(style_blend_weights) == len(style_image_list), \
+            "-style_blend_weights and -style_images must have the same number of elements!"
+
+    # Normalize the style blending weights so they sum to 1
+    style_blend_sum = 0
+    for i, blend_weights in enumerate(style_blend_weights):
+        style_blend_weights[i] = float(style_blend_weights[i])
+        style_blend_sum = float(style_blend_sum) + style_blend_weights[i]
+    for i, blend_weights in enumerate(style_blend_weights):
+        style_blend_weights[i] = float(style_blend_weights[i]) / float(style_blend_sum)
+
+    content_layers = params.content_layers.split(',')
+    style_layers = params.style_layers.split(',')
+
+    # Set up the network, inserting style and content loss modules
+    cnn = copy.deepcopy(cnn)
+    content_losses, style_losses, tv_losses = [], [], []
+    next_content_idx, next_style_idx = 1, 1
+    net = nn.Sequential()
+    c, r = 0, 0
+    if params.tv_weight > 0:
+        tv_mod = TVLoss(params.tv_weight).type(dtype)
+        net.add_module(str(len(net)), tv_mod)
+        tv_losses.append(tv_mod)
+
+    for i, layer in enumerate(list(cnn), 1):
+        if next_content_idx <= len(content_layers) or next_style_idx <= len(style_layers):
+            if isinstance(layer, nn.Conv2d):
+                net.add_module(str(len(net)), layer)
+
+                if layerList['C'][c] in content_layers:
+                    print("Setting up content layer " + str(i) + ": " + str(layerList['C'][c]))
+                    loss_module = ContentLoss(params.content_weight, params.normalize_gradients)
+                    net.add_module(str(len(net)), loss_module)
+                    content_losses.append(loss_module)
+
+                if layerList['C'][c] in style_layers:
+                    print("Setting up style layer " + str(i) + ": " + str(layerList['C'][c]))
+                    loss_module = StyleLoss(params.style_weight, params.normalize_gradients)
+                    net.add_module(str(len(net)), loss_module)
+                    style_losses.append(loss_module)
+                c += 1
+
+            if isinstance(layer, nn.ReLU):
+                net.add_module(str(len(net)), layer)
+
+                if layerList['R'][r] in content_layers:
+                    print("Setting up content layer " + str(i) + ": " + str(layerList['R'][r]))
+                    loss_module = ContentLoss(params.content_weight, params.normalize_gradients)
+                    net.add_module(str(len(net)), loss_module)
+                    content_losses.append(loss_module)
+                    next_content_idx += 1
+
+                if layerList['R'][r] in style_layers:
+                    print("Setting up style layer " + str(i) + ": " + str(layerList['R'][r]))
+                    loss_module = StyleLoss(params.style_weight, params.normalize_gradients)
+                    net.add_module(str(len(net)), loss_module)
+                    style_losses.append(loss_module)
+                    next_style_idx += 1
+                r += 1
+
+            if isinstance(layer, nn.MaxPool2d) or isinstance(layer, nn.AvgPool2d):
+                net.add_module(str(len(net)), layer)
+
+    if multidevice:
+        net = setup_multi_device(net)
+
+    # Capture content targets
+    for i in content_losses:
+        i.mode = 'capture'
+    print("Capturing content targets")
+    print_torch(net, multidevice)
+    net(content_image)
+
+    # Capture style targets
+    for i in content_losses:
+        i.mode = 'None'
+
+    for i, image in enumerate(style_images_caffe):
+        print("Capturing style target " + str(i + 1))
+        for j in style_losses:
+            j.mode = 'capture'
+            j.blend_weight = style_blend_weights[i]
+        net(style_images_caffe[i])
+
+    # Set all loss modules to loss mode
+    for i in content_losses:
+        i.mode = 'loss'
+    for i in style_losses:
+        i.mode = 'loss'
+
+    # Maybe normalize content and style weights
+    if params.normalize_weights:
+        normalize_weights(content_losses, style_losses)
+
+    # Freeze the network in order to prevent
+    # unnecessary gradient calculations
+    for param in net.parameters():
+        param.requires_grad = False
+
+    # Initialize the image
+    if params.seed >= 0:
+        torch.manual_seed(params.seed)
+        torch.cuda.manual_seed_all(params.seed)
+        torch.backends.cudnn.deterministic = True
+    if params.init == 'random':
+        B, C, H, W = content_image.size()
+        img = torch.randn(C, H, W).mul(0.001).unsqueeze(0).type(dtype)
+    elif params.init == 'image':
+        if params.init_image != None:
+            img = init_image.clone()
+        else:
+            img = content_image.clone()
+    img = nn.Parameter(img)
+
+    def maybe_print(t, loss):
+        if params.print_iter > 0 and t % params.print_iter == 0:
+            print("Iteration " + str(t) + " / " + str(params.num_iterations))
+            for i, loss_module in enumerate(content_losses):
+                print("  Content " + str(i + 1) + " loss: " + str(loss_module.loss.item()))
+            for i, loss_module in enumerate(style_losses):
+                print("  Style " + str(i + 1) + " loss: " + str(loss_module.loss.item()))
+            print("  Total loss: " + str(loss.item()))
+
+    '''
+    def maybe_save(t):
+        should_save = params.save_iter > 0 and t % params.save_iter == 0
+        should_save = should_save or t == params.num_iterations
+        if should_save:
+            output_filename, file_extension = os.path.splitext(params.output_image)
+            if t == params.num_iterations:
+                filename = output_filename + str(file_extension)
+            else:
+                filename = str(output_filename) + "_" + str(t) + str(file_extension)
+            disp = deprocess(img.clone())
+
+            # Maybe perform postprocessing for color-independent style transfer
+            if params.original_colors == 1:
+                disp = original_colors(deprocess(content_image.clone()), disp)
+
+            disp.save(str(filename))
+    '''
+
+    def maybe_save(t):
+        should_save = params.save_iter > 0 and t % params.save_iter == 0
+        should_save = should_save or t == params.num_iterations
+        if should_save:
+            output_filename, file_extension = os.path.splitext(params.output_image)
+            if t == params.num_iterations:
+                filename = os.path.join('static', output_filename + file_extension)
+            else:
+                filename = os.path.join('static', output_filename + "_" + str(t) + file_extension)
+            disp = deprocess(img.clone())
+
+            # Maybe perform postprocessing for color-independent style transfer
+            if params.original_colors == 1:
+                disp = original_colors(deprocess(content_image.clone()), disp)
+
+            disp.save(filename)
+
+    # Function to evaluate loss and gradient. We run the net forward and
+    # backward to get the gradient, and sum up losses from the loss modules.
+    # optim.lbfgs internally handles iteration and calls this function many
+    # times, so we manually count the number of iterations to handle printing
+    # and saving intermediate results.
+    num_calls = [0]
+
+    def feval():
+        num_calls[0] += 1
+        optimizer.zero_grad()
+        net(img)
+        loss = 0
+
+        for mod in content_losses:
+            loss += mod.loss.to(backward_device)
+        for mod in style_losses:
+            loss += mod.loss.to(backward_device)
+        if params.tv_weight > 0:
+            for mod in tv_losses:
+                loss += mod.loss.to(backward_device)
+
+        loss.backward()
+
+        maybe_save(num_calls[0])
+        maybe_print(num_calls[0], loss)
+
+        return loss
+
+    optimizer, loopVal = setup_optimizer(img)
+    while num_calls[0] <= loopVal:
+        optimizer.step(feval)
+
+
+# Configure the optimizer
+def setup_optimizer(img):
+    if params.optimizer == 'lbfgs':
+        print("Running optimization with L-BFGS")
+        optim_state = {
+            'max_iter': params.num_iterations,
+            'tolerance_change': -1,
+            'tolerance_grad': -1,
+        }
+        if params.lbfgs_num_correction != 100:
+            optim_state['history_size'] = params.lbfgs_num_correction
+        optimizer = optim.LBFGS([img], **optim_state)
+        loopVal = 1
+    elif params.optimizer == 'adam':
+        print("Running optimization with ADAM")
+        optimizer = optim.Adam([img], lr=params.learning_rate)
+        loopVal = params.num_iterations - 1
+    return optimizer, loopVal
+
+
+def setup_gpu():
+    def setup_cuda():
+        if 'cudnn' in params.backend:
+            torch.backends.cudnn.enabled = True
+            if params.cudnn_autotune:
+                torch.backends.cudnn.benchmark = True
+        else:
+            torch.backends.cudnn.enabled = False
+
+    def setup_cpu():
+        if 'mkl' in params.backend and 'mkldnn' not in params.backend:
+            torch.backends.mkl.enabled = True
+        elif 'mkldnn' in params.backend:
+            raise ValueError("MKL-DNN is not supported yet.")
+        elif 'openmp' in params.backend:
+            torch.backends.openmp.enabled = True
+
+    multidevice = False
+    if "," in str(params.gpu):
+        devices = params.gpu.split(',')
+        multidevice = True
+
+        if 'c' in str(devices[0]).lower():
+            backward_device = "cpu"
+            setup_cuda(), setup_cpu()
+        else:
+            backward_device = "cuda:" + devices[0]
+            setup_cuda()
+        dtype = torch.FloatTensor
+
+    elif "c" not in str(params.gpu).lower():
+        setup_cuda()
+        dtype, backward_device = torch.cuda.FloatTensor, "cuda:" + str(params.gpu)
+    else:
+        setup_cpu()
+        dtype, backward_device = torch.FloatTensor, "cpu"
+    return dtype, multidevice, backward_device
+
+
+def setup_multi_device(net):
+    assert len(params.gpu.split(',')) - 1 == len(params.multidevice_strategy.split(',')), \
+        "The number of -multidevice_strategy layer indices minus 1, must be equal to the number of -gpu devices."
+
+    new_net = ModelParallel(net, params.gpu, params.multidevice_strategy)
+    return new_net
+
+
+# Preprocess an image before passing it to a model.
+# We need to rescale from [0, 1] to [0, 255], convert from RGB to BGR,
+# and subtract the mean pixel.
+def preprocess(image_name, image_size):
+    image = Image.open(image_name).convert('RGB')
+    if type(image_size) is not tuple:
+        image_size = tuple([int((float(image_size) / max(image.size)) * x) for x in (image.height, image.width)])
+    Loader = transforms.Compose([transforms.Resize(image_size), transforms.ToTensor()])
+    rgb2bgr = transforms.Compose([transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])])])
+    Normalize = transforms.Compose([transforms.Normalize(mean=[103.939, 116.779, 123.68], std=[1, 1, 1])])
+    tensor = Normalize(rgb2bgr(Loader(image) * 255)).unsqueeze(0)
+    return tensor
+
+
+#  Undo the above preprocessing.
+def deprocess(output_tensor):
+    Normalize = transforms.Compose([transforms.Normalize(mean=[-103.939, -116.779, -123.68], std=[1, 1, 1])])
+    bgr2rgb = transforms.Compose([transforms.Lambda(lambda x: x[torch.LongTensor([2, 1, 0])])])
+    output_tensor = bgr2rgb(Normalize(output_tensor.squeeze(0).cpu())) / 255
+    output_tensor.clamp_(0, 1)
+    Image2PIL = transforms.ToPILImage()
+    image = Image2PIL(output_tensor.cpu())
+    return image
+
+
+# Combine the Y channel of the generated image and the UV/CbCr channels of the
+# content image to perform color-independent style transfer.
+def original_colors(content, generated):
+    content_channels = list(content.convert('YCbCr').split())
+    generated_channels = list(generated.convert('YCbCr').split())
+    content_channels[0] = generated_channels[0]
+    return Image.merge('YCbCr', content_channels).convert('RGB')
+
+
+# Print like Lua/Torch7
+def print_torch(net, multidevice):
+    if multidevice:
+        return
+    simplelist = ""
+    for i, layer in enumerate(net, 1):
+        simplelist = simplelist + "(" + str(i) + ") -> "
+    print("nn.Sequential ( \n  [input -> " + simplelist + "output]")
+
+    def strip(x):
+        return str(x).replace(", ", ',').replace("(", '').replace(")", '') + ", "
+
+    def n():
+        return "  (" + str(i) + "): " + "nn." + str(l).split("(", 1)[0]
+
+    for i, l in enumerate(net, 1):
+        if "2d" in str(l):
+            ks, st, pd = strip(l.kernel_size), strip(l.stride), strip(l.padding)
+            if "Conv2d" in str(l):
+                ch = str(l.in_channels) + " -> " + str(l.out_channels)
+                print(n() + "(" + ch + ", " + (ks).replace(",", 'x', 1) + st + pd.replace(", ", ')'))
+            elif "Pool2d" in str(l):
+                st = st.replace("  ", ' ') + st.replace(", ", ')')
+                print(n() + "(" + ((ks).replace(",", 'x' + ks, 1) + st).replace(", ", ','))
+        else:
+            print(n())
+    print(")")
+
+
+# Divide weights by channel size
+def normalize_weights(content_losses, style_losses):
+    for n, i in enumerate(content_losses):
+        i.strength = i.strength / max(i.target.size())
+    for n, i in enumerate(style_losses):
+        i.strength = i.strength / max(i.target.size())
+
+
+# Scale gradients in the backward pass
+class ScaleGradients(torch.autograd.Function):
+    @staticmethod
+    def forward(self, input_tensor, strength):
+        self.strength = strength
+        return input_tensor
+
+    @staticmethod
+    def backward(self, grad_output):
+        grad_input = grad_output.clone()
+        grad_input = grad_input / (torch.norm(grad_input, keepdim=True) + 1e-8)
+        return grad_input * self.strength * self.strength, None
+
+
+# Define an nn Module to compute content loss
+class ContentLoss(nn.Module):
+
+    def __init__(self, strength, normalize):
+        super(ContentLoss, self).__init__()
+        self.strength = strength
+        self.crit = nn.MSELoss()
+        self.mode = 'None'
+        self.normalize = normalize
+
+    def forward(self, input):
+        if self.mode == 'loss':
+            loss = self.crit(input, self.target)
+            if self.normalize:
+                loss = ScaleGradients.apply(loss, self.strength)
+            self.loss = loss * self.strength
+        elif self.mode == 'capture':
+            self.target = input.detach()
+        return input
+
+
+class GramMatrix(nn.Module):
+
+    def forward(self, input):
+        B, C, H, W = input.size()
+        x_flat = input.view(C, H * W)
+        return torch.mm(x_flat, x_flat.t())
+
+
+# Define an nn Module to compute style loss
+class StyleLoss(nn.Module):
+
+    def __init__(self, strength, normalize):
+        super(StyleLoss, self).__init__()
+        self.target = torch.Tensor()
+        self.strength = strength
+        self.gram = GramMatrix()
+        self.crit = nn.MSELoss()
+        self.mode = 'None'
+        self.blend_weight = None
+        self.normalize = normalize
+
+    def forward(self, input):
+        self.G = self.gram(input)
+        self.G = self.G.div(input.nelement())
+        if self.mode == 'capture':
+            if self.blend_weight == None:
+                self.target = self.G.detach()
+            elif self.target.nelement() == 0:
+                self.target = self.G.detach().mul(self.blend_weight)
+            else:
+                self.target = self.target.add(self.blend_weight, self.G.detach())
+        elif self.mode == 'loss':
+            loss = self.crit(self.G, self.target)
+            if self.normalize:
+                loss = ScaleGradients.apply(loss, self.strength)
+            self.loss = self.strength * loss
+        return input
+
+
+class TVLoss(nn.Module):
+
+    def __init__(self, strength):
+        super(TVLoss, self).__init__()
+        self.strength = strength
+
+    def forward(self, input):
+        self.x_diff = input[:, :, 1:, :] - input[:, :, :-1, :]
+        self.y_diff = input[:, :, :, 1:] - input[:, :, :, :-1]
+        self.loss = self.strength * (torch.sum(torch.abs(self.x_diff)) + torch.sum(torch.abs(self.y_diff)))
+        return input
+
+# if __name__ == "__main__":
+#     main()