{"record":{"id":"72b65f6115fc50ba","repo":"Comfy-Org/ComfyUI","slug":"image-token-count-image-idx-shape-0-vit-outp","errorCode":null,"errorMessage":"Image-token count {image_idx.shape[0]} != ViT output count {image_embeds.shape[0]}; check tokenizer/processor alignment.","messagePattern":"Image-token count (.+?) != ViT output count (.+?); check tokenizer/processor alignment\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"critical","filePath":"comfy/ldm/hidream_o1/model.py","lineNumber":181,"sourceCode":"            # ViT output is constant across sampling steps within a generation\n            # identity-key by the input tensor so refs don't recompute every step.\n            cached = self._visual_cache\n            if cached is not None and cached[0] is ref_pixel_values:\n                image_embeds = cached[1]\n            else:\n                ref_pv = ref_pixel_values.to(inputs_embeds.device)\n                ref_grid = ref_image_grid_thw.to(inputs_embeds.device).long()\n                # extra_conds wraps with a leading batch dim; refs are model-level so [0] always recovers them.\n                if ref_pv.dim() == 3:\n                    ref_pv = ref_pv[0]\n                if ref_grid.dim() == 3:\n                    ref_grid = ref_grid[0]\n                image_embeds = self.visual(ref_pv, ref_grid).to(inputs_embeds.dtype)\n                self._visual_cache = (ref_pixel_values, image_embeds)\n            # image_pad positions identical across batch (input_ids shared cond/uncond).\n            image_idx = (input_ids[0] == IMAGE_TOKEN_ID).nonzero(as_tuple=True)[0]\n            if image_idx.shape[0] != image_embeds.shape[0]:\n                raise ValueError(\n                    f\"Image-token count {image_idx.shape[0]} != ViT output count \"\n                    f\"{image_embeds.shape[0]}; check tokenizer/processor alignment.\"\n                )\n            inputs_embeds[:, image_idx] = image_embeds.unsqueeze(0).expand(B, -1, -1)\n\n        sigma = timesteps.float() / 1000.0\n        t_pixeldit = 1.0 - sigma\n        t_emb = self.t_embedder1(t_pixeldit * 1000, inputs_embeds.dtype)\n        tms_mask_3d = (input_ids == self.tms_token_id).unsqueeze(-1).expand_as(inputs_embeds)\n        inputs_embeds = torch.where(tms_mask_3d, t_emb.unsqueeze(1).expand_as(inputs_embeds), inputs_embeds)\n\n        vinputs_embedded = self.x_embedder(vinputs.to(inputs_embeds.dtype))\n        inputs_embeds = torch.cat([inputs_embeds, vinputs_embedded], dim=1)\n\n        # extra_conds stores position_ids as (1, 3, T); process_cond repeats dim 0 to B. Take row 0.\n        freqs_cis = self.language_model.compute_freqs_cis(position_ids[0].to(x.device), x.device)\n        freqs_cis = tuple(t.to(x.dtype) for t in freqs_cis)\n","sourceCodeStart":163,"sourceCodeEnd":199,"githubUrl":"https://github.com/Comfy-Org/ComfyUI/blob/1c6d8d45b3693bfbb32385b410d813a7fd6be216/comfy/ldm/hidream_o1/model.py#L163-L199","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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)","Re-run the tokenizer/processor on the same image so the number of inserted image tokens equals ViT token count","Clear self._visual_cache (or force a new forward without it) if you changed the reference image or input_ids between calls","Check the processor's image resizing/patch/merge settings against the checkpoint's visual config"],"exampleFix":"# before: grid from a different (larger) image\nref_grid = torch.tensor([[3, 16, 16]])  # 256 tokens\n# input_ids has 576 image tokens -> ValueError\n\n# after: recompute grid from the actual pixel values via the processor\ninputs = processor(images=pil_image, text=prompt, return_tensors=\"pt\")\nref_pv = inputs.pixel_values\nref_grid = inputs.image_grid_thw  # token counts now align","handlingStrategy":"validation","validationCode":"n_img_tokens = int((input_ids[0] == IMAGE_TOKEN_ID).sum())\nn_vit_tokens = int(ref_image_grid_thw[0].prod()) // MERGE  # per your merge factor\nassert n_img_tokens == n_vit_tokens, (n_img_tokens, n_vit_tokens)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always produce pixel_values and image_grid_thw from the same processor call as input_ids","Invalidate cached visual embeddings when the reference image changes"],"tags":["hidream","multimodal","tokenizer","vit","shape-mismatch"],"backgroundTag":null,"analyzedSha":"1c6d8d45b3693bfbb32385b410d813a7fd6be216","analyzedAt":"2026-08-14T19:37:18.893Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}