import svetoviz_webgpu as sv

# 1. Prepare 3D MRI volume (32 slices for clean striding)
ds = load_dataset("g4m3r/T1w_MRI_Brain_Slices", split="train", streaming=True)
slices = [np.array(ex['image'].resize((128, 128)).convert("L")) for ex in ds.skip(2).take(32)]
input_volume = torch.from_numpy(np.stack(slices)).unsqueeze(0).unsqueeze(0)

# 2. Encoder: Downsample (32, 128, 128) -> (16, 64, 64)
encoder = nn.Conv3d(in_channels=1, out_channels=8, kernel_size=3, stride=2, padding=1)

# 3. Decoder: Reconstruction with ConvTranspose3d
decoder = nn.ConvTranspose3d(in_channels=8, out_channels=1, kernel_size=4, stride=2, padding=1)

def terminal_callback(buffer, message, images, files):
    latent = encoder(input_volume)
    output = decoder(latent)
    buffer.send_system_message(f"Latent shape: {list(latent.shape)}")
    buffer.send_system_message(f"Reconstructed shape: {list(output.shape)}")

# 4. Start interactive session
sv.pytorch_web(module=decoder, terminal_callback=terminal_callback)