Views
No views yet
| Content Image | Stylized Output |
|---|---|
![]() | ![]() |
huggingface_hub utility library. The script automatically downloads your weights file directly from the cloud and applies the necessary ImageNet normalization matching the training routine.pip install torch torchvision pillow huggingface_hubinference.py)inference.py. You can run it via terminal with python inference.py your_image.jpg.1import sys
2import torch
3import torch.nn as nn
4from PIL import Image
5from torchvision import transforms
6from torchvision.utils import save_image
7from huggingface_hub import hf_hub_download
8
9# ── CONFIG ───────────────────────────────────────────────────
10REPO_ID = "Rohanify/Brawnz-StyleTransferSN"
11FILENAME = "pytorch_model.bin"
12IMG_SIZE = 512
13DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
14# ─────────────────────────────────────────────────────────────
15
16# ── NATIVE PYTORCH NETWORK DEFINITION ────────────────────────
17
18def conv_bn_relu(in_c, out_c, k, stride=1, pad=0):
19 return nn.Sequential(
20 nn.ReflectionPad2d(pad),
21 nn.Conv2d(in_c, out_c, k, stride),
22 nn.InstanceNorm2d(out_c),
23 nn.ReLU(inplace=True),
24 )
25
26class ResBlock(nn.Module):
27 def __init__(self, c):
28 super().__init__()
29 self.block = nn.Sequential(
30 nn.ReflectionPad2d(1),
31 nn.Conv2d(c, c, 3),
32 nn.InstanceNorm2d(c),
33 nn.ReLU(inplace=True),
34 nn.ReflectionPad2d(1),
35 nn.Conv2d(c, c, 3),
36 nn.InstanceNorm2d(c),
37 )
38 def forward(self, x):
39 return x + self.block(x)
40
41class TransformNet(nn.Module):
42 def __init__(self):
43 super().__init__()
44 self.net = nn.Sequential(
45 conv_bn_relu(3, 32, 9, pad=4),
46 conv_bn_relu(32, 64, 3, stride=2, pad=1),
47 conv_bn_relu(64, 128, 3, stride=2, pad=1),
48 ResBlock(128), ResBlock(128), ResBlock(128),
49 ResBlock(128), ResBlock(128),
50 nn.Upsample(scale_factor=2, mode="nearest"),
51 conv_bn_relu(128, 64, 3, pad=1),
52 nn.Upsample(scale_factor=2, mode="nearest"),
53 conv_bn_relu(64, 32, 3, pad=1),
54 nn.ReflectionPad2d(4),
55 nn.Conv2d(32, 3, 9),
56 nn.Tanh(),
57 )
58 def forward(self, x):
59 return self.net(x)
60
61# ── LOAD INPUT IMAGE ─────────────────────────────────────────
62if len(sys.argv) < 2:
63 print("Usage: python inference.py path_to_input_image.jpg")
64 sys.exit(1)
65
66input_path = sys.argv[1]
67output_path = "output_styled.jpg"
68
69transform = transforms.Compose([
70 transforms.Resize((IMG_SIZE, IMG_SIZE)),
71 transforms.ToTensor(),
72 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
73])
74
75img = Image.open(input_path).convert("RGB")
76x = transform(img).unsqueeze(0).to(DEVICE)
77
78# ── SECURE FILE DOWNLOAD & STATE LOAD ────────────────────────
79print("Downloading weights from Hugging Face Hub...")
80weights_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
81
82model = TransformNet().to(DEVICE)
83model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
84model.eval()
85print(f"Weights successfully loaded on: {DEVICE}")
86
87# ── RUN INFERENCE ────────────────────────────────────────────
88print("Processing style transfer...")
89with torch.no_grad():
90 out = model(x)
91
92save_image(out[0] * 0.5 + 0.5, output_path)
93print(f"Success! Styled image saved to: {output_path}")
94