Update app.py

app.py

@@ -50,6 +50,16 @@ model = ConvVAE()
 model.load_state_dict(torch.load("vae_supertux.pth", map_location=torch.device("cpu")))
 model.eval()
 
+# Sampling
+def sample_with_temperature(probs, temperature=1.2):
+    logits = torch.log(probs + 1e-8) / temperature
+    scaled_probs = torch.softmax(logits, dim=1)
+    batch, channels, height, width = scaled_probs.shape
+    scaled_probs = scaled_probs.permute(0, 2, 3, 1).contiguous().view(-1, channels)
+    sampled = torch.multinomial(scaled_probs, num_samples=1)
+    sampled = sampled.view(batch, height, width)
+    return sampled
+    
 def generate_map(seed: int = None):
     model.eval()
     if seed is None: