import svetoviz_webgpu as sv

# 1. Prepare 2D image data
img = Image.open("sample_image.jpg").convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
input_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0)

# 2. Downsample with Conv2d (Stride=2)
encoder = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1)

# 3. Upsample with ConvTranspose2d (Stride=2)
decoder = nn.ConvTranspose2d(in_channels=16, out_channels=3, kernel_size=4, stride=2, padding=1)

def terminal_callback(buffer, message, images, files):
    # Forward Pass
    latent_features = encoder(input_tensor)
    reconstruction = decoder(latent_features)

    buffer.send_system_message(f"Input: {list(input_tensor.shape)}")
    buffer.send_system_message(f"Latent: {list(latent_features.shape)}")
    buffer.send_system_message(f"Output: {list(reconstruction.shape)}")

# 4. Visualizing the Decoder module
sv.pytorch_web(module=decoder, terminal_callback=terminal_callback)