Style Transfer in PyTorch

A high-performance PyTorch implementation of Neural Style Transfer (Gatys et al., 2015), optimized for both consumer GPUs and multi-GPU setups.

This implementation employs a coarse-to-fine multi-scale optimization pyramid with momentum interpolation and iterate averaging, producing crisp, high-resolution artwork (up to print resolution) with minimal edge and high-frequency artifacts.

---

Example Outputs

Content Image	Style Image	Stylized Output
Golden Gate Bridge	The Starry Night (van Gogh)	Multi-Scale Result (512px)

---

Key Enhancements over Standard Gatys (2015)

Compared to the original formulation in the literature, this implementation introduces several key improvements:

• Multi-Scale Pyramid: Coarse-to-fine optimization scaling by factors of $\sqrt{2}$ eliminates stroke repetition and preserves global structure at higher resolutions.
• Warm-Started Adam Optimizer: Bilinear and bicubic spatial interpolation of first and second Adam moment vectors (exp_avg, exp_avg_sq) across pyramid scales prevents transient noise bursts.
• Iterate Averaging (EMA): Exponential moving average across optimization steps reduces iterate variance and stroke flutter.
• Scaled MSE Loss (Gradient Normalization): Approximates an $L_1$ gradient norm of $\approx 1$, preventing dominant layers from overwhelming subtle stylistic textures.
• Spatial Gram Matrix Normalization: Gram matrices are normalized by spatial dimension ($H \times W$) rather than total element count ($C \times H \times W$), preserving channel energy balance.
• Replicate Padding: First convolution layer uses replicate padding rather than zero padding, completely eliminating border discoloration and edge halos.
• Non-Uniform Style Layer Weighting: Style layers scale exponentially ($[256, 64, 16, 4, 1]$) from low to high levels, yielding richer texture synthesis.

---

Visual Comparison of Key Parameters

1. Content Weight (--content-weight / -cw)

Controls the trade-off between semantic preservation of the content and stylistic distortion.

Low (-cw 0.002)	Default (-cw 0.015)	High (-cw 0.080)
Heavy abstraction, style dominates	Balanced structure and style	Sharp photographic structure preserved

---

2. Style Scale Factor (--style-scale-fac)

Scales the style image relative to the content canvas, adjusting the physical size of synthesized brush strokes.

Fine Strokes (--style-scale-fac 0.5)	Standard (--style-scale-fac 1.0)	Broad Strokes (--style-scale-fac 1.8)
Dense, detailed, intricate textures	Standard brush stroke scale	Large, sweeping painterly swirls

---

3. Pooling Mode (--pooling)

Swaps the pooling operation used inside the pre-trained VGG-19 feature extractor.

Max Pooling (--pooling max, Default)	Average Pooling (--pooling average)
Crisp edges, localized high-contrast features	Softer gradients, blended artistic brushwork

---

Installation

Prerequisites

• Python 3.10+ (tested on Python 3.12)
• PyTorch >= 2.0 with CUDA support

Setup Environment

# 1. Create and activate a conda environment
conda create -n nst python=3.12 -y
conda activate nst

# 2. Install PyTorch with CUDA (e.g. CUDA 12.x / 13.x)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# 3. Install remaining dependencies
pip install -r requirements.txt

Note: On first launch, the pre-trained VGG-19 weights (~548 MB) will be downloaded automatically by torchvision.

---

Quick Start

Basic Stylization

python stylize.py examples/golden_gate.jpg examples/starry_night.jpg -o out.png

Multi-Style Blending

Blend multiple artworks with explicit relative weights:

python stylize.py content.jpg style1.jpg style2.jpg -sw 0.7 0.3 -o blended.png

High-Resolution Generation

For high-resolution rendering, append + to compute an optimal aspect-ratio safe resolution:

python stylize.py content.jpg style.jpg -s 1024+ -o high_res.png

Batch Stylization

To process an entire folder of content images against a target style without reloading the model:

python tools/batch_style.py --content-dir ./inputs --styles ./styles/starry_night.jpg --output-dir ./batch_outputs

---

Complete CLI Options Reference

All arguments supported by stylize.py:

1. Positional Arguments

Argument	Type	Description
content	str	Path to the content image.
style [style ...]	str	Paths to one or more style images.

2. General & I/O Options

