Views
No views yet
1import torch
2from huggingface_hub import hf_hub_download
3import sys, os
4import torchvision.transforms as transforms
5from PIL import Image
6
7# Download model file
8model_file = hf_hub_download(
9 repo_id="hajar001/fast-neural-style-transfer",
10 filename="style_transfer_model.py"
11)
12sys.path.insert(0, os.path.dirname(model_file))
13from style_transfer_model import StyleTransferNet
14
15transform = transforms.Compose([
16 transforms.Resize((256, 256)),
17 transforms.ToTensor(),
18])
19
20# Load model
21model = StyleTransferNet.from_pretrained("hajar001/fast-neural-style-transfer")
22device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
23model = model.to(device)
24model.eval()
25
26with torch.no_grad():
27 # Load and preprocess content image
28 test_image = Image.open("content_image.jpg").convert("RGB")
29 test_tensor = transform(test_image).unsqueeze(0).to(device)
30
31 # Generate stylized image
32 stylized_tensor = model(test_tensor)
33
34 # Denormalize and convert to PIL
35 denorm = transforms.Normalize(
36 mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
37 std=[1/0.229, 1/0.224, 1/0.225]
38 )
39 stylized_tensor = denorm(stylized_tensor[0])
40 stylized_tensor = torch.clamp(stylized_tensor, 0, 1)
41
42 stylized_img = transforms.ToPILImage()(stylized_tensor.cpu())
43 stylized_img.save("stylized_image.jpg")