Factor out candidate generation and post-processing in discrete samplers to reduce boilerplate. This makes the update methods focus purely on computing the routing probabilities.

PiperOrigin-RevId: 908000169

Author: vdeborto

hackable_diffusion/lib/sampling/__init__.py

@@ -31,6 +31,8 @@
 from hackable_diffusion.lib.sampling.discrete_step_sampler import NoRemaskingFn
 from hackable_diffusion.lib.sampling.discrete_step_sampler import RemaskingFn
 from hackable_diffusion.lib.sampling.discrete_step_sampler import RescaledRemaskingFn
+from hackable_diffusion.lib.sampling.discrete_step_sampler import RoutingStrategy
+from hackable_diffusion.lib.sampling.discrete_step_sampler import RoutingWeights
 from hackable_diffusion.lib.sampling.discrete_step_sampler import UnMaskingStep
 from hackable_diffusion.lib.sampling.gaussian_step_sampler import AdjustedDDIMStep
 from hackable_diffusion.lib.sampling.gaussian_step_sampler import DDIMStep