(
xy: Float[Tensor, "batch height width 2"],
color_function: ColorFunction,
scale: float,
subdivision: int,
remaining_passes: int,
device: torch.device,
batch_size: int = int(2**16),
)
| 74 | |
| 75 | @torch.no_grad() |
| 76 | def run_msaa_pass( |
| 77 | xy: Float[Tensor, "batch height width 2"], |
| 78 | color_function: ColorFunction, |
| 79 | scale: float, |
| 80 | subdivision: int, |
| 81 | remaining_passes: int, |
| 82 | device: torch.device, |
| 83 | batch_size: int = int(2**16), |
| 84 | ) -> Float[Tensor, "batch 4 height width"]: # color (RGBA with straight alpha) |
| 85 | # Sample the color function. |
| 86 | b, h, w, _ = xy.shape |
| 87 | color = [ |
| 88 | color_function(batch) |
| 89 | for batch in rearrange(xy, "b h w xy -> (b h w) xy").split(batch_size) |
| 90 | ] |
| 91 | color = torch.cat(color, dim=0) |
| 92 | color = rearrange(color, "(b h w) c -> b c h w", b=b, h=h, w=w) |
| 93 | |
| 94 | # If any MSAA passes remain, subdivide. |
| 95 | if remaining_passes > 0: |
| 96 | mask = detect_msaa_pixels(color) |
| 97 | batch_index, row_index, col_index = torch.where(mask) |
| 98 | xy = xy[batch_index, row_index, col_index] |
| 99 | |
| 100 | offsets = generate_sample_grid((subdivision, subdivision), device) |
| 101 | offsets = (offsets / subdivision - 0.5) * scale |
| 102 | |
| 103 | color_fine = run_msaa_pass( |
| 104 | xy[:, None, None] + offsets, |
| 105 | color_function, |
| 106 | scale / subdivision, |
| 107 | subdivision, |
| 108 | remaining_passes - 1, |
| 109 | device, |
| 110 | batch_size=batch_size, |
| 111 | ) |
| 112 | color[batch_index, :, row_index, col_index] = reduce_straight_alpha(color_fine) |
| 113 | |
| 114 | return color |
| 115 | |
| 116 | |
| 117 | @torch.no_grad() |
no test coverage detected