feat(neopp): support conv pixel head (ConvDecoder) fm_head

[fuheaven]

Adapt the neopp generation head to the new 8+8 neo model whose fm_head is a convolutional pixel head (PixelShuffle + Conv2d) decoding hidden states directly to RGB, instead of the legacy per-token MLP head. Guarded by use_pixel_head so the legacy MLP-head models keep working.

• module_io: add token_h/token_w to carry 2D grid dims to the head
• pre_infer: populate token_h/token_w
• transformer_weights: register conv1/conv2 when use_pixel_head else MLP head
• transformer_infer: dual-branch _fm_head + _fm_head_pixel + _patchify_pixels

Numerically verified equivalent to the reference ConvDecoder (max abs diff 0.0).

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, you can upgrade your account or add credits to your account and enable them for code reviews in your settings.

[gemini-code-assist]

Code Review

This pull request introduces a convolutional 'pixel head' variant to the Neopp transformer inference pipeline, allowing token hidden states to be decoded directly to RGB pixels via PixelShuffle and Conv2d operations before being re-patchified. The changes include tracking token grid dimensions in the pre-inference output, implementing the pixel head decoding and patchifying logic in the transformer inference, and updating weight loading to support the new Conv2d layers. The review feedback recommends adding defensive validation checks to prevent runtime errors: specifically, verifying that the token dimensions are positive and match the sequence length in _fm_head_pixel, and ensuring that the image dimensions are perfectly divisible by the patch size in _patchify_pixels.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When use_pixel_head is enabled, token_h and token_w are required to reshape the flat token sequence back to a 2D feature map. If they are unpopulated (defaulting to 0) or invalid, the .view() operation will fail with a cryptic runtime error. Adding defensive validation checks ensures that token_h and token_w are positive integers and that their product matches the sequence length of hidden_states.

Suggested change

	def _fm_head_pixel(self, fm_head_w, hidden_states, token_h, token_w):
	"""Conv "pixel head" (ConvDecoder): decode token hidden states directly to RGB
	pixels via PixelShuffle + Conv2d, then re-patchify to the [L, patch_dim] token
	space so the downstream v_pred / Euler / unpatchify path stays unchanged.
	hidden_states: [L, hidden] with L == token_h * token_w in row-major order.
	(Assumes non seq-parallel: full token sequence present on this rank.)
	"""
	# [L, hidden] -> [1, hidden, token_h, token_w]
	x = hidden_states.view(token_h, token_w, -1).permute(2, 0, 1).unsqueeze(0).contiguous()
	def _fm_head_pixel(self, fm_head_w, hidden_states, token_h, token_w):
	"""Conv "pixel head" (ConvDecoder): decode token hidden states directly to RGB
	pixels via PixelShuffle + Conv2d, then re-patchify to the [L, patch_dim] token
	space so the downstream v_pred / Euler / unpatchify path stays unchanged.
	hidden_states: [L, hidden] with L == token_h * token_w in row-major order.
	(Assumes non seq-parallel: full token sequence present on this rank.)
	"""
	if not token_h or not token_w:
	raise ValueError(f"token_h and token_w must be positive integers, got {token_h} and {token_w}")
	if hidden_states.shape[0] != token_h * token_w:
	raise ValueError(f"hidden_states sequence length {hidden_states.shape[0]} does not match token_h * token_w ({token_h} * {token_w})")
	# [L, hidden] -> [1, hidden, token_h, token_w]
	x = hidden_states.view(token_h, token_w, -1).permute(2, 0, 1).unsqueeze(0).contiguous()

[gemini-code-assist]

In _patchify_pixels, the spatial dimensions of the images tensor are divided by patch_size and then reshaped. If the spatial dimensions are not perfectly divisible by patch_size, the .reshape() operation will fail with a shape mismatch error. Adding a defensive check/assertion ensures that the input dimensions are valid and provides a clear error message if they are not.

Suggested change

	def _patchify_pixels(images, patch_size):
	# images: [N, 3, H, W] -> [N, (H/ps)*(W/ps), ps*ps*3], row-major (h outer, w inner)
	h, w = images.shape[2] // patch_size, images.shape[3] // patch_size
	x = images.reshape(images.shape[0], 3, h, patch_size, w, patch_size)
	@staticmethod
	def _patchify_pixels(images, patch_size):
	# images: [N, 3, H, W] -> [N, (H/ps)*(W/ps), ps*ps*3], row-major (h outer, w inner)
	if images.shape[2] % patch_size != 0 or images.shape[3] % patch_size != 0:
	raise ValueError(
	f"Image dimensions ({images.shape[2]}x{images.shape[3]}) must be "
	f"divisible by patch_size ({patch_size})"
	)
	h, w = images.shape[2] // patch_size, images.shape[3] // patch_size
	x = images.reshape(images.shape[0], 3, h, patch_size, w, patch_size)