Flag	Default	Description & Effect
-o, --output	out.png	Output path. Supports .png, .jpg, .webp, and 16-bit .tiff (uncompressed print quality).
-sw, --style-weights	None	Relative weights for multiple styles (e.g. -sw 0.7 0.3). Normalized automatically.
-r, --random-seed	0	Random seed for stochastic operations and noise initializations.
--proof	None	Optional path to a CMYK ICC profile for soft-proofing content and style colors.
--save-every	0	Iteration interval to save intermediate snapshot previews (e.g. 50). Set 0 to disable.

3. Loss & Optimization Weights

Flag	Default	Description & Effect
-cw, --content-weight	0.015	Weight of content feature reconstruction. Higher preserves original structure; lower increases stylization.
-tw, --tv-weight	2.0	Total Variation regularization weight. Higher eliminates high-frequency salt-and-pepper noise.
-ss, --step-size	0.02	Adam learning rate.
-ad, --avg-decay	0.99	Exponential moving average (EMA) decay factor for iterate stabilization.

4. Scale & Resolution Control

Flag	Default	Description & Effect
-s, --end-scale	512	Final maximum image dimension in pixels. Append + (e.g. 1024+) to adapt scale safely to GPU memory.
-ms, --min-scale	128	Initial coarse scale in pixels. Stylization scales by $\sqrt{2}$ until reaching end-scale.
-i, --iterations	500	Number of iterations per pyramid scale.
-ii, --initial-iterations	1000	Number of iterations allocated to the coarsest base scale.
--style-scale-fac	1.0	Scale factor applied to the style image relative to content canvas (controls brush stroke size).
--style-size	None	Fixed pixel dimension for style image across all scales (decouples style scale from content size).

5. Model & Initialization Options

Flag	Default	Description & Effect
--init	content	Initialization canvas. Choices: content (fastest convergence), style_mean (color palette of style), gray, uniform.
--pooling	max	Pooling replacement in VGG-19: max (crisp), average (smooth transitions), or l2.
--devices	Auto	PyTorch devices (e.g. cuda:0, or cuda:0 cuda:1 for model parallel layer distribution, or cpu).

---

Project Structure

neural-style-transfer/
├── stylize.py               # Main CLI entry point for neural style transfer
├── requirements.txt         # Project package dependencies
├── README.md                # Project documentation and visual guides
├── core/                    # Core algorithm implementation
│   ├── model.py             # VGG-19 feature extractor & pooling adaptations
│   ├── losses.py            # ScaledMSE, Content, Style, and TV loss functions
│   └── style_transfer.py    # Multi-scale pyramid scheduler, Adam scaling, EMA
├── utils/                   # Shared utility modules
│   ├── image.py             # Robust PIL/TIFF I/O, ICC color conversion, scale calculation
│   ├── device.py            # Hardware verification & GPU memory diagnostics
│   └── sRGB Profile.icc     # Standard sRGB ICC profile for color correctness
└── tools/                   # Automation and batch processing utilities
    └── batch_style.py       # High-throughput batch stylization tool

---

References

1. Foundational Neural Style Transfer Papers

• L. A. Gatys, A. S. Ecker, M. Bethge (2015). "A Neural Algorithm of Artistic Style". arXiv:1508.06576. [Paper]
• L. A. Gatys, A. S. Ecker, M. Bethge, A. Hertzmann, E. Shechtman (2016). "Controlling Perceptual Factors in Neural Style Transfer". CVPR 2017 / arXiv:1611.07865. [Paper]

2. Perceptual Losses & Image Inversion

• J. Johnson, A. Alahi, L. Fei-Fei (2016). "Perceptual Losses for Real-Time Style Transfer and Super-Resolution". ECCV 2016 / arXiv:1603.08155. [Paper]
• A. Mahendran, A. Vedaldi (2014). "Understanding Deep Image Representations by Inverting Them". CVPR 2015 / arXiv:1412.0035. [Paper]

3. Neural Architecture & Optimization

• K. Simonyan, A. Zisserman (2014). "Very Deep Convolutional Networks for Large-Scale Image Recognition". ICLR 2015 / arXiv:1409.1556. [Paper]
• D. P. Kingma, J. Ba (2014). "Adam: A Method for Stochastic Optimization". ICLR 2015 / arXiv:1412.6980. [Paper]