Upload model.py with huggingface_hub

model.py

@@ -0,0 +1,72 @@
+"""Fully-convolutional autoencoder over OpenFront tile states.
+
+Input encoding per tile:
+  - owner slot -> learned embedding (static slot per player for a whole game,
+    slot 0 = unowned, slots 1..MAX_SLOTS-1 assigned by smallID order at spawn)
+  - terrain scalars: land flag, normalized magnitude, fallout flag
+
+The encoder downsamples by 16x into a spatial latent grid (LATENT_C channels
+per 16x16 tile region), so any map size divisible by 16 works. The decoder
+reconstructs per-tile owner-slot logits.
+"""
+
+import torch
+import torch.nn as nn
+
+MAX_SLOTS = 128  # owner classes: 0 = unowned, 1..127 player slots
+OWNER_EMB_DIM = 8
+TERRAIN_CHANNELS = 3  # land, magnitude, fallout
+
+
+def conv_block(c_in: int, c_out: int, stride: int) -> nn.Sequential:
+    return nn.Sequential(
+        nn.Conv2d(c_in, c_out, kernel_size=3, stride=stride, padding=1),
+        nn.GroupNorm(8, c_out),
+        nn.SiLU(),
+    )
+
+
+def deconv_block(c_in: int, c_out: int) -> nn.Sequential:
+    return nn.Sequential(
+        nn.ConvTranspose2d(c_in, c_out, kernel_size=4, stride=2, padding=1),
+        nn.GroupNorm(8, c_out),
+        nn.SiLU(),
+    )
+
+
+class TileAutoencoder(nn.Module):
+    def __init__(self, latent_c: int = 64):
+        super().__init__()
+        self.owner_emb = nn.Embedding(MAX_SLOTS, OWNER_EMB_DIM)
+        c_in = OWNER_EMB_DIM + TERRAIN_CHANNELS
+
+        self.encoder = nn.Sequential(
+            conv_block(c_in, 32, stride=1),
+            conv_block(32, 64, stride=2),  # /2
+            conv_block(64, 96, stride=2),  # /4
+            conv_block(96, 128, stride=2),  # /8
+            conv_block(128, 128, stride=2),  # /16
+            nn.Conv2d(128, latent_c, kernel_size=1),
+        )
+
+        self.decoder = nn.Sequential(
+            conv_block(latent_c, 128, stride=1),
+            deconv_block(128, 128),  # /8
+            deconv_block(128, 96),  # /4
+            deconv_block(96, 64),  # /2
+            deconv_block(64, 32),  # /1
+            nn.Conv2d(32, MAX_SLOTS, kernel_size=1),
+        )
+
+    def encode(self, owners: torch.Tensor, terrain: torch.Tensor) -> torch.Tensor:
+        """owners: (B, H, W) int64 slots; terrain: (B, 3, H, W) float."""
+        emb = self.owner_emb(owners).permute(0, 3, 1, 2)  # (B, E, H, W)
+        x = torch.cat([emb, terrain], dim=1)
+        return self.encoder(x)
+
+    def forward(
+        self, owners: torch.Tensor, terrain: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        z = self.encode(owners, terrain)
+        logits = self.decoder(z)  # (B, MAX_SLOTS, H, W)
+        return logits, z