From dc7de65b70c8f15a614020f9a38af36f7f7a29ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20CREPY?= Date: Thu, 16 Jul 2026 02:28:02 -0700 Subject: [PATCH] Cleanup protocols naming and convert Union types to Protocols. PiperOrigin-RevId: 948845677 --- .../hackable_diffusion_adapter/hd/hd_gemma_network.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py b/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py index bb44dc42..3e6b98ed 100644 --- a/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py +++ b/gemma/diffusion/hackable_diffusion_adapter/hd/hd_gemma_network.py @@ -35,7 +35,7 @@ TargetInfoTree = hd_typing.TargetInfoTree -BaseDiffusionNetwork = diffusion_network.BaseDiffusionNetwork +StandardDiffusionNetwork = diffusion_network.StandardDiffusionNetwork DiffusionGemmaModel = gemma_diffusion.DiffusionGemma_26B_A4B @@ -112,8 +112,8 @@ def prefill_kv_cache_with_encoder( ################################################################################ -class WrappedDiffusionGemmaNetwork(nn.Module, BaseDiffusionNetwork): - """Wraps a Diffusion Gemma model as an HD ``BaseDiffusionNetwork``. +class WrappedDiffusionGemmaNetwork(nn.Module, StandardDiffusionNetwork): + """Wraps a Diffusion Gemma model as an HD ``DiffusionNetwork``. The wrapper handles the hackable diffusion interfaces and Gemma specifics. For self conditioning, it uses `call_with_self_conditioning` function from Gemma.