Comfy-Org/ComfyUI · critical · ValueError

Image-token count {image_idx.shape[0]} != ViT output count {

Error message

Image-token count {image_idx.shape[0]} != ViT output count {image_embeds.shape[0]}; check tokenizer/processor alignment.

What it means

Raised during HiDream-O1 multimodal forward when the number of IMAGE_TOKEN_ID placeholders in input_ids does not match the number of patch tokens the ViT visual tower produced for the reference image. inputs_embeds[:, image_idx] would then scatter-assign mismatched tensors, so the model refuses rather than corrupt the sequence. Root cause is nearly always a processor/tokenizer and ViT grid computation disagreeing on image resolution/patch grid.

Source

Thrown at comfy/ldm/hidream_o1/model.py:181

            # ViT output is constant across sampling steps within a generation
            # identity-key by the input tensor so refs don't recompute every step.
            cached = self._visual_cache
            if cached is not None and cached[0] is ref_pixel_values:
                image_embeds = cached[1]
            else:
                ref_pv = ref_pixel_values.to(inputs_embeds.device)
                ref_grid = ref_image_grid_thw.to(inputs_embeds.device).long()
                # extra_conds wraps with a leading batch dim; refs are model-level so [0] always recovers them.
                if ref_pv.dim() == 3:
                    ref_pv = ref_pv[0]
                if ref_grid.dim() == 3:
                    ref_grid = ref_grid[0]
                image_embeds = self.visual(ref_pv, ref_grid).to(inputs_embeds.dtype)
                self._visual_cache = (ref_pixel_values, image_embeds)
            # image_pad positions identical across batch (input_ids shared cond/uncond).
            image_idx = (input_ids[0] == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0]
            if image_idx.shape[0] != image_embeds.shape[0]:
                raise ValueError(
                    f"Image-token count {image_idx.shape[0]} != ViT output count "
                    f"{image_embeds.shape[0]}; check tokenizer/processor alignment."
                )
            inputs_embeds[:, image_idx] = image_embeds.unsqueeze(0).expand(B, -1, -1)

        sigma = timesteps.float() / 1000.0
        t_pixeldit = 1.0 - sigma
        t_emb = self.t_embedder1(t_pixeldit * 1000, inputs_embeds.dtype)
        tms_mask_3d = (input_ids == self.tms_token_id).unsqueeze(-1).expand_as(inputs_embeds)
        inputs_embeds = torch.where(tms_mask_3d, t_emb.unsqueeze(1).expand_as(inputs_embeds), inputs_embeds)

        vinputs_embedded = self.x_embedder(vinputs.to(inputs_embeds.dtype))
        inputs_embeds = torch.cat([inputs_embeds, vinputs_embedded], dim=1)

        # extra_conds stores position_ids as (1, 3, T); process_cond repeats dim 0 to B. Take row 0.
        freqs_cis = self.language_model.compute_freqs_cis(position_ids[0].to(x.device), x.device)
        freqs_cis = tuple(t.to(x.dtype) for t in freqs_cis)

View on GitHub (pinned to 1c6d8d45b3)

Solutions

1. Verify ref_image_grid_thw matches the resolution of ref_pixel_values (t*h/patch * t*w/patch after spatial merge equals ViT output rows)
2. Re-run the tokenizer/processor on the same image so the number of inserted image tokens equals ViT token count
3. Clear self._visual_cache (or force a new forward without it) if you changed the reference image or input_ids between calls
4. Check the processor's image resizing/patch/merge settings against the checkpoint's visual config

Example fix

# before: grid from a different (larger) image
ref_grid = torch.tensor([[3, 16, 16]])  # 256 tokens
# input_ids has 576 image tokens -> ValueError

# after: recompute grid from the actual pixel values via the processor
inputs = processor(images=pil_image, text=prompt, return_tensors="pt")
ref_pv = inputs.pixel_values
ref_grid = inputs.image_grid_thw  # token counts now align

n_img_tokens = int((input_ids[0] == IMAGE_TOKEN_ID).sum())
n_vit_tokens = int(ref_image_grid_thw[0].prod()) // MERGE  # per your merge factor
assert n_img_tokens == n_vit_tokens, (n_img_tokens, n_vit_tokens)

When it happens

Trigger: Calling the model forward with ref_pixel_values + ref_image_grid_thw whose grid produces more/fewer ViT tokens than the <image> tokens embedded in input_ids by the tokenizer/processor; e.g. processor padding/cropping config differs from the grid used at build time, or a cached _visual_cache from a different image being reused with new input_ids.

Common situations: Swapping reference images or resolutions without re-tokenizing, using a processor config (image size, patch size, merge ops) that does not match the checkpoint's visual config, mixing cond/uncond batches where the image cache path diverges, or manually editing input_ids.

Related errors

• {type(model).__name__} must implement map_context_window_to_
• Unknown patch method: {self.patch_method}
• insupportable scoring function for MoE gating: {self.scoring
• attention_mask must be {expected}, got {tuple(attention_mask
• MiniMax Music3 prompt has {prompt_tokens} tokens; maximum is

AI-assisted analysis of Comfy-Org/ComfyUI@1c6d8d45b3 (2026-08-14). Data as JSON: /api/errors/72b65f6115fc50ba. Report an issue: GitHub.