Taking up the text-to-image example (Imagen-style, frozen BERT encoder)
I'd like to take up the "Text-to-image" item from the call for contributions list.
Some background on me: for my senior project I built latent diffusion from scratch and
trained it on TPU pods with JAX. The only pretrained parts I used were the CLIP text
encoder and SD3 components from KerasHub. For this example the only pretrained piece
would be the frozen BERT text encoder from a KerasHub preset, everything else written
from scratch.
Plan
Following Imagen, the example trains a text-conditioned U-Net denoiser directly in
pixel space at 64x64x3, no VAE. The noise schedule, sampler, denoiser and generation
loop are all implemented in the example itself, in backend-agnostic Keras 3.
For super-resolution I don't plan to include SR training code in the example. Instead
it loads my pretrained SR diffusion model (64 to 256) and runs it at generation time,
so readers get the full pipeline output at 256x256. The SR training script gets linked
as an external reference for anyone who wants to reproduce it. This keeps the example
focused on the base model while still showing the cascade working end to end.
Training and results
The results shown in the example will come from a long run on a large image-text
dataset (around 2M pairs, I'll list the sources for licensing review) on TPU v5e-8
via Kaggle with the JAX backend. The code in the example itself defaults to a small
dataset and a short schedule so readers can actually run it end to end. As a side
effect it should also work as a practical reference for Keras 3 + JAX on TPU v5e-8.
One question
I'll aim for the 300-line guideline, but this covers schedule, sampler, denoiser,
training and cascade generation. If it comes down to it, is going over by up to 100
lines acceptable? I haven't started writing yet and would rather know the hard
ceiling now than find out at review.
Taking up the text-to-image example (Imagen-style, frozen BERT encoder)
I'd like to take up the "Text-to-image" item from the call for contributions list.
Some background on me: for my senior project I built latent diffusion from scratch and
trained it on TPU pods with JAX. The only pretrained parts I used were the CLIP text
encoder and SD3 components from KerasHub. For this example the only pretrained piece
would be the frozen BERT text encoder from a KerasHub preset, everything else written
from scratch.
Plan
Following Imagen, the example trains a text-conditioned U-Net denoiser directly in
pixel space at 64x64x3, no VAE. The noise schedule, sampler, denoiser and generation
loop are all implemented in the example itself, in backend-agnostic Keras 3.
For super-resolution I don't plan to include SR training code in the example. Instead
it loads my pretrained SR diffusion model (64 to 256) and runs it at generation time,
so readers get the full pipeline output at 256x256. The SR training script gets linked
as an external reference for anyone who wants to reproduce it. This keeps the example
focused on the base model while still showing the cascade working end to end.
Training and results
The results shown in the example will come from a long run on a large image-text
dataset (around 2M pairs, I'll list the sources for licensing review) on TPU v5e-8
via Kaggle with the JAX backend. The code in the example itself defaults to a small
dataset and a short schedule so readers can actually run it end to end. As a side
effect it should also work as a practical reference for Keras 3 + JAX on TPU v5e-8.
One question
I'll aim for the 300-line guideline, but this covers schedule, sampler, denoiser,
training and cascade generation. If it comes down to it, is going over by up to 100
lines acceptable? I haven't started writing yet and would rather know the hard
ceiling now than find out at review.