ModelTC · fuheaven · Jul 6, 2026 · gemini-code-assist · Jul 6, 2026 · gemini-code-assist
diff --git a/lightx2v/models/networks/neopp/infer/module_io.py b/lightx2v/models/networks/neopp/infer/module_io.py
@@ -15,3 +15,7 @@ class NeoppPreInferModuleOutput:
     # 在 _infer_t2i_i2i 中预计算，避免 _infer_pass 中反复 chunk/restore
     image_embeds_cond: Optional[torch.Tensor] = None
     image_embeds_uncond: Optional[torch.Tensor] = None
+    # Token grid dims (row-major: h outer, w inner). Needed by the conv pixel head
+    # to reshape the flat token sequence back to a 2D feature map.
+    token_h: int = 0
+    token_w: int = 0
diff --git a/lightx2v/models/networks/neopp/infer/pre_infer.py b/lightx2v/models/networks/neopp/infer/pre_infer.py
@@ -136,7 +136,15 @@ def infer(self, weights):
             timestep_embeddings += noise_embeddings
         image_embeds = image_embeds + timestep_embeddings
 
-        return NeoppPreInferModuleOutput(image_embeds=image_embeds, t=t, z=z, image_token_num=token_h * token_w, timestep_embeddings=timestep_embeddings)
+        return NeoppPreInferModuleOutput(
+            image_embeds=image_embeds,
+            t=t,
+            z=z,
+            image_token_num=token_h * token_w,
+            timestep_embeddings=timestep_embeddings,
+            token_h=token_h,
+            token_w=token_w,
+        )
 
     def patchify(self, images, patch_size, channel_first=False):
         """

diff --git a/lightx2v/models/networks/neopp/infer/transformer_infer.py b/lightx2v/models/networks/neopp/infer/transformer_infer.py
@@ -117,6 +117,11 @@ def __init__(self, config):
         self.scaling = self.head_dim**-0.5
         self.use_triton_qknorm_rope = config.get("use_triton_qknorm_rope", True)
         self.version = config.get("version", "moe")
+        # Conv "pixel head" variant. patch_size/merge_size mirror NeoppPreInfer so the
+        # conv head can reshape tokens -> 2D map and re-patchify back to token space.
+        self.use_pixel_head = config.get("use_pixel_head", False)
+        self.patch_size = config.get("patch_size", 16)
+        self.merge_size = 2
         self.fi_moe_autotune = MoeFiAutotune.from_neopp_config(config)
         if self.version == "moe":
             self.num_experts_per_tok = llm_config["num_experts_per_tok"]
@@ -164,7 +169,7 @@ def infer(self, weights, pre_infer_out, inputs):
         hidden_states = self.infer_without_offload(weights.blocks, hidden_states, cos_sin, past_key_values)
 
         hidden_states = weights.norm_mot_gen.apply(hidden_states)
-        hidden_states = self._fm_head(weights.fm_head, hidden_states)
+        hidden_states = self._fm_head(weights.fm_head, hidden_states, pre_infer_out.token_h, pre_infer_out.token_w)
         return hidden_states.unsqueeze(0)
 
     def infer_without_offload(self, blocks, hidden_states, cos_sin, past_key_values):
@@ -436,12 +441,44 @@ def _sparse_moe(self, moe_w, hidden_states):
         return output
 
     # @ProfilingContext4DebugL1("FM Head")
-    def _fm_head(self, fm_head_w, hidden_states):
+    def _fm_head(self, fm_head_w, hidden_states, token_h=None, token_w=None):
+        if self.use_pixel_head:
+            return self._fm_head_pixel(fm_head_w, hidden_states, token_h, token_w)
         hidden_states = fm_head_w.fm_head_0.apply(hidden_states)
         hidden_states = F.gelu(hidden_states)
         hidden_states = fm_head_w.fm_head_2.apply(hidden_states)
         return hidden_states
 
+    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.)
-        """
-        # [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()
-    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()
+        # ConvDecoder: ps(2) -> conv1 -> gelu -> ps(2) -> conv2 -> ps(8) == x32 upsample
+        x = F.pixel_shuffle(x, 2)
+        x = fm_head_w.conv1.apply(x)
+        x = F.gelu(x)
+        x = F.pixel_shuffle(x, 2)
+        x = fm_head_w.conv2.apply(x)
+        x = F.pixel_shuffle(x, 8)  # [1, 3, H, W]
+        # Re-patchify to [1, L, (patch_size*merge_size)**2 * 3], mirrors NeoppPreInfer.patchify
+        x = self._patchify_pixels(x, self.patch_size * self.merge_size)
+        return x.squeeze(0)
+
+    @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)
+        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)
-    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)
-    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)
+        x = torch.einsum("nchpwq->nhwpqc", x)
+        x = x.reshape(images.shape[0], h * w, patch_size * patch_size * 3)
+        return x
+
     def _dense_mlp(self, mlp_w, hidden_states):
         up_states = mlp_w.up_proj.apply(hidden_states)
         gate_states = mlp_w.gate_proj.apply(hidden_states)

