Views
No views yet
1Total Parameters: 493,876
2Starting training on cuda...
3Step [100/10000] | Loss: 0.001546
4...
5Step [1500/10000] | Loss: 0.000405
6...
7Step [3000/10000] | Loss: 0.000444
8...
9Step [6000/10000] | Loss: 0.000317
10...
11Step [9000/10000] | Loss: 0.000135
12...
13Step [10000/10000] | Loss: 0.000038
14Model saved as ae_model_720_rgba.pt
151import torch
2import torch.nn as nn
3from PIL import Image
4import torchvision.transforms as T
5from huggingface_hub import hf_hub_download
6
7# 1. THE ACTUAL ARCHITECTURE (AlphaAutoencoder)
8class AlphaAutoencoder(nn.Module):
9 def __init__(self):
10 super().__init__()
11 # Encoder: 720 -> 360 -> 180 -> 90 -> 45
12 self.encoder = nn.Sequential(
13 nn.Conv2d(4, 32, 3, stride=2, padding=1),
14 nn.LeakyReLU(0.2),
15 nn.Conv2d(32, 64, 3, stride=2, padding=1),
16 nn.LeakyReLU(0.2),
17 nn.Conv2d(64, 128, 3, stride=2, padding=1),
18 nn.LeakyReLU(0.2),
19 nn.Conv2d(128, 256, 3, stride=2, padding=1),
20 nn.LeakyReLU(0.2),
21 nn.Conv2d(256, 4, 1)
22 )
23 # Decoder: 45 -> 90 -> 180 -> 360 -> 720
24 self.decoder = nn.Sequential(
25 nn.Conv2d(4, 256, 3, padding=1),
26 nn.PixelShuffle(2),
27 nn.LeakyReLU(0.2),
28 nn.Conv2d(64, 128, 3, padding=1),
29 nn.PixelShuffle(2),
30 nn.LeakyReLU(0.2),
31 nn.Conv2d(32, 64, 3, padding=1),
32 nn.PixelShuffle(2),
33 nn.LeakyReLU(0.2),
34 nn.Conv2d(16, 16, 3, padding=1),
35 nn.PixelShuffle(2),
36 nn.Sigmoid()
37 )
38
39 def forward(self, x):
40 return self.decoder(self.encoder(x))
41
42# 2. SETUP & DOWNLOAD
43device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
44REPO_ID = "Parallax-labs-1/parallax_VISION-boxes-RGBA"
45FILENAME = "model.pt" # Using your exact filename
46
47print(f"Fetching weights from {REPO_ID}...")
48model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
49
50# 3. INITIALIZE AND LOAD
51model = AlphaAutoencoder().to(device)
52model.load_state_dict(torch.load(model_path, map_location=device))
53model.eval()
54print("Model AlphaAutoencoder is live and ready.")
55
56# 4. INFERENCE FUNCTION
57def run_parallax_inference(img_path):
58 img = Image.open(img_path).convert("RGBA").resize((720, 720))
59 transform = T.Compose([T.ToTensor()])
60 input_tensor = transform(img).unsqueeze(0).to(device)
61
62 with torch.no_grad():
63 reconstructed = model(input_tensor)
64
65 # Convert back to PIL
66 output_img = T.ToPILImage()(reconstructed.squeeze(0).cpu())
67 return output_img
68
69print("Inference function 'run_parallax_inference' is ready to use.")"If a model can reconstruct a human face using only what it learned from squares, imagine what it can do once you actually show it the world."