Update README.md

README.md

@@ -21,7 +21,7 @@ This repository contains weights for a Fast Neural Style Transfer network based
 
 ## How to Use Programmatically
 
-You can run inference using the Hugging Face `transformers` library. The script automatically downloads your weights and applies the necessary ImageNet normalization matching the training routine.
+You can run inference using the official `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.
 
 ### Dependencies
 Ensure you have the required packages installed:
@@ -35,18 +35,65 @@ Save the following code as `inference.py`. You can run it via terminal with `pyt
 ```python
 import sys
 import torch
+import torch.nn as nn
 from PIL import Image
 from torchvision import transforms
 from torchvision.utils import save_image
-from transformers import AutoModel
+from huggingface_hub import hf_hub_download
 
 # ── CONFIG ───────────────────────────────────────────────────
-REPO_ID  = "Rohanify/Brawnz-StyleTransferSN"  # Your repository ID
+REPO_ID  = "Rohanify/Brawnz-StyleTransferSN"
+FILENAME = "pytorch_model.bin"
 IMG_SIZE = 512
 DEVICE   = "cuda" if torch.cuda.is_available() else "cpu"
 # ─────────────────────────────────────────────────────────────
 
-# ── LOAD IMAGE ───────────────────────────────────────────────
+# ── NATIVE PYTORCH NETWORK DEFINITION ────────────────────────
+
+def conv_bn_relu(in_c, out_c, k, stride=1, pad=0):
+    return nn.Sequential(
+        nn.ReflectionPad2d(pad),
+        nn.Conv2d(in_c, out_c, k, stride),
+        nn.InstanceNorm2d(out_c),
+        nn.ReLU(inplace=True),
+    )
+
+class ResBlock(nn.Module):
+    def __init__(self, c):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(c, c, 3),
+            nn.InstanceNorm2d(c),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(c, c, 3),
+            nn.InstanceNorm2d(c),
+        )
+    def forward(self, x):
+        return x + self.block(x)
+
+class TransformNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.net = nn.Sequential(
+            conv_bn_relu(3,   32,  9, pad=4),
+            conv_bn_relu(32,  64,  3, stride=2, pad=1),
+            conv_bn_relu(64, 128,  3, stride=2, pad=1),
+            ResBlock(128), ResBlock(128), ResBlock(128),
+            ResBlock(128), ResBlock(128),
+            nn.Upsample(scale_factor=2, mode="nearest"),
+            conv_bn_relu(128, 64, 3, pad=1),
+            nn.Upsample(scale_factor=2, mode="nearest"),
+            conv_bn_relu(64, 32, 3, pad=1),
+            nn.ReflectionPad2d(4),
+            nn.Conv2d(32, 3, 9),
+            nn.Tanh(),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+# ── LOAD INPUT IMAGE ─────────────────────────────────────────
 if len(sys.argv) < 2:
     print("Usage: python inference.py path_to_input_image.jpg")
     sys.exit(1)
@@ -57,24 +104,27 @@ output_path = "output_styled.jpg"
 transform = transforms.Compose([
     transforms.Resize((IMG_SIZE, IMG_SIZE)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225]),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
 ])
 
 img = Image.open(input_path).convert("RGB")
 x   = transform(img).unsqueeze(0).to(DEVICE)
 
-# ── LOAD MODEL FROM HUGGING FACE ─────────────────────────────
-print("Loading model weights from Hugging Face Hub...")
-model = AutoModel.from_pretrained(REPO_ID, trust_remote_code=True).to(DEVICE)
+# ── SECURE FILE DOWNLOAD & STATE LOAD ────────────────────────
+print("Downloading weights from Hugging Face Hub...")
+weights_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
+
+model = TransformNet().to(DEVICE)
+model.load_state_dict(torch.load(weights_path, map_location=DEVICE))
 model.eval()
+print(f"Weights successfully loaded on: {DEVICE}")
 
 # ── RUN INFERENCE ────────────────────────────────────────────
-print(f"Running inference on device: {DEVICE}...")
+print("Processing style transfer...")
 with torch.no_grad():
     out = model(x)
 
-# De-normalize and save output image
 save_image(out[0] * 0.5 + 0.5, output_path)
 print(f"Success! Styled image saved to: {output_path}")
+
 ```