CDDIMInterpolantSampler.py

import tensorflow as tf
from .CBasicInterpolantSampler import CBasicInterpolantSampler
from .CDDIMSamplingAlgorithm import CDDIMSamplingAlgorithm

class CDDIMInterpolantSampler(CBasicInterpolantSampler):
  def __init__(
    self, interpolant,
    stochasticity, noiseProvider, schedule, steps, clipping, projectNoise
  ):
    super().__init__(
      interpolant=interpolant,
      algorithm=CDDIMSamplingAlgorithm(
        stochasticity=stochasticity,
        noiseProvider=noiseProvider,
        schedule=schedule,
        steps=steps,
        clipping=clipping,
        projectNoise=projectNoise
      )
    )
    self._schedule = schedule
    return
  
  def train(self, x0, x1, T, xT=None, model=None):
    B = tf.shape(x0)[0]
    tf.assert_equal(tf.shape(T), (B, 1))
    tf.assert_equal(tf.shape(x0), tf.shape(x1))
    
    T = self._schedule.to_discrete(T, lastStep=True)
    # apply training procedure from interpolant
    alpha_hat_t = self._schedule.parametersForT(T).alphaHat
    tf.assert_equal(tf.shape(alpha_hat_t), (B, 1))
    
    if (xT is not None):
      REPLACE_NOISE = True
      if REPLACE_NOISE:
        x1 = xT
      else:
        x1 = None
      pass

    trainData = self._interpolant.train(x0=x0, x1=x1, T=alpha_hat_t, xT=xT)
    return {
      **trainData,
      'alphaHat': alpha_hat_t,
      'T': self._schedule.to_continuous(T),
    }
# End of CDDIMInterpolantSampler