Views
No views yet
1import torch
2import torch.nn as nn
3import os
4import requests
5
6# 1. DOWNLOAD MODELS
7def download_file(url, filename):
8 if not os.path.exists(filename):
9 print(f"Downloading {filename}...")
10 r = requests.get(url, allow_redirects=True)
11 with open(filename, 'wb') as f:
12 f.write(r.content)
13
14urls = {
15 "ae_global.pt": "https://huggingface.co/Parallax-labs-1/parallax_VIDEO-Boxes/resolve/main/ae_global.pt",
16 "predictor.pt": "https://huggingface.co/Parallax-labs-1/parallax_VIDEO-Boxes/resolve/main/predictor.pt",
17 "vision_base.pt": "https://huggingface.co/Parallax-labs-1/parallax_VISION-boxes-RGBA/resolve/main/model.pt"
18}
19
20for name, url in urls.items():
21 download_file(url, name)
22
23# 2. DEFINE ARCHITECTURES
24class GlobalHCA_AE(nn.Module):
25 def __init__(self):
26 super().__init__()
27 self.encoder = nn.Sequential(
28 nn.Conv2d(4, 16, 3, stride=2, padding=1), nn.ReLU(),
29 nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU(),
30 nn.Flatten(),
31 nn.Linear(32 * 12 * 12, 2048), nn.ReLU(),
32 nn.Linear(2048, 1012 * 2)
33 )
34 self.decoder = nn.Sequential(
35 nn.Linear(1012, 2048), nn.ReLU(),
36 nn.Linear(2048, 32 * 12 * 12), nn.ReLU(),
37 nn.Unflatten(1, (32, 12, 12)),
38 nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1), nn.ReLU(),
39 nn.ConvTranspose2d(16, 4, 3, stride=2, padding=1), nn.Sigmoid()
40 )
41
42 def forward(self, x):
43 h = self.encoder(x)
44 mu, _ = h.chunk(2, dim=-1)
45 return self.decoder(mu), mu
46
47class LatentPredictor(nn.Module):
48 def __init__(self):
49 super().__init__()
50 self.net = nn.Sequential(
51 nn.Linear(1012, 2048), nn.ReLU(),
52 nn.Linear(2048, 4096), nn.ReLU(),
53 nn.Linear(4096, 8100),
54 nn.Unflatten(1, (4, 45, 45))
55 )
56 def forward(self, z):
57 return self.net(z)
58
59class AlphaAutoencoder(nn.Module):
60 def __init__(self):
61 super().__init__()
62 self.encoder = nn.Sequential(
63 nn.Conv2d(4, 32, 3, stride=2, padding=1),
64 nn.LeakyReLU(0.2),
65 nn.Conv2d(32, 64, 3, stride=2, padding=1),
66 nn.LeakyReLU(0.2),
67 nn.Conv2d(64, 128, 3, stride=2, padding=1),
68 nn.LeakyReLU(0.2),
69 nn.Conv2d(128, 256, 3, stride=2, padding=1),
70 nn.LeakyReLU(0.2),
71 nn.Conv2d(256, 4, 1)
72 )
73 self.decoder = nn.Sequential(
74 nn.Conv2d(4, 256, 3, padding=1),
75 nn.PixelShuffle(2),
76 nn.LeakyReLU(0.2),
77 nn.Conv2d(64, 128, 3, padding=1),
78 nn.PixelShuffle(2),
79 nn.LeakyReLU(0.2),
80 nn.Conv2d(32, 64, 3, padding=1),
81 nn.PixelShuffle(2),
82 nn.LeakyReLU(0.2),
83 nn.Conv2d(16, 16, 3, padding=1),
84 nn.PixelShuffle(2),
85 nn.Sigmoid()
86 )
87 def forward(self, x):
88 z = self.encoder(x)
89 return self.decoder(z), z
90
91# 3. LOAD WEIGHTS
92device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
93
94# Initialize models
95ae_global = GlobalHCA_AE().to(device)
96predictor = LatentPredictor().to(device)
97vision_base = AlphaAutoencoder().to(device)
98
99# Load weights
100ae_global.load_state_dict(torch.load("ae_global.pt", map_location=device))
101predictor.load_state_dict(torch.load("predictor.pt", map_location=device))
102vision_base.load_state_dict(torch.load("vision_base.pt", map_location=device))
103
104print("All Parallax-Labs models (Global AE, Predictor, and Vision Base) loaded successfully.")| Phase | Metric | Value |
|---|---|---|
| Autoencoder | Reconstruction MSE | 0.001722 |
| Predictor | Latent-to-Pixel MSE | 0.000150 |
| Robustness | Max Noise Stability | 0.30 |