Set enable_hlfb to true for PaliGemma image encoder

- XNNPACK supports all-zeros mask not passed by cl/722748270.
- Calculate pixel size from config, not from a command flag.
- Don't assume the first dimension of pixel value is always 1 which is not the case for Qwen VL, for example.

PiperOrigin-RevId: 725867638

Author: ai-edge-bot

ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py

@@ -56,11 +56,6 @@
     1280,
     'The maximum size of KV cache buffer, including both prefill and decode.',
 )
-_PIXEL_VALUES_SIZE = flags.DEFINE_multi_integer(
-    'pixel_values_size',
-    [3, 224, 224],
-    'The size of prefill pixel values except the batch dimension.',
-)
 _QUANTIZE = flags.DEFINE_bool(
     'quantize',
     True,
@@ -75,12 +70,15 @@ def main(_):
       kv_cache_max_len=_KV_CACHE_MAX_LEN.value,
   )
 
+  config = pytorch_model.image_encoder.config.image_embedding
   converter.convert_to_tflite(
       pytorch_model,
       output_path=_OUTPUT_PATH.value,
       output_name_prefix=f'{_OUTPUT_NAME_PREFIX.value}_{_VERSION.value}',
       prefill_seq_len=_PREFILL_SEQ_LEN.value,
-      pixel_values_size=torch.Size(_PIXEL_VALUES_SIZE.value),
+      pixel_values_size=torch.Size(
+          [1, config.channels, config.image_size, config.image_size]
+      ),
       quantize=_QUANTIZE.value,
       config=pytorch_model.config.decoder_config,
       export_config=ExportConfig(),

ai_edge_torch/generative/examples/paligemma/image_encoder.py

@@ -136,9 +136,7 @@ def get_image_encoder_config() -> cfg.ModelConfig:
       image_embedding=image_embedding_config,
       block_configs=block_config,
       final_norm_config=norm_config,
-      # TODO: b/377051577 - Once RemoveSDPACompositeZeroMaskPass is removed,
-      # enable_hlfb can be set to True. See b/383865404#comment3 for details.
-      # enable_hlfb=True,
+      enable_hlfb=True,
   )
   return config
 

ai_edge_torch/generative/utilities/converter.py

@@ -145,7 +145,7 @@ def _export_helper(
     prefill_input_pos_list.append(torch.arange(0, seq_len, dtype=torch.int))
 
   prefill_pixel_values = (
-      torch.full((1,) + pixel_values_size, 0, dtype=torch.float32)
+      torch.full(pixel_values_size, 0, dtype=torch.float32)
       if pixel_values_size
       else None
   )