content: # This file includes code derived from the SiT project (https://github.com/willisma/SiT),
  # which is licensed under the MIT License.
  #
  # MIT License
  #
  # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # Permission is hereby granted, free of charge, to any person obtaining a copy
  # of this software and associated documentation files (the "Software"), to deal
  # in the Software without restriction, including without limitation the rights
  # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  # copies of the Software, and to permit persons to whom the Software is
  # furnished to do so, subject to the following conditions:
  #
  # The above copyright notice and this permission notice shall be included in all
  # copies or substantial portions of the Software.
  #
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  # SOFTWARE.
  
  from .transport import Transport, ModelType, WeightType, PathType, Sampler
  
  
  def create_transport(
      path_type='Linear',
      prediction="velocity",
      loss_weight=None,
      train_eps=None,
      sample_eps=None,
      train_sample_type="uniform",
      mean = 0.0,
      std = 1.0,
      shift_scale = 1.0,
  ):
      """function for creating Transport object
      **Note**: model prediction defaults to velocity
      Args:
      - path_type: type of path to use; default to linear
      - learn_score: set model prediction to score
      - learn_noise: set model prediction to noise
      - velocity_weighted: weight loss by velocity weight
      - likelihood_weighted: weight loss by likelihood weight
      - train_eps: small epsilon for avoiding instability during training
      - sample_eps: small epsilon for avoiding instability during sampling
      """
  
      if prediction == "noise":
          model_type = ModelType.NOISE
      elif prediction == "score":
          model_type = ModelType.SCORE
      else:
          model_type = ModelType.VELOCITY
  
      if loss_weight == "velocity":
          loss_type = WeightType.VELOCITY
      elif loss_weight == "likelihood":
          loss_type = WeightType.LIKELIHOOD
      else:
          loss_type = WeightType.NONE
  
      path_choice = {
          "Linear": PathType.LINEAR,
          "GVP": PathType.GVP,
          "VP": PathType.VP,
      }
  
      path_type = path_choice[path_type]
  
      if (path_type in [PathType.VP]):
          train_eps = 1e-5 if train_eps is None else train_eps
          sample_eps = 1e-3 if train_eps is None else sample_eps
      elif (path_type in [PathType.GVP, PathType.LINEAR] and model_type != ModelType.VELOCITY):
          train_eps = 1e-3 if train_eps is None else train_eps
          sample_eps = 1e-3 if train_eps is None else sample_eps
      else:  # velocity & [GVP, LINEAR] is stable everywhere
          train_eps = 0
          sample_eps = 0
  
      # create flow state
      state = Transport(
          model_type=model_type,
          path_type=path_type,
          loss_type=loss_type,
          train_eps=train_eps,
          sample_eps=sample_eps,
          train_sample_type=train_sample_type,
          mean=mean,
          std=std,
          shift_scale =shift_scale,
      )
  
      return state
  