import svetoviz_webgpu as sv

# 1. Parameters
upscale_factor = 2
# Input must have C * r^2 channels (e.g., 3 * 2^2 = 12)
simulated_input = torch.cat([input_tensor] * (upscale_factor ** 2), dim=1)

# 2. Define Module
# Rearranges elements in a tensor of shape (*, C*r^2, H, W) to (*, C, H*r, W*r)
shuffle = nn.PixelShuffle(upscale_factor=upscale_factor)

def terminal_callback(buffer, message, images, files):
    output = shuffle(simulated_input)

    buffer.send_system_message(f"Input Channels: {simulated_input.shape[1]}")
    buffer.send_system_message(f"Upscale Factor: {upscale_factor}")
    buffer.send_system_message(f"Output Shape: {list(output.shape)}")
    buffer.send_system_message("Result: Channels interleaved into higher spatial resolution.")

# 3. Start the interactive session
sv.pytorch_web(module=shuffle, terminal_callback=terminal_callback)