Add Diffusion Gemma model to Hub#2804
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for discrete block-diffusion language models in KerasHub, specifically implementing the BlockDiffusionLM and Gemma4BlockDiffusionLM tasks, their preprocessors, the EntropyBoundSampler, and a self-conditioning layer. It also updates the Gemma4 backbone and decoder blocks to support diffusion-specific mechanisms like dual layer scalars and bidirectional canvas attention, and adds a checkpoint conversion script for DiffusionGemma. The review feedback is highly constructive, pointing out critical robustness issues: ensuring task models can generate immediately after instantiation without explicit compilation, resetting the stateful sampler variable on new generation runs to prevent shape mismatch crashes across varying batch sizes, and strictly validating all multimodal inputs in the text-only preprocessor path to avoid silent failures.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| super().__init__(*args, **kwargs) | ||
| self.canvas_length = canvas_length | ||
| self.max_denoising_steps = max_denoising_steps | ||
| self.t_min = t_min | ||
| self.t_max = t_max |
There was a problem hiding this comment.
In KerasHub, task models should be callable for generation immediately after instantiation (e.g., when loaded from a preset) without requiring an explicit compile() call. Currently, self.sampler is only initialized inside compile(), which will cause generate() to raise an AttributeError if called on an uncompiled model.
Initializing self.sampler to a default EntropyBoundSampler in __init__ resolves this issue and ensures a seamless user experience.
super().__init__(*args, **kwargs)
self.canvas_length = canvas_length
self.max_denoising_steps = max_denoising_steps
self.t_min = t_min
self.t_max = t_max
from keras_hub.src.samplers.entropy_bound_sampler import (
EntropyBoundSampler,
)
self.sampler = EntropyBoundSampler(
vocabulary_size=self.backbone.vocabulary_size
)References
- Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. (link)
| def __call__(self, canvas, logits, step): | ||
| logits = ops.cast(logits, "float32") |
There was a problem hiding this comment.
self._prev_argmax is a stateful keras.Variable that is lazily initialized on the first call. However, if the batch size changes between independent generate() calls, the shape of self._prev_argmax from the previous run will mismatch with the new batch size, causing a shape mismatch crash during assign or equal operations.
Resetting self._prev_argmax to None at step == 0 (the start of a new generation run) ensures that the variable is correctly re-initialized with the new batch shape, preventing any potential crashes.
| def __call__(self, canvas, logits, step): | |
| logits = ops.cast(logits, "float32") | |
| def __call__(self, canvas, logits, step): | |
| if step == 0: | |
| self._prev_argmax = None | |
| logits = ops.cast(logits, "float32") |
References
- Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. (link)
| if self.text_only_model and audio is not None: | ||
| raise ValueError( | ||
| "The initialized preprocessor/model is text-only, but " | ||
| "`audio` is not `None`." | ||
| ) |
There was a problem hiding this comment.
For a text-only preprocessor, passing multimodal inputs during training (call()) should raise a ValueError to prevent silently ignoring them. Currently, only audio is checked, while other multimodal inputs like images, videos, or pixel_values are silently ignored.
Updating the check to include all multimodal inputs ensures consistency with generate_preprocess() and prevents silent failures.
| if self.text_only_model and audio is not None: | |
| raise ValueError( | |
| "The initialized preprocessor/model is text-only, but " | |
| "`audio` is not `None`." | |
| ) | |
| if self.text_only_model and ( | |
| images is not None | |
| or videos is not None | |
| or pixel_values is not None | |
| or audio is not None | |
| ): | |
| raise ValueError( | |
| "The initialized preprocessor/model is text-only, but " | |
| "multimodal inputs (images, videos, pixel_values, or audio) " | |
| "were provided." | |
| ) |
References
- Demand Robustness: Do not accept fragile code. If the proposed code is not robust enough or lacks proper error handling, explicitly tell the author why the current approach is brittle and what must be done to reinforce it. (link)
1 participant
Description of the change
Reference
Colab Notebook
Checklist