diff --git a/lightx2v/models/networks/neopp/weights/transformer_weights.py b/lightx2v/models/networks/neopp/weights/transformer_weights.py
@@ -11,6 +11,7 @@
 from lightx2v.common.ops.norm.rms_norm_weight import RMSWeightFusedQKNorm3DRope
 from lightx2v.utils.registry_factory import (
     ATTN_WEIGHT_REGISTER,
+    CONV2D_WEIGHT_REGISTER,
     MM_WEIGHT_REGISTER,
     RMS_WEIGHT_REGISTER,
 )
@@ -44,7 +45,7 @@ def __init__(self, config, lazy_load_path=None, lora_path=None):
 
         self.add_module(
             "fm_head",
-            NeoppFmHeadWeights(self.mm_type),
+            NeoppFmHeadWeights(self.mm_type, config=self.config),
         )
 
 
@@ -212,23 +213,49 @@ def __init__(self, block_index, mm_type, lora_path=None):
 
 
 class NeoppFmHeadWeights(WeightModule):
-    def __init__(self, mm_type):
+    def __init__(self, mm_type, config=None):
         super().__init__()
         lora_prefix = "fm_modules"
-        self.add_module(
-            "fm_head_0",
-            MM_WEIGHT_REGISTER["Default"](
-                "fm_modules.fm_head.0.weight",
-                "fm_modules.fm_head.0.bias",
-                lora_prefix=lora_prefix,
-            ),
-        )
+        # New "pixel head" variant (use_pixel_head=True): fm_head is a ConvDecoder
+        # (PixelShuffle + Conv2d) that decodes hidden states directly to RGB pixels,
+        # instead of the legacy per-token MLP head (fm_head.0 / fm_head.2).
+        self.use_pixel_head = bool(config.get("use_pixel_head", False)) if config is not None else False
+        if self.use_pixel_head:
+            # Conv2d(k=3, p=1). in/out channels are inferred from PixelShuffle upstream,
+            # weights carry the true shapes: conv1[1024,1024,3,3], conv2[192,256,3,3].
+            self.add_module(
+                "conv1",
+                CONV2D_WEIGHT_REGISTER["Default"](
+                    "fm_modules.fm_head.conv1.weight",
+                    "fm_modules.fm_head.conv1.bias",
+                    stride=1,
+                    padding=1,
+                ),
+            )
+            self.add_module(
+                "conv2",
+                CONV2D_WEIGHT_REGISTER["Default"](
+                    "fm_modules.fm_head.conv2.weight",
+                    "fm_modules.fm_head.conv2.bias",
+                    stride=1,
+                    padding=1,
+                ),
+            )
+        else:
+            self.add_module(
+                "fm_head_0",
+                MM_WEIGHT_REGISTER["Default"](
+                    "fm_modules.fm_head.0.weight",
+                    "fm_modules.fm_head.0.bias",
+                    lora_prefix=lora_prefix,
+                ),
+            )
 
-        self.add_module(
-            "fm_head_2",
-            MM_WEIGHT_REGISTER["Default"](
-                "fm_modules.fm_head.2.weight",
-                "fm_modules.fm_head.2.bias",
-                lora_prefix=lora_prefix,
-            ),
-        )
+            self.add_module(
+                "fm_head_2",
+                MM_WEIGHT_REGISTER["Default"](
+                    "fm_modules.fm_head.2.weight",
+                    "fm_modules.fm_head.2.bias",
+                    lora_prefix=lora_prefix,
+                ),
+            )