model:
  target: flae.stegan.STEGAN
  params:
    cover_key: image
    secret_key: secret
    secret_len: 100
    resolution: 256
    use_ema: false
    lr_scheduler: CosineAnnealingRestartCyclicLR
    secret_encoder_config:
      target: flae.unet.Unet1
      params:
        width: 32
        ndown: 4
        nmiddle: 1
        activ: silu
        secret_len: 100
        resolution: 256
    secret_decoder_config:
      target: flae.models.SecretDecoder
      params:
        arch: resnet50
        secret_len: 100
        resolution: 224
    discriminator_config:
      target: flae.munit.MsDCDisGP
      params:
        gp_weight: 10.0
        num_scales: 1
        norm: none
    loss_config:
      target: cldm.loss.ImageSecretLoss
      params:
        recon_type: ffl+yuv
        recon_weight: 1.5
        perceptual_weight: 1.0
        kl_weight: 0.0
        secret_weight: 20.0
        generator_weight: 0.5
        discriminator_weight: 1.0
        generator_update_freq: 2
        max_image_weight_ratio: 15.0
    noise_config:
      target: cldm.transformations2.TransformNet
      params:
        ramp: 10000
        severity: medium
        crop_mode: resized_crop
        gamma: false
data:
  target: tools.imgcap_dataset.DataModuleFromConfig
  params:
    batch_size: 32
    num_workers: 4
    wrap: false
    use_worker_init_fn: true
    train:
      target: tools.image_dataset.ImageFolder
      params:
        data_dir: /mnt/fast/nobackup/scratch4weeks/tb0035/projects/diffsteg/data/mir1M/images
        data_list: /mnt/fast/nobackup/users/tb0035/projects/diffsteg/ControlNet/prep_data/mir_train2.csv
        resize: 256
        secret_len: 100
    validation:
      target: tools.image_dataset.ImageFolder
      params:
        data_dir: /mnt/fast/nobackup/scratch4weeks/tb0035/projects/diffsteg/data/mir1M/images
        data_list: /mnt/fast/nobackup/users/tb0035/projects/diffsteg/ControlNet/prep_data/mir_val2.csv
        resize: 256
        secret_len: 100
lightning:
  callbacks:
    image_logger:
      target: cldm.logger.ImageLogger
      params:
        batch_frequency: 5000
        max_images: 4
        increase_log_steps: false
        fixed_input: true
    progress_bar:
      target: lightning.pytorch.callbacks.TQDMProgressBar
      params:
        refresh_rate: 100
    checkpoint:
      target: lightning.pytorch.callbacks.ModelCheckpoint
      params:
        verbose: true
        filename: '{epoch:06}-{step:09}'
        every_n_train_steps: 10000
  trainer:
    benchmark: true
    base_learning_rate: 4.0e-06
    lr_mult: true
    accumulate_grad_batches: 1
    max_epochs: 150
