google-deepmind / google-deepmind/gemma

Remove TensorFlow dependency from multimodal image preprocessing

Open
#547 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
5.7k
Forks
1k
Avg merge
10h 33m
Merged PRs (30d)
2

Description

Current Behavior
Right now in
gemma/multimodal/image.py
, the code uses TensorFlow's tf.image.decode_jpeg to handle JPEG decoding. This means anyone using the multimodal features has to install the full TensorFlow package, which is pretty heavy (several hundred MB) just for that one image decoding function.

There's actually a TODO comment at line 73 that says: # TODO(eyvinec): we should remove tf dependency.

Expected Behavior
For a library that's built on JAX and Flax, it would make way more sense to use a lighter image library like PIL (Pillow) for the JPEG decoding. This would let people work with multimodal models without needing to install TensorFlow at all.

Why This Matters
Smaller install size: Removing TensorFlow as a dependency would make the install much lighter
Fewer conflicts: TensorFlow can sometimes have CUDA version conflicts with other libraries
Better dev experience: People who just want to use Gemma with JAX shouldn't need to install another whole ML framework
What I'm Thinking
Replace the tf.image.decode_jpeg call with PIL's Image.open and convert the result to a JAX array. The output format would stay exactly the same (H, W, 3 shape with the right dtype), so this shouldn't break anything for existing users.

I'd be happy to put together a PR for this if it sounds good. I'm currently working on a GSoC project proposal around Gemma and came across this while setting up my environment.

Environment
Gemma: main branch
Python: 3.10+
Let me know what you think!

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.