Views
No views yet
torch, torchvision, matplotlib и huggingface_hub.1import torch
2import torch.nn as nn
3import torchvision.utils as vutils
4import matplotlib.pyplot as plt
5from huggingface_hub import hf_hub_download
6
7# 1. Архитектура Генератора
8class Generator(nn.Module):
9 def __init__(self):
10 super(Generator, self).__init__()
11 self.main = nn.Sequential(
12 nn.ConvTranspose2d(100, 64 * 8, 4, 1, 0, bias=False),
13 nn.BatchNorm2d(64 * 8), nn.ReLU(True),
14 nn.ConvTranspose2d(64 * 8, 64 * 4, 4, 2, 1, bias=False),
15 nn.BatchNorm2d(64 * 4), nn.ReLU(True),
16 nn.ConvTranspose2d(64 * 4, 64 * 2, 4, 2, 1, bias=False),
17 nn.BatchNorm2d(64 * 2), nn.ReLU(True),
18 nn.ConvTranspose2d(64 * 2, 64, 4, 2, 1, bias=False),
19 nn.BatchNorm2d(64), nn.ReLU(True),
20 nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False),
21 nn.Tanh()
22 )
23 def forward(self, input): return self.main(input)
24
25# 2. Загрузка весов
26device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
27repo_id = "prostochel097/alphagpt-image" # Твой ID репозитория
28
29weights_path = hf_hub_download(repo_id=repo_id, filename="generator.pth")
30model = Generator().to(device)
31model.load_state_dict(torch.load(weights_path, map_location=device))
32model.eval()
33
34# 3. Генерация
35noise = torch.randn(1, 100, 1, 1, device=device)
36with torch.no_grad():
37 fake = model(noise).detach().cpu()
38
39plt.figure(figsize=(5,5))
40plt.axis("off")
41plt.imshow(vutils.make_grid(fake, padding=2, normalize=True).permute(1,2,0))
42plt.show()