In the forward pass, which we recommend defining as __call__, you have complete creative freedom to add whatever feature you’d like. For our amazing one-step pipeline, create a random image and only call the unet and scheduler once by setting timestep=1: