File size: 2,110 Bytes
6e204a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """
Unit tests for FloorGen Generative Models (Stage 3).
"""
import torch
import pytest
from floorgen.models.gan_baseline.house_gan_pp import HouseGANppGenerator, HouseGANppDiscriminator
from floorgen.models.diffusion_core.house_diffusion import DDPMScheduler, HouseDiffusionDenoiser
from floorgen.models.diffusion_core.rag_diffusion import RAGFloorplanDiffusion
from floorgen.models.shared.representation import floorplan_geometric_loss
def test_house_gan_pp():
B, N = 2, 8
room_types = torch.randint(0, 10, (B, N))
adj = torch.eye(N).unsqueeze(0).repeat(B, 1, 1)
mask = torch.ones((B, N), dtype=torch.bool)
boundary = torch.tensor([[-1.0, -1.0, 1.0, 1.0], [-1.0, -1.0, 1.0, 1.0]])
netG = HouseGANppGenerator(hidden_dim=64, num_layers=2)
netD = HouseGANppDiscriminator(hidden_dim=64, num_layers=2)
fake_boxes = netG(room_types, adj, mask, boundary)
assert fake_boxes.shape == (B, N, 4)
logits = netD(fake_boxes, room_types, adj, mask)
assert logits.shape == (B, 1)
def test_rag_diffusion_model():
B, N = 2, 8
room_types = torch.randint(0, 10, (B, N))
adj = torch.eye(N).unsqueeze(0).repeat(B, 1, 1)
mask = torch.ones((B, N), dtype=torch.bool)
noisy_boxes = torch.randn((B, N, 4))
timesteps = torch.tensor([100, 200])
exemplars = torch.randn((B, 3, N, 4)) # K=3 exemplars
model = RAGFloorplanDiffusion(hidden_dim=64, num_layers=2, num_heads=2)
scheduler = DDPMScheduler(num_timesteps=100)
# Forward prediction
eps_pred = model(
x_t=noisy_boxes,
timesteps=timesteps,
room_types=room_types,
adj_matrix=adj,
room_mask=mask,
retrieved_exemplars=exemplars
)
assert eps_pred.shape == (B, N, 4)
# Denoising sampling
sampled_boxes = model.sample(
scheduler=scheduler,
room_types=room_types,
adj_matrix=adj,
room_mask=mask,
retrieved_exemplars=exemplars,
num_inference_steps=5
)
assert sampled_boxes.shape == (B, N, 4)
assert torch.all(sampled_boxes >= -1.0) and torch.all(sampled_boxes <= 1.0)
|