可以本地跑的 Python 图像风格迁移（Neural Style Transfer）

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from PIL import Image
import copy

# ---- 设置设备 ----
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---- 加载和预处理图像 ----
def image_loader(image_name, imsize):
    loader = transforms.Compose([
        transforms.Resize(imsize),
        transforms.ToTensor()
    ])
    image = Image.open(image_name)
    image = loader(image).unsqueeze(0)
    return image.to(device, torch.float)

# ---- 定义内容损失 ----
class ContentLoss(nn.Module):
    def __init__(self, target):
        super(ContentLoss, self).__init__()
        self.target = target.detach()
    def forward(self, x):
        self.loss = nn.functional.mse_loss(x, self.target)
        return x

# ---- 定义风格损失 ----
def gram_matrix(input):
    a, b, c, d = input.size()  # a=batch_size(=1)
    features = input.view(a * b, c * d)
    G = torch.mm(features, features.t())
    return G.div(a * b * c * d)

class StyleLoss(nn.Module):
    def __init__(self, target_feature):
        super(StyleLoss, self).__init__()
        self.target = gram_matrix(target_feature).detach()
    def forward(self, x):
        G = gram_matrix(x)
        self.loss = nn.functional.mse_loss(G, self.target)
        return x

# ---- 加载模型 ----
cnn = models.vgg19(pretrained=True).features.to(device).eval()

# ---- 设置图像路径 ----
content_img = image_loader("content.jpg", imsize=512)
style_img = image_loader("style.jpg", imsize=512)
input_img = content_img.clone()

# ---- 损失层添加位置 ----
content_layers = ['conv_4']
style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']

# ---- 创建模型并插入损失模块 ----
cnn = copy.deepcopy(cnn)
content_losses = []
style_losses = []

model = nn.Sequential()
i = 0
for layer in cnn.children():
    if isinstance(layer, nn.Conv2d):
        i += 1
        name = 'conv_{}'.format(i)
    elif isinstance(layer, nn.ReLU):
        name = 'relu_{}'.format(i)
        layer = nn.ReLU(inplace=False)
    elif isinstance(layer, nn.MaxPool2d):
        name = 'pool_{}'.format(i)
    elif isinstance(layer, nn.BatchNorm2d):
        name = 'bn_{}'.format(i)
    else:
        continue

    model.add_module(name, layer)

    if name in content_layers:
        target = model(content_img).detach()
        content_loss = ContentLoss(target)
        model.add_module("content_loss_{}".format(i), content_loss)
        content_losses.append(content_loss)

    if name in style_layers:
        target = model(style_img).detach()
        style_loss = StyleLoss(target)
        model.add_module("style_loss_{}".format(i), style_loss)
        style_losses.append(style_loss)

# ---- 优化器 ----
optimizer = optim.LBFGS([input_img.requires_grad_()])

# ---- 风格迁移循环 ----
num_steps = 300
style_weight = 1e6
content_weight = 1

run = [0]
while run[0] <= num_steps:

    def closure():
        optimizer.zero_grad()
        model(input_img)
        style_score = sum(sl.loss for sl in style_losses)
        content_score = sum(cl.loss for cl in content_losses)

        loss = style_weight * style_score + content_weight * content_score
        loss.backward()

        if run[0] % 50 == 0:
            print(f"Step {run[0]}: Style Loss: {style_score.item():4f}, Content Loss: {content_score.item():4f}")

        run[0] += 1
        return loss

    optimizer.step(closure)

# ---- 保存输出结果 ----
output = input_img.cpu().clone().squeeze(0)
output = transforms.ToPILImage()(output.clamp(0, 1))
output.save("result.jpg")
print("✅ 风格迁移完成，结果保存在 result.jpg")