epochs: 100
output_dir: output_dir

dataset:
  train:
    name: FirstOrderDataset
    batch_size: 8
    num_workers: 4
    use_shared_memory: False
    phase: train
    dataroot: data/first_order/Voxceleb/
    frame_shape: [256, 256, 3]
    id_sampling: True
    pairs_list: None
    time_flip: True
    num_repeats: 75
    create_frames_folder: False
    transforms:
      - name: PairedRandomHorizontalFlip
        prob: 0.5
        keys: [image, image]
      - name: PairedColorJitter
        brightness: 0.1
        contrast: 0.1
        saturation: 0.1
        hue: 0.1
        keys: [image, image]
  test:
    name: FirstOrderDataset
    dataroot: data/first_order/Voxceleb/
    phase: test
    batch_size: 1
    num_workers: 1
    time_flip: False
    id_sampling: False
    create_frames_folder: False
    frame_shape: [ 256, 256, 3 ]


model:
  name: FirstOrderModelMobile
  mode: generator # should be kp_detector, generator, both
  kp_weight_path: None
  gen_weight_path: None
  common_params:
    num_kp: 10
    num_channels: 3
    estimate_jacobian: True
  generator:
    name: FirstOrderGenerator
    kp_detector_cfg:
      temperature: 0.1
      block_expansion: 32
      max_features: 256
      scale_factor: 0.25
      num_blocks: 5
      mobile_net: True
    generator_cfg:
      block_expansion: 32
      max_features: 256
      num_down_blocks: 2
      num_bottleneck_blocks: 6
      estimate_occlusion_map: True
      dense_motion_params:
        block_expansion: 32
        max_features: 256
        num_blocks: 5
        scale_factor: 0.25
      mobile_net: True
  generator_ori:
    name: FirstOrderGenerator
    kp_detector_cfg:
      temperature: 0.1
      block_expansion: 32
      max_features: 1024
      scale_factor: 0.25
      num_blocks: 5
    generator_cfg:
      block_expansion: 64
      max_features: 512
      num_down_blocks: 2
      num_bottleneck_blocks: 6
      estimate_occlusion_map: True
      dense_motion_params:
        block_expansion: 64
        max_features: 1024
        num_blocks: 5
        scale_factor: 0.25
  discriminator:
    name: FirstOrderDiscriminator
    discriminator_cfg:
      scales: [1]
      block_expansion: 32
      max_features: 512
      num_blocks: 4
      sn: True
  train_params:
    num_epochs: 100
    scales: [1, 0.5, 0.25, 0.125]
    checkpoint_freq: 50
    transform_params:
      sigma_affine: 0.05
      sigma_tps: 0.005
      points_tps: 5
    loss_weights:
      generator_gan: 1
      discriminator_gan: 1
      feature_matching: [10, 10, 10, 10]
      perceptual: [10, 10, 10, 10, 10]
      equivariance_value: 10
      equivariance_jacobian: 10

lr_scheduler:
  name: MultiStepDecay
  epoch_milestones: [2000000000000] #just fix learning rate, and change into [237360, 356040] for "both" mode fine tune
  lr_generator: 2.0e-4
  lr_discriminator: 2.0e-4
  lr_kp_detector: 2.0e-4

reconstruction_params:
  num_videos: 1000
  format: '.mp4'

animate_params:
  num_pairs: 50
  format: '.mp4'
  normalization_params:
    adapt_movement_scale: False
    use_relative_movement: True
    use_relative_jacobian: True

visualizer_params:
  kp_size: 5
  draw_border: True
  colormap: 'gist_rainbow'

log_config:
  interval: 10
  visiual_interval: 10

validate:
  interval: 20000000000  #close validate step to speed up training, if you need to see the process, change it into 20000
  save_img: true

snapshot_config:
  interval: 5

optimizer:
  name: Adam

export_model:
  - {